speed up function that takes a function as argument with numba

Question

I am trying to use numba to speed up a function that takes another function as argument. A minimal example would be the following:

import numba as nb

def f(x):
    return x*x

@nb.jit(nopython=True)
def call_func(func,x):
    return func(x)

if __name__ == '__main__':
    print(call_func(f,5))

This, however, doesn't work, as apparently numba doesn't know what to do with that function argument. The traceback is quite long:

Traceback (most recent call last):
  File "numba_function.py", line 15, in <module>
    print(call_func(f,5))
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/dispatcher.py", line 330, in _compile_for_args
    raise e
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/dispatcher.py", line 307, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/dispatcher.py", line 579, in compile
    cres = self._compiler.compile(args, return_type)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/dispatcher.py", line 80, in compile
    flags=flags, locals=self.locals)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 740, in compile_extra
    return pipeline.compile_extra(func)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 360, in compile_extra
    return self._compile_bytecode()
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 699, in _compile_bytecode
    return self._compile_core()
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 686, in _compile_core
    res = pm.run(self.status)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 246, in run
    raise patched_exception
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 238, in run
    stage()
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 452, in stage_nopython_frontend
    self.locals)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/compiler.py", line 841, in type_inference_stage
    infer.propagate()
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/typeinfer.py", line 773, in propagate
    raise errors[0]
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/typeinfer.py", line 129, in propagate
    constraint(typeinfer)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/typeinfer.py", line 380, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/opt/local/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/numba/typeinfer.py", line 402, in resolve
    raise TypingError(msg, loc=self.loc)
numba.errors.TypingError: Failed at nopython (nopython frontend)
Invalid usage of pyobject with parameters (int64)
No type info available for pyobject as a callable.
File "numba_function.py", line 10
[1] During: resolving callee type: pyobject
[2] During: typing of call at numba_function.py (10)

This error may have been caused by the following argument(s):
- argument 0: cannot determine Numba type of <class 'function'>

Is there a way to fix this?

Answer 1

It depends if the func you pass to call_func can be compiled in nopython mode.

If it can't be compiled in nopython mode then it's impossible because numba doesn't support python calls inside a nopython function (that's the reason why it's called nopython).

However if it can be compiled in nopython mode you can use a closure:

import numba as nb

def f(x):
    return x*x

def call_func(func, x):
    func = nb.njit(func)   # compile func in nopython mode!
    @nb.njit
    def inner(x):
        return func(x)
    return inner(x)

if __name__ == '__main__':
    print(call_func(f,5))

That approach has some obvious downsides because it needs to compile func and inner every time you call call_func. That means it's only viable if the speedup by compiling the function is bigger than the compilation cost. You can mitigate that overhead if you call call_func with the same function several times:

import numba as nb

def f(x):
    return x*x

def call_func(func):  # only take func
    func = nb.njit(func)   # compile func in nopython mode!
    @nb.njit
    def inner(x):
        return func(x)
    return inner  # return the closure

if __name__ == '__main__':
    call_func_with_f = call_func(f)   # compile once
    print(call_func_with_f(5))        # call the compiled version
    print(call_func_with_f(5))        # call the compiled version
    print(call_func_with_f(5))        # call the compiled version
    print(call_func_with_f(5))        # call the compiled version
    print(call_func_with_f(5))        # call the compiled version

Just a general note: I wouldn't create numba functions that take a function argument. If you can't hardcode the function numba can't produce really fast functions and if you also include the compilation cost for closures it's mostly just not worth it.

Answer 2

As suggested by the error message, Numba cannot deal with values of type function. You can check in the documentation what types can Numba work with. The reason is that Numba cannot in general optimize (jit-compile) arbitrary functions in noptyhon mode, they are considered basically a black box (in fact, the passed function could even be a native one!).

The usual approach would be to ask Numba to optimize the called function instead. If you cannot add the decorator to the function (e.g. because it is not part of your source code), you can still use it manually like:

import numba as nb

def f(x):
    return x*x

if __name__ == '__main__':
    f_opt = nb.jit(nopython=True)(f)
    print(f_opt(5))

Obviously it will still fail if f cannot be compiled by Numba either, but in that case there's not much you can do anyway.