使用 @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_type
将 cffi
类型转换为 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 结构体的指针将作为 void 指针返回。
注意
从 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 回调有意义。您通常应将自己限制为标量类型(例如 int8
或 float64
)、指向它们的指针(例如 types.CPointer(types.int8)
),或指向 Record
类型的指针。
编译选项
多个仅限关键字的参数可以传递给 @cfunc
装饰器:nopython
和 cache
。它们的含义与 @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