Cache njit function at import or install

Numba Support: How do I do ...?
1

I have a package that contains many public njit functions. Of course, these functions are only compiled just-in-time as needed but I was wondering if there was a way to compile these functions either at installation OR at import? Is there a standard approach for this? I know that there is a cache=True option but this still requires the function to be called at least one time by the user. I was thinking that I could sneak this into __init__.py by calling the functions there but any suggestions would be greatly appreciated.

1 Like
2

@seanlaw I would like to know that as well. I think the only way right now to do that is to use AOT compilation https://numba.pydata.org/numba-doc/dev/user/pycc.html

I was thinking: if calling those njit functions from a background thread will it work for the first time, in order to do not freeze the app at startup?

3

@rpopovici Yes, I’ve read that too but AOT sounds like a worse option as it is not optimized for the local machine and the code is pretty ugly from a maintenance perspective. I think performing the JIT compilation and placing the function call inside if __init__.py is another option. This means that import time may be high as it will include compile time but, for a data science workflow (i.e., a Jupyter notebook where import happens once), this might be fine but the solution needs to be easily portable as it is not just for me and my machine.

4

I recently became aware of the approach used by Qibo to force compilation at install time. Each decorator has a signature and cache=True, e.g.

@njit("int64(int64, int32[:])",
      cache=True)
def multicontrol_index(g, qubits):

(from gates.py)

Then after installation in setup.py, the library is imported:

class CustomInstall(install):
    def run(self):
        install.run(self)
        from qibo import K

This forces the compilation and caching of all functions.

Note that this will break binary distribution when you do python setup.py bdist, so Qibo uploads sdists to PyPI instead, creating them with python setup.py sdist.

(Many thanks to Stefano Carrazza for pointing me to this approach).

1 Like