Accessing numba function directly

[aseyboldt]

I can see quite a few use cases for accessing the generated numba functions directly, instead of using them through the aesara function wrapper. Mostly if we want to create a function pointer that we can call from other compiled code, using numba.cfunc.

I'm using this is a nuts sampler implemented in rust, and it would also be very useful for sunode, where we could compile the right hand side of an ODE as aesara function, compile that and pass a pointer to the function to the c library sundials to do the actual work.

In nutpie, I currently compile a function in numba mode, then access the optimized fgraph and then use numba_funcify to get an optimized numba function:
https://github.com/aseyboldt/nutpie/blob/master/nutpie/compile_pymc.py#L124

This works, but seems like a rather roundabout way of doing it. I also think it ends up compiling the numba functions twice, and thus slows things down a lot.

Did anyone else do something like this already? Should I just call the the optimization code myself, or should there be an interface for it?

[brandonwillard]

The JIT compiled functions produced by JITLinker are made available through a .jit_fn attribute. More specifically, the Numba/JAX compilation result is available as fn.vm.jit_fn for an aesara.function-compiled function fn.

[aseyboldt]

@brandonwillard Yes, it gets embedded into something else, namely a numba.cfunc (https://github.com/aseyboldt/nutpie/blob/master/nutpie/compile_pymc.py#L169).
This is why I would like to avoid any previous compilation steps, because (I think) those will not be used when the cfunc is compiled.

[brandonwillard]

compiler_machinery is a component of the Numba compilation pipeline, so the next step would involve determining when/where/how often Numba compilation is occuring, and then whether or not any of that is performed redundantly due to the Aesara compilation/transpilation pipeline.

The Aesara compilation/transpilation pipeline shouldn't be calling any of the dispatch functions that perform intermediate Numba compilations, or the final, entire-graph compilations, more than once for a given node/graph, so, in a sense, this compilation latency should be close to the latency one would see when compiling the same graph by hand with Numba.

As I mentioned previously, we are adding some additional compilation passes/settings to the standard Numba compilation pipeline (e.g. in order to get performant vectorized operations in certain cases), so keep that in mind.

[brandonwillard]

Also, the Numba cache should be enabled by default, but I don't know how much of the graph you're using is covered by Numba's caching. This is something we haven't considered—but need to—regarding the type of Python functions we generate in Aesara for Numba (i.e. we might be producing graphs that can't be adequately cached when we could be, or not caching with enough sub-graph granularity to help).

[aseyboldt]

Can you point me to those additional passes? I can't find them...

I played a little bit with the cache argument, but so far I never observed any difference, no clue why.

[brandonwillard]

Anywhere where aesara.link.numba.dispatch.basic.use_optimized_cheap_pass is used will invoke an extra optimization pass in Numba. Again, this is all related to performance concerns for element-wise reduction operations in Numba (see #529).