如何在 Numba @jit 中调用使用指针返回值的 ctypes 函数

Question

例如，这是一个使用指针返回值的简单 c 函数：

void add(double x, double y, double *r)
{
    *r = x + y;
}

我想为add()两个数组中的每个元素调用函数，并通过 numba @jit 函数收集结果。

先编译c代码：

!gcc -c -fpic func.c
!gcc -shared -o func.so func.o

并通过 ctypes 加载它：

lib = ctypes.cdll.LoadLibrary("./func.so")
add = lib.add
add.argtypes = ctypes.c_double, ctypes.c_double, ctypes.c_void_p
add.restype = None

然后是 numba 函数：

from numba import jit, float64

@jit(float64(float64[:], float64[:]))
def f(x, y):
    z = np.zeros_like(x)
    for i in range(x.shape[0]):
        add(x[i], y[i], &z[i]) # here I want to pass the address of z[i]
    return z

但是 numba 没有 addressof 运算符或函数。

目前我正在使用以下方法。但是这个方法不能用于nopython模式，也不知道for循环中的代码有没有python对象。

@jit(float64(float64[:], float64[:]))
def f(x, y):
    z = np.zeros_like(x)
    tmp = ctypes.c_double(0.0)
    addr = intp(ctypes.addressof(tmp))
    val = carray(ctypes.pointer(tmp), 1)
    for i in range(x.shape[0]):
        add(x[i], y[i], addr)
        z[i] = val[0]
    return z

Answer 1

我找不到传递确切元素的引用的精确方法，但以下似乎有效：

@nb.jit(nb.float64[:](nb.float64[:], nb.float64[:]))
def f(x, y):
    z = np.zeros_like(x)
    for i in range(x.shape[0]):
        add(x[i], y[i], z[i:].ctypes.data) # here I want to pass the address of z[i]
    return z

z基本上，您可以获取使用的数据指针z.ctypes.data，但这只是给您第一个元素。这感觉很黑客，但基本上我只是取一个切片，所以我想要的内存地址位于切片的开头。

我不确定是否有更好的选择。