使用 @cfunc 创建C回调函数

与一些原生库(例如用C或C++编写的库)交互可能需要编写原生回调函数,以便向库提供业务逻辑。numba.cfunc() 装饰器创建一个可以从外部C代码调用的编译函数,使用您选择的签名。

基本用法

@cfunc 装饰器与 @jit 的使用方式相似,但有一个重要的区别:必须传递一个单一的签名。它决定了C回调的可见签名:

from numba import cfunc

@cfunc("float64(float64, float64)")
def add(x, y):
    return x + y

C 函数对象将编译后的 C 回调地址作为 address 属性公开,以便您可以将其传递给任何外部 C 或 C++ 库。它还公开了一个指向该回调的 ctypes 回调对象;该对象也可以从 Python 中调用,从而可以轻松检查编译后的代码:

@cfunc("float64(float64, float64)")
def add(x, y):
    return x + y

print(add.ctypes(4.0, 5.0))  # prints "9.0"

示例

在这个例子中,我们将使用 scipy.integrate.quad 函数。该函数接受一个常规的 Python 回调函数或一个包装在 ctypes 回调对象中的 C 回调函数。

让我们定义一个纯Python的被积函数,并将其编译为C回调:

>>> import numpy as np
>>> from numba import cfunc
>>> def integrand(t):
        return np.exp(-t) / t**2
   ...:
>>> nb_integrand = cfunc("float64(float64)")(integrand)

我们可以将 nb_integrand 对象的 ctypes 回调传递给 scipy.integrate.quad,并检查结果是否与纯 Python 函数相同:

>>> import scipy.integrate as si
>>> def do_integrate(func):
        """
        Integrate the given function from 1.0 to +inf.
        """
        return si.quad(func, 1, np.inf)
   ...:
>>> do_integrate(integrand)
(0.14849550677592208, 3.8736750296130505e-10)
>>> do_integrate(nb_integrand.ctypes)
(0.14849550677592208, 3.8736750296130505e-10)

使用编译后的回调函数,积分函数在每次评估被积函数时不会调用 Python 解释器。在我们的例子中,积分速度提高了 18 倍:

>>> %timeit do_integrate(integrand)
1000 loops, best of 3: 242 µs per loop
>>> %timeit do_integrate(nb_integrand.ctypes)
100000 loops, best of 3: 13.5 µs per loop

处理指针和数组内存

C回调函数的一个不那么简单的用例涉及对调用者传递的数据数组进行操作。由于C语言没有类似于Numpy数组的高级抽象,C回调函数的签名将传递低级指针和大小参数。尽管如此,回调函数的Python代码将期望利用Numpy数组的强大功能和表达能力。

在以下示例中,C回调函数预期操作二维数组,其签名是 void(double *input, double *output, int m, int n)。你可以这样实现这样的回调:

from numba import cfunc, types, carray

c_sig = types.void(types.CPointer(types.double),
                   types.CPointer(types.double),
                   types.intc, types.intc)

@cfunc(c_sig)
def my_callback(in_, out, m, n):
    in_array = carray(in_, (m, n))
    out_array = carray(out, (m, n))
    for i in range(m):
        for j in range(n):
            out_array[i, j] = 2 * in_array[i, j]

函数 numba.carray() 接受一个数据指针和一个形状作为输入,并返回该数据上给定形状的数组视图。假定数据按C顺序排列。如果数据按Fortran顺序排列,则应使用 numba.farray()

处理C结构体

使用 CFFI

对于有很多状态的应用程序,在C结构中传递数据是有用的。为了简化与C代码的互操作性,numba可以使用 numba.core.typing.cffi_utils.map_typecffi 类型转换为 numba 的 Record 类型:

from numba.core.typing import cffi_utils

nbtype = cffi_utils.map_type(cffi_type, use_record_dtype=True)

备注

use_record_dtype=True 是必需的,否则指向C结构的指针将作为空指针返回。

备注

从 v0.49 开始,numba.cffi_support 模块已被淘汰,取而代之的是 numba.core.typing.cffi_utils

例如:

from cffi import FFI

src = """

/* Define the C struct */
typedef struct my_struct {
   int    i1;
   float  f2;
   double d3;
   float  af4[7]; // arrays are supported
} my_struct;

/* Define a callback function */
typedef double (*my_func)(my_struct*, size_t);
"""

ffi = FFI()
ffi.cdef(src)

# Get the function signature from *my_func*
sig = cffi_utils.map_type(ffi.typeof('my_func'), use_record_dtype=True)

# Make the cfunc
from numba import cfunc, carray

@cfunc(sig)
def foo(ptr, n):
   base = carray(ptr, n)  # view pointer as an array of my_struct
   tmp = 0
   for i in range(n):
      tmp += base[i].i1 * base[i].f2 / base[i].d3
      tmp += base[i].af4.sum()  # nested arrays are like normal NumPy arrays
   return tmp

使用 numba.types.Record.make_c_struct

numba.types.Record 类型可以手动创建以遵循 C 结构的布局。为此,请使用 Record.make_c_struct,例如:

my_struct = types.Record.make_c_struct([
   # Provides a sequence of 2-tuples i.e. (name:str, type:Type)
   ('i1', types.int32),
   ('f2', types.float32),
   ('d3', types.float64),
   ('af4', types.NestedArray(dtype=types.float32, shape=(7,))),
])

由于 ABI 限制,结构体应通过指针传递,使用 types.CPointer(my_struct) 作为参数类型。在 cfunc 主体内,可以通过 carray 访问 my_struct*

完整示例

参见完整示例在 examples/notebooks/Accessing C Struct Data.ipynb

签名规范

显式的 @cfunc 签名可以使用任何 Numba 类型,但对于 C 回调来说,只有其中一部分是有意义的。通常你应该限制自己使用标量类型(如 int8float64)、指向它们的指针(例如 types.CPointer(types.int8)),或指向 Record 类型的指针。

编译选项

@cfunc 装饰器可以传递一些仅限关键字的参数:nopythoncache。它们的含义与 @jit 装饰器中的类似。

从 Numba 调用 C 代码

也可以从 Numba @jit 函数中调用 C 代码。在这个例子中,我们将编译一个简单的函数 sum ,它将两个整数相加,并在 Numba @jit 代码中调用它。

备注

下面的示例在Linux上进行了测试,并且很可能在类Unix操作系统上也能工作。

#include <stdint.h>

int64_t sum(int64_t a, int64_t b){
   return a + b;
}

使用 gcc lib.c -fPIC -shared -o shared_library.so 编译代码以生成共享库。

from numba import njit
from numba.core import types, typing
from llvmlite import binding
import os

# load the library into LLVM
path = os.path.abspath('./shared_library.so')
binding.load_library_permanently(path)

# Adds typing information
c_func_name = 'sum'
return_type = types.int64
argty = types.int64
c_sig = typing.signature(return_type, argty, argty)
c_func = types.ExternalFunction(c_func_name, c_sig)

@njit
def example(x, y):
   return c_func(x, y)

print(example(3, 4)) # 7

也可以使用 ctypes 来调用 C 函数。使用 ctypes 的优点是它不受 JIT 装饰器使用的影响。

from numba import njit
import ctypes
DSO = ctypes.CDLL('./shared_library.so')

# Add typing information
c_func = DSO.sum
c_func.restype = ctypes.c_int
c_func.argtypes = [ctypes.c_int, ctypes.c_int]

@njit
def example(x, y):
   return c_func(x, y)

print(example(3, 4)) # 7
print(example.py_func(3, 4)) # 7