Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

For example, here is a simple c function that use pointer to return value:

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

I want to call add() function for every elements in two arrays, and collect the result by numba @jit function.

Compile the c code first:

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

And load it by 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

then the numba function:

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

But numba has no addressof operator or function.

Currently I am using the following method. But this method can't be used for nopython mode, and I don't know whether the code in the for-loop has python object or not.

@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
like image 874
HYRY Avatar asked May 31 '17 13:05

HYRY


1 Answers

I couldn't find a precise way to pass the reference of the exact element, but the following appears to work:

@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

Basically, you can get the data pointer of z using z.ctypes.data, but that just gives you the first element. It feels hack-ish, but basically I just take a slice so the memory address I want is at the start of the slice.

I'm not sure if there is a better alternative.

like image 61
JoshAdel Avatar answered Oct 13 '22 06:10

JoshAdel