2

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?

Achxy_
  • 1,171
  • 1
  • 5
  • 25
  • this will happen on the call to njit (on module initialization) if you specify the argument types and return types of the function, the reason why it delays compilation to when it is called is because it doesn't know the argument types, refer to numba documentation on how to specify argument types. – Ahmed AEK Jul 10 '22 at 11:09

1 Answers1

3

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.

Jérôme Richard
  • 41,678
  • 6
  • 29
  • 59
  • 1
    Thanks! But this does not seem to work for functions with keyword arguments in their signature, for example [this](https://paste.gg/p/anonymous/8176c9b53f6f4cf49dec7323c42b8d19) fails. So how do we go about decorating signature for kwargs? – Achxy_ Jul 10 '22 at 15:41
  • Yeah because `Omitted` or `optional` should be used in this case. See: [this post](https://stackoverflow.com/questions/46123657/numba-calling-jit-with-explicit-signature-using-arguments-with-default-values). However, It looks like it does not work anymore with the last version of Numba: I got a low-level IR error which look like a regression/bug... There is an [opened related issue](https://github.com/numba/numba/issues/7501) on this for the AOT. You could try to use generated_jit but it is cumbersome. A basic python wraper could also work if you do not call the function from Numba. – Jérôme Richard Jul 10 '22 at 16:39