Reputation: 4666
I know that a Numba-jitted function calling another jitted function will recognize this and automatically use a fast C calling convention rather than going through the Python object layer, and therefore avoid the high Python function call overhead:
import numba
@numba.jit
def foo(x):
return x**2
@numba.jit
def bar(x):
return 4 * foo(x) # this will be a fast function call
My question is whether the same is true if I call a Cython function from Numba. So let's say I have a Cython module, foo.pyx
:
cpdef double foo(double x):
return x**2
As well as a standard Python module bar.py
:
import numba
import foo
@numba.jit
def bar(x):
return 4 * foo.foo(x) # will this be a fast function call?
Will Numba recognize foo.foo
as a C-callable function automatically or do I need to tell it manually by, say, setting up a CFFI wrapper?
EDIT: Upon further reflection, Cython functions are just standard "builtin" functions from the view of the Python interpreter. So the question can be made more general: does Numba optimize calls to builtin functions and methods to bypass the Python calling overhead?
Upvotes: 10
Views: 3537
Reputation: 34337
It is possible to use Cython's cpdef
/cdef
-functions (but not the def
-functions) in nopython-numba:
cdef
/cpdef
function must be marked as api
in the Cython code.numba.extending.get_cython_function_address
can be used to get the address of the cpdef-function.ctypes
can be used to create a CFunction
from the address of the cpdef-function, which can be used in numba-nopython code.Read on for a more detailed explanation.
Even if the builtin-functions (PyCFunction
, the same as Cython's def
-functions) are written in C, they don't have a signature which could be used by nopython-numba-code.
For example the acos
function from the math
-module, doesn't have the signature
`double acos(double)`
as one could expect, but its signature is
static PyObject * math_acos(PyObject *self, PyObject *args)
So basically in order to call this function, numba would need to build a Python-float from the C-float at hand, but this is prohibited by nopython=True
.
However, Cythons cpdef
-functions are a little bit different: it is a small wrapper around a real cdef
-function, whose arguments are raw C-types like double
, int
and so on. This cdef
-function could be used by numba, only if its address were known.
Cython offers a way to find out the addresses of cdef
-functions in a portable way: the addresses can be found in the attribute __pyx_capi__
of the cythonized module.
However, not all cdef
and cpdef
functions are exposed in this way, but only ones which are explicitly marked as C-api declarations or implicitly by being shared through a pxd
-file.
Once the function foo
of the foomodule
is marked as api
:
cpdef api double foo(double x):
return x*x
the address of the cpdef-function foo
can be found in foomodule.__pyx_capi__
-dictionary:
import foomodule
foomodule.__pyx_capi
# {'foo': <capsule object "double (double)" at 0x7fe0a46f0360>}
It is surprisingly hard to extract the address from a PyCapsule
in Python. One possibility is to use ctypes.pythonapi
, another (maybe easier one) is to utilize Cython to access Python's C-API:
%%cython
from cpython.pycapsule cimport PyCapsule_GetPointer, PyCapsule_GetName
def address_from_capsule(object capsule):
name = PyCapsule_GetName(capsule)
return <unsigned long long int> PyCapsule_GetPointer(capsule, name)
which can be used as:
addr = address_from_capsule(foomodule.__pyx_capi__['foo'])
However, numba offers a similar functionality out of the box - get_cython_function_address
:
from numba.extending import get_cython_function_address
addr = get_cython_function_address("foomodule", "foo")
Once we got the address of the c-function, we can construct a ctypes
-function:
import ctypes
foo_functype = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)
foo_for_numba = foo_functype(addr)
This function can be utilized for example as follows from nopython-numba:
from numba import njit
@njit
def use_foo(x):
return foo_for_numba(x)
and now:
use_foo(5)
# 25.0
yields the expected result.
Arguments of ctype-functions which are understood by numba out of the box can be found here and are:
_FROM_CTYPES = {
ctypes.c_bool: types.boolean,
ctypes.c_int8: types.int8,
ctypes.c_int16: types.int16,
ctypes.c_int32: types.int32,
ctypes.c_int64: types.int64,
ctypes.c_uint8: types.uint8,
ctypes.c_uint16: types.uint16,
ctypes.c_uint32: types.uint32,
ctypes.c_uint64: types.uint64,
ctypes.c_float: types.float32,
ctypes.c_double: types.float64,
ctypes.c_void_p: types.voidptr,
ctypes.py_object: types.ffi_forced_object,
}
Upvotes: 12
Reputation: 68682
There is a limited set of builtin functions (from both the python standard library and numpy) that numba knows how to translate into native code:
Anything else will not be able to be jitted by Numba in nopython
mode, thus resorting to objectmode
which is much slower.
There is no direct way to pass a cython function to Numba and have it be recognized in nopython
mode. Numba does have hooks for cffi:
http://numba.pydata.org/numba-doc/latest/reference/pysupported.html#cffi
that can be leveraged to call outside C code, that you might be able to rig up to call cython if you could create a low level wrapper at the C-level; I'm not 100% sure if this is possible though. I wrote about doing this for calling RMath functions from Numba:
It might be helpful in getting you started if you go that route.
Upvotes: 5