Home > Software engineering >  How do I make numba compile python functions before initial invocation?
How do I make numba compile python functions before initial invocation?

Time:07-11

I am aware of Ahead of Time compilation (AOT) but it has certain limitations which I believe is a trade-off for accommodating portability of compiled objects. I am also aware of cache kwarg in jit decorator but it has significant overhead too.

Ideally I wanted a function which gets compiled not during the initial invocation but right after calling numba.njit
For example :

from numba import njit
from testthis import test_this


@test_this
@njit
def foo(x):
    return x   1


foo(1)
foo(2)
foo(3)

Which results in :

1st call took 1003.4722999989754 msec and returned 2
2nd call took 0.011500000255182385 msec and returned 3
3rd call took 0.0030000010156072676 msec and returned 4

Clearly, the first call took significantly longer than the subsequent calls. But how to make it so that compilation is done before invocation without the aforementioned limitation (as portability is not of concern)

Example of AOT but with limitations :

from testthis import test_this
from numba.pycc import CC

cc = CC(__name__)


@test_this
@cc.export("foo", "i8(i8)")
def foo(x):
    return x   1

Which results in :

1st call took 0.005500005499925464 msec and returned 2
2nd call took 0.00549999822396785 msec and returned 3
3rd call took 0.0027999994927085936 msec and returned 4

I saw a compile method for numba.core.registry.CPUDispatcher which takes the signature, but this doesn't help alot and only reduces the compilation time by a bit for the initial call.

What would be a way to go about this where first function call doesn't have to suffer overhead whilst not having the limitations of aot?

CodePudding user response:

What you want is called eager compilation as opposed to lazy compilation (which is the default mode). As the documentation states, you need to provide the signature. Multiple signature can be provided if needed. Signature a mandatory for a compilation before the first call because type are mandatory so for the type inference to work and fast instructions to be generated.

  • Related