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没有运营商或职能的地址。

目前我正在使用以下方法。但是这种方法不能用于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的数据指针，但这只是给你第一个元素。它感觉很糟糕，但基本上我只需要一个切片，所以我想要的内存地址就在切片的开头。

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

如何在Numba @jit

1 个答案: