从Numba jitted代码调用Cython函数

Question

我知道调用另一个jitted函数的Numba-jitted函数会识别这个并自动使用快速C调用约定而不是通过Python对象层,因此避免了高Python函数调用开销:

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

我的问题是,如果我从Numba调用Cython函数,是否也是如此.所以假设我有一个Cython模块,foo.pyx:

cpdef double foo(double x):
    return x**2

以及标准的Python模块bar.py:

import numba
import foo

@numba.jit
def bar(x):
    return 4 * foo.foo(x)   # will this be a fast function call?

Numba会自动识别foo.foo为C可调用函数,还是需要通过设置CFFI包装器手动告诉它？

编辑:进一步反思,从Python解释器的角度来看,Cython函数只是标准的"内置"函数.所以这个问题可以更加通用:Numba是否优化对内置函数和方法的调用以绕过Python调用开销？

Answer 1

可以在 nopython-numba 中使用 Cython 的cpdef/ -functions cdef（但不是def-functions）：

1. step: cdef/cpdef函数必须在 Cython 代码中标记api。
2. step：numba.extending.get_cython_function_address可以用来获取cpdef函数的地址。
3. step：ctypes可以用来CFunction从cpdef-function的地址创建一个，可以在numba-nopython代码中使用。

继续阅读以获得更详细的解释。

即使内置函数（PyCFunction，与 Cython 的def函数相同）是用 C 编写的，它们也没有可供 nopython-numba-code 使用的签名。

例如，-module 中的acos函数math没有签名

`double acos(double)`

正如人们所料，但它的签名是

static PyObject * math_acos(PyObject *self, PyObject *args)

所以基本上为了调用这个函数，numba 需要从手头的 C-float 构建一个 Python-float，但这被nopython=True.

然而，Cythons -functionscpdef有点不同：它是一个真正的-function的小包装cdef，其参数是原始的 C 类型，如double，int等等。cdef只有在知道地址的情况下，numba 才能使用此- 函数。

Cython 提供了一种以cdef可移植的方式找出 -functions 的地址的方法：地址可以__pyx_capi__在 cythonized 模块的属性中找到。

然而，并不是所有的cdef和cpdef函数都以这种方式公开，只有那些明确标记为C-api 声明或通过pxd-file隐式共享的函数。

一旦功能foo的foomodule标记为api：

cpdef api double foo(double x):
    return x*x

cpdef 函数的地址foo可以在foomodule.__pyx_capi__-dictionary 中找到：

import foomodule
foomodule.__pyx_capi
# {'foo': <capsule object "double (double)" at 0x7fe0a46f0360>}

从PyCapsulePython 中的a中提取地址非常困难。一种可能性是使用ctypes.pythonapi，另一种（可能更简单）是利用 Cython 来访问 Python 的 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)

可以用作：

addr = address_from_capsule(foomodule.__pyx_capi__['foo'])

但是，numba 提供了一个开箱即用的类似功能 - get_cython_function_address：

from numba.extending import get_cython_function_address
addr = get_cython_function_address("foomodule", "foo")

一旦我们得到了 c-function 的地址，我们就可以构造一个ctypes-function：

import ctypes
foo_functype = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)
foo_for_numba = foo_functype(addr)

例如，可以从 nopython-numba 使用此函数，如下所示：

from numba import njit
@njit
def use_foo(x):
    return foo_for_numba(x)

现在：

use_foo(5)
# 25.0

产生预期的结果。

Answer 2

numba 知道如何将一组有限的内置函数（来自 python 标准库和 numpy）转换为本机代码：

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

其他任何东西都无法在nopython模式下被 Numba 处理，因此诉诸objectmode其中的速度要慢得多。

没有直接的方法可以将 cython 函数传递给 Numba 并使其在nopython模式下被识别。Numba 确实有 cffi 的钩子：

http://numba.pydata.org/numba-doc/latest/reference/pysupported.html#cffi

可以利用它来调用外部 C 代码，如果您可以在 C 级别创建一个低级包装器，您可能能够调用 cython；不过，我不是 100% 确定这是否可能。我写过这样做是为了从 Numba 调用 RMath 函数：

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

如果你走那条路，它可能会帮助你开始。