numba-cfunc

参考文档: https://numba.pydata.org/numba-doc/latest/user/cfunc.html

cfunc

创建 C/C++回调函数
与 jit 相似，有一个不同点是，cfunc 强制传递一个签名，用来确定 C 回调的可见签名。
cfunc 对象暴漏出编译后的 C 回调地址，以便可以传递给任何的外部 C/C++ 库。他还暴漏了一个 ctypes 回调对象指向该回调函数，这个对象也可以在 python 中调用，方便检查编译后的代码。

import numba as nb


@nb.cfunc("int32(int32, int32)")
def add(x, y):
    return x + y


print("add.ctypes(1, 2): ", add.ctypes(1, 2))

处理指针和数组内存

C 回调的一个不太简单的用例涉及对某些 调用方传递的数据数组。由于 C 没有高级 类似于 Numpy 数组的抽象，C 回调的签名将通过 低级指针和大小参数。尽管如此，Python 代码 回调将期望利用 Numpy 的力量和表现力 阵 列。

在以下示例中，C 回调应在二维数组上运行， 带有签名 . 您可以这样实现这样的回调：void(double *input, double *output, int m, int n)

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


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

完整代码

import numba as nb


@nb.cfunc("int32(int32, int32)")
def add(x, y):
    return x + y


print("add.ctypes(1, 2): ", add.ctypes(1, 2))


print("-" * 20)

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


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