Improve interface Python - Rust

[felixhekhorn]

• Identify weak points and improvements between numba and Rust
• Identify weak points and improvements between Python and numba and Rust
• Identify weak points and improvements between Rust/eko and Rust/ekore
• Implement improvements
• Update architecture documentation Document current architecture #516

Things to consider:

• build system
• passing numba function pointer around Split ev_op/init #526
• memory allocation and unravelling
• caching strategy
• more optimisation closer to hardware

[felixhekhorn]

@AkshatRai07
I suggest to have a copy of the existing quad_ker_qcd where I replace these lines (and only these lines)

eko/src/eko/evolution_operator/__init__.py

Lines 390 to 401 in 0287086

	if is_polarized:
	if is_time_like:
	raise NotImplementedError("Polarized, time-like is not implemented")
	else:
	gamma_singlet = ad_ps.gamma_singlet(order, ker_base.n, nf)
	else:
	if is_time_like:
	gamma_singlet = ad_ut.gamma_singlet(order, ker_base.n, nf)
	else:
	gamma_singlet = ad_us.gamma_singlet(
	order, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv
	)

with a single call to the eko crate gamma_singlet = ekors.rust_quad_ker(is_polarized, order, blub).

In the first iteration ekors.rust_quad_ker will be: 1) call ekore.blub and 2) return gamma from there. So no more calling python inside Rust, but just returning the ekore computation; also SV or exponentiation remain identical to the current Python implementation.

In a second step we can move as a next candidate SV from Python to Rust (while leave exp untouched). This way we can shift pieces slowly one at a time.

[felixhekhorn]

Step 1:

• split ev_op/init.py into two files
• ev_op/quad_ker.py will hold all quad_ker* related code (=functions); current content can be dropped; this will leave the (whole) library unchanged but for a few (single?) new import(s) in ev_op/init.py
• the new Rust logic flow will then patch quad_ker.py (possibly instead of current ev_op/init.py)
• the definition of the integration path will remain in Python for the moment and will be deactivated for now on the Rust side
• decide something for OME

this will remove the passing of the Python function pointer

[AkshatRai07]

@felixhekhorn I debugged as to why this new code in #526 is slower than the original one.
These are stats for this code (i.e. before #526 merger)

====================== 3 passed, 12 deselected, 1 warning in 517.12s (0:08:37) ======================
        Command being timed: "poe lha -m nnlo and sv"
        User time (seconds): 667.09
        System time (seconds): 6.06
        Percent of CPU this job got: 99%
        Elapsed (wall clock) time (h:mm:ss or m:ss): 11:16.11
        Average shared text size (kbytes): 0
        Average unshared data size (kbytes): 0
        Average stack size (kbytes): 0
        Average total size (kbytes): 0
        Maximum resident set size (kbytes): 616324
        Average resident set size (kbytes): 0
        Major (requiring I/O) page faults: 1410
        Minor (reclaiming a frame) page faults: 573231
        Voluntary context switches: 3869
        Involuntary context switches: 55162
        Swaps: 0
        File system inputs: 351544
        File system outputs: 193368
        Socket messages sent: 0
        Socket messages received: 0
        Signals delivered: 0
        Page size (bytes): 4096
        Exit status: 0

I converted the nb.objectmode into C-Types function, but the speed was at par with the original code. Then I did some changes in memory allocation. Basically the arrays were getting malloc for every scipy call (100s of them for one point), so I passed them from outside the integrate.quad, no malloc in the hot-path. Even then, the speed was same as before.

Main culprit: LowLevelCallable. I wrapped my quad_ker_ad and quad_ker_ome into 2 nb.cfunc, passed them to quad using LLC, and these are the stats:

====================== 3 passed, 12 deselected, 3 warnings in 828.57s (0:13:48)  ======================
        Command being timed: "poe lha -m nnlo and sv"
        User time (seconds): 1214.81
        System time (seconds): 11.79
        Percent of CPU this job got: 100%
        Elapsed (wall clock) time (h:mm:ss or m:ss): 20:24.52
        Average shared text size (kbytes): 0
        Average unshared data size (kbytes): 0
        Average stack size (kbytes): 0
        Average total size (kbytes): 0
        Maximum resident set size (kbytes): 1952660
        Average resident set size (kbytes): 0
        Major (requiring I/O) page faults: 18
        Minor (reclaiming a frame) page faults: 3282977
        Voluntary context switches: 530
        Involuntary context switches: 14346
        Swaps: 0
        File system inputs: 4104
        File system outputs: 312184
        Socket messages sent: 0
        Socket messages received: 0
        Signals delivered: 0
        Page size (bytes): 4096
        Exit status: 0

The speed increased drastically. I am still trying to figure out why it's taking 2x time (if we remove compilation time, it's still taking around 1.6x more time) and 3x RAM when compared to the version before, will let you know once it's done.

[felixhekhorn]

The speed increased drastically.

the runtime increased? the speed decreased, no?

so you're saying nb.objmode is not the problem? but we are using LLC on master, no? I don't understand ...

Anyway, it might be worth to keep an official record of precisely these numbers to track what we are doing. For example we can make a document (e.g. next to the architecture file):

# performance statistics

## 2c49156022235da1bc3582b4b61bd323e42d8d7f
this is after #522 but before #526

[past numbers here]

## commit hash
this is after #526 

[bla]

I think poe lha -m nnlo and sv can serve as reference point - it is sufficiently complicated and a somewhat real life example

[felixhekhorn]

Then I did some changes in memory allocation. Basically the arrays were getting malloc for every scipy call (100s of them for one point), so I passed them from outside the integrate.quad, no malloc in the hot-path.

if we want to optimize memory allocation I guess this would be the optimal solution, right?

[AkshatRai07]

Sorry for the confusion caused, I misspoke.

To clarify my initial thought process: The raw code in #526 was extremely slow. When I wrapped the kernels in nb.cfunc and passed them via LLC, the runtime improved significantly compared to the raw #526, but as you noted, it is still taking ~20 mins (almost 2x the time of the original 2c49156... baseline).

And yes, nb.objmode isn't the primary bottleneck here.

Also, performance tracking document is a great thing, will implement it in a future PR.