Avoid recompilation of functions that take other functions as arguments

(this is not a question about the persistent cache, which has been answered here)

Please consider this simplified example:

import numba as nb

t0 = 0.
y0 = 1.

@nb.njit
def f(t, y):
    return -y

@nb.njit
def g(t, y):
    return -y

f(t0, y0)
g(t0, y0)

print(f.nopython_signatures)
print(g.nopython_signatures)

yields:

[(float64, array(float64, 1d, C)) -> array(float64, 1d, C)]
[(float64, array(float64, 1d, C)) -> array(float64, 1d, C)]

So both functions have the same signature. I was expecting that the signature of a function that take either f or g does not depend on being f or g. But:

@nb.njit
def step(rhs, t, y):
    return rhs(t+1, 2*y)

step(f, t0, y0)
step(g, t0, y0)

print(step.nopython_signatures)

yields two different signatures:

[
(type(CPUDispatcher(<function f at 0x7fb92e3271f0>)), float64, array(float64, 1d, C)) -> array(float64, 1d, C), 
(type(CPUDispatcher(<function g at 0x7fb92fbf90d0>)), float64, array(float64, 1d, C)) -> array(float64, 1d, C)
]

Possible inlining is the only reason that I can think of this behaviour. Is this the reason?

This means that step is recompiled every time with the associated impact on execution time. Is there a way to avoid this behaviour?

hi @hgrecco, I’ve been working on the topic of higher-order functions (e.g. Semantics of first-class functions) and I have the same question, which I posted on gitter last week (no answer yet).
I’m refining the proposal above, and I need to know whether there’s any performance advantage (like inlining) in re-compiling for every specific function being passed, which is the same question you have.

Regarding how to avoid it, I have created this PR (https://github.com/numba/numba/pull/5579), which would allow you to annotate step as receiving a function argument, and would avoid the duplicated compilation. Whether that’s desirable, it depends on whether there’s a performance advantage or not.

Thanks for posting, it’s good to know I’m not the only one having this problem.

Cheers,
Luk

Hi @luk-f-a Thanks for writing back and providing these useful links

I do think that there might be a performance advantage in certain cases (e.g. inlining a function in a tight loop). However, I also think that the ability to have a “function pointer” in which only the signature is important will allow numba users to build up libraries of algorithms that are useful for other users as well.

Therefore, I think that making this feature it opt-in via annotations is the way to go.

Hernán

update on my question on gitter: sklam confirmed that using “function pointers” might prevent inlining, although LLVM might be able to do it anyway depending on the context.