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

Question

例如，这是一个简单的 c 函数，它使用指向 return 值的指针：

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-loop中的代码是否有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.ctypes.data 获取 z 的数据指针，但这只会给您第一个元素。感觉有点乱，但基本上我只取一片，所以我想要的内存地址在片的开头。

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

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

How to call ctypes functions that use pointer to return value in Numba @jit

python

ctypes

numba