Caching redefined functions

Numba Community Support
1

In the following example, the cache fails for the first function, but not the second:

from numba import njit

@njit(cache=True)
def incr_basic(x):
    return x + 1

def create_foo():
    @njit(cache=True)
    def incr_basic(x):
        return x + 1

    @njit(cache=True)
    def foo(x):
        return incr_basic(x)

    return foo

def create_foo_with_global():
    @njit(cache=True)
    def foo(x):
        return incr_basic(x)

    return foo

assert create_foo()(3) == create_foo_with_global()(3)
%timeit create_foo()(3)  # 32.8 ms ± 2.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit create_foo_with_global()(3)  # 4.27 ms ± 353 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

This is an issue for us at pytensor as we compile functions on demand, using a pattern that looks like create_foo. AFAICT this happens because the cache key hashes the closure of foo and that is always new for a newly defined function. The relevant line is in: numba/numba/core/caching.py at 5e0ceb077fb219a21ac8a50e5234ca643aaf96f9 · numba/numba · GitHub

Is there a way for us to easily override how the key is defined? We are already using a custom CacheLocator for handling string-generated/compiled functions, but numba will still refuse to use the cached files because of the way the key is defined.

PS: The timings were done in ipython for ease of the %timeit magic, but the issue persists (and is perhaps only relevant) in a regular python script.

Feedback on custom caching strategy in PyTensor
2

Have you considered using compile + exec sequence to compile the function’s code text in the module’s file and then executing it in the module’s namespace?

3

Not sure that makes sense for us? We’re compiling custom functions for users on demand, not a finite set of predefined functions. We already do compile + exec for the user defined functions, the issue is not that, but the inner functions in the closure not hashing to the same when we redefine them.

4

The key to utilizing ‘compile + exec’ for the purposes such as the one you described is to compile and execute the code in the module’s file and the global namespace, which is what addresses the issue of the closure you have encountered. Try this (a somewhat laborious example just to illustrate the point, feel free to shape it into something more streamlined for your needs),

from inspect import getfile, getmodule
from numba import njit  # noqa: F401


def _anchor():
    pass


def create_foo_exec():
    incr_basic_txt = """
@njit(cache=True)
def incr_basic_(x):
    return x + 1  
"""
    inc_basic_code = compile(incr_basic_txt, getfile(_anchor), mode="exec")
    ns = getmodule(_anchor).__dict__
    exec(inc_basic_code, ns)

    foo_txt = """
@njit(cache=True)
def foo_(x):
    return incr_basic_(x)
"""
    foo_code = compile(foo_txt, getfile(_anchor), mode="exec")
    exec(foo_code)

    return locals()["foo_"]


if __name__ == "__main__":
    print(create_foo_exec()(3))
2 Likes
5

Thanks a lot, I’ll give it some thought

6

Another option was to use register_jitable instead of njit for the inner function. That way the original function is used for the cache key which serializes the same way across multiple definitions.

2 Likes
7

Interesting observation - right, the original python function ‘incr_basic’ is given back by the ‘register_jitable’ (which does the ‘overload’ for the jitted context behind the scenes), and having python functions in the closure doesn’t invalidate the cache (it would have been quite inconvenient indeed if it did!). Thanks for sharing!

P.S. I have noticed ‘register_jitable‘ to be less flexible in some cases, for instance, one cannot directly compile the function with it for the desired signature.

2 Likes