python - 从 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/cdef 函数(但不是 def 函数):

1. 步骤:cdef/cpdef 函数必须是marked as api在 Cython 代码中。
2. 步骤:numba.extending.get_cython_function_address可用于获取 cpdef 函数的地址。
3. 步骤:ctypes 可用于从 cpdef 函数的地址创建一个 CFunction，可在 numba-nopython 代码中使用。

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

---

即使内置函数( PyCFunction ，与 Cython 的 def 函数相同)是用 C 语言编写的，它们也没有可供 nopython-numba- 使用的签名代码。

例如 math 模块中的 acos 函数没有签名

`double acos(double)`

正如人们所预料的那样，但是 its signature是

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

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

然而，Cythons cpdef 函数有点不同:它是一个真正的 cdef 函数的小包装器，其参数是原始 C 类型，如 double, int 等。这个 cdef 函数可以被 numba 使用，前提是它的地址是已知的。

Cython 提供了一种以可移植的方式查找 cdef 函数地址的方法:地址可以在 cython 化模块的属性 __pyx_capi__ 中找到。

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

一旦 foomodule 的函数 foo 被标记为 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>}

从 PyCapsule 中提取地址出奇地困难在 Python 中。一种可能性是使用 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 函数的地址，我们就可以构造一个 ctypes 函数:

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

产生预期的结果。

---

可以找到开箱即用的 numba 理解的 ctype 函数的参数 here并且是:

_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,
}