Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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?

like image 651
cfh Avatar asked Mar 25 '17 12:03

cfh


People also ask

Should I use Cython or Numba?

Both Cython and Numba speeds up Python code even small number of operations. More the number of operations more is the speed up. However, performance gain by Cython saturates at around 100-150 times of Python. On the other hand, speed up gain by Numba increases steadily with number of operations.

Does Numba work with CPU?

Numba offers a range of options for parallelizing your code for CPUs and GPUs, often with only minor code changes.

What is Nopython in Numba?

The behaviour of the nopython compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement of the Python interpreter. This is the recommended and best-practice way to use the Numba jit decorator as it leads to the best performance.


2 Answers

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 explanaton.


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.

like image 129
ead Avatar answered Sep 28 '22 11:09

ead


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:

  • http://numba.pydata.org/numba-doc/latest/reference/pysupported.html
  • http://numba.pydata.org/numba-doc/latest/reference/numpysupported.html

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.

like image 42
JoshAdel Avatar answered Sep 28 '22 09:09

JoshAdel