I'd like to create a Numba-compiled function that takes a pointer or the memory address of an array as an argument and does calculations on it, e.g., modifies the underlying data.
The pure-python version to illustrate this looks like this:
import ctypes
import numba as nb
import numpy as np
arr = np.arange(5).astype(np.double) # create arbitrary numpy array
def modify_data(addr):
""" a function taking the memory address of an array to modify it """
ptr = ctypes.c_void_p(addr)
data = nb.carray(ptr, arr.shape, dtype=arr.dtype)
data += 2
addr = arr.ctypes.data
modify_data(addr)
arr
# >>> array([2., 3., 4., 5., 6.])
As you can see in the example, the array arr
got modified without passing it to the function explicitly. In my use case, the shape and dtype of the array are known and will remain unchanged at all times, which should simplify the interface.
1. Attempt: Naive jitting
I now tried to compile the modify_data
function, but failed. My first attempt was to use
shape = arr.shape
dtype = arr.dtype
@nb.njit
def modify_data_nb(ptr):
data = nb.carray(ptr, shape, dtype=dtype)
data += 2
ptr = ctypes.c_void_p(addr)
modify_data_nb(ptr) # <<< error
This failed with cannot determine Numba type of <class 'ctypes.c_void_p'>
, i.e., it does not know how to interpret the pointer.
2. Attempt: Explicit types
I tried putting explicit types,
arr_ptr_type = nb.types.CPointer(nb.float64)
shape = arr.shape
@nb.njit(nb.types.void(arr_ptr_type))
def modify_data_nb(ptr):
""" a function taking the memory address of an array to modify it """
data = nb.carray(ptr, shape)
data += 2
but this did not help. It did not throw any errors, but I do not know how to call the function modify_data_nb
. I tried the following options
modify_data_nb(arr.ctypes.data)
# TypeError: No matching definition for argument type(s) int64
ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
modify_data_nb(ptr)
# TypeError: No matching definition for argument type(s) pyobject
ptr = ctypes.c_void_p(arr.ctypes.data)
modify_data_nb(ptr)
# TypeError: No matching definition for argument type(s) pyobject
Is there a way to obtain the correct pointer format from arr
so I can pass it to the Numba-compiled modify_data_nb
function? Alternatively, is there another way of passing the memory location to function.
3. Attempt: Using scipy.LowLevelCallable
I made some progress by using scipy.LowLevelCallable
and its magic:
arr = np.arange(3).astype(np.double)
print(arr)
# >>> array([0., 1., 2.])
# create the function taking a pointer
shape = arr.shape
dtype = arr.dtype
@nb.cfunc(nb.types.void(nb.types.CPointer(nb.types.double)))
def modify_data(ptr):
data = nb.carray(ptr, shape, dtype=dtype)
data += 2
modify_data_llc = LowLevelCallable(modify_data.ctypes).function
# create pointer to array
ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
# call the function only with the pointer
modify_data_llc(ptr)
# check whether array got modified
print(arr)
# >>> array([2., 3., 4.])
I can now call a function to access the array, but this function is no longer a Numba function. In particular, it cannot be used in other Numba functions.
See Question&Answers more detail:
os