cfh
cfh

Reputation: 4666

Calling Cython functions from Numba jitted code

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

Answers (2)

ead
ead

Reputation: 34337

It is possible to use Cython's cpdef/cdef-functions (but not the def-functions) in nopython-numba:

  1. step: cdef/cpdef function must be marked as api in the Cython code.
  2. step: numba.extending.get_cython_function_address can be used to get the address of the cpdef-function.
  3. step: 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

JoshAdel
JoshAdel

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:

https://web.archive.org/web/20160611082327/https://www.continuum.io/blog/developer-blog/calling-c-libraries-numba-using-cffi

It might be helpful in getting you started if you go that route.

Upvotes: 5

Related Questions