I thought about an old answer of mine on https://stackoverflow.com/a/58561573/4045774 where I showed a bit tedicious way how to pass args to scipy integrate using a Numba cfunc.
Here I want to put a method for discussion including
- Passing any numba compile-able function with signature double(double, Tuple_of_args_Numba_can_handle) by using a generated cfunc (everything cacheable)
- Calling a function pointer (in this case a ctypes annotated Python function) in a cachable way
Of course there is plenty room for improvement, but are there any major problems?
- Passing a pointer to a Numba tuple on the stack looks for example a bit error prone
- Inserting function pointers this way could maybe also have unwanted side effects
Example of usage
from gen_cfunc import cast_tuple_to_voidptr, gen_cfunc_with_args
from scipy_integrate import integrate
import numpy as np
import numba as nb
#Define some args
a=3
foo = np.arange(200, dtype=np.float64).reshape(2, -1)
bar = np.arange(600, dtype=np.float64).reshape(2, -1)
args = (a,foo,bar)
#Define the integrand
def function_using_arrays(x1, args):
x2,array1, array2 = args
res1 = np.interp(x1, array1[0], array1[1])
res2 = np.interp(x2, array2[0], array2[1])
return res1 + res2
#Generate cfunc
cfunc_address = gen_cfunc_with_args(function_using_arrays,args,cache=True)
#Testing, including the ability of caching
@nb.njit(cache=True)
def Test(cfunc_address,a,b,args):
return integrate(cfunc_address,a,b,args)
print(Test(cfunc_address,0.,1.,args))
print("Cache hits:", Test.stats.cache_hits)
print("Cache misses:", Test.stats.cache_misses)
#Cache hits: Counter({(int64, float64, float64, Tuple((int64, Array(float64, 2, 'C', False, aligned=True), Array(float64, 2, 'C', False, aligned=True)))): 1})
#Cache misses: Counter()
gen_cfunc.py
- Intrinsic to get a Pointer to a Numba tuple
- Generator function to get a compiled cfunc which can handle the pointer to some tuple
import numba as nb
from numba.extending import intrinsic
from numba.core import cgutils
from numba import types
import sys
@intrinsic
def cast_tuple_to_voidptr(typingctx, args):
"""
Casts a llvm struct (Numba Tuple) to a void pointer
"""
if not isinstance(args, types.Tuple):
raise TypeError("Argument must be a tuple")
#Set return Type
sig = types.voidptr(args)
def codegen(context, builder, signature, args):
[val] = args
ty = signature.args[0]
tup_ptr = cgutils.alloca_once_value(builder, val)
voidptr_ty = context.get_value_type(types.voidptr)
cast_ptr = builder.bitcast(tup_ptr, voidptr_ty)
return cast_ptr
return sig, codegen
def compile_inner(nb_func,sig,cache):
@nb.cfunc(sig,error_model="numpy",cache = cache)
def c_func(x,args_ptr):
return nb_func(x,args_ptr[0])
return c_func.address
def gen_cfunc_with_args(func,args,cache=False):
"""
Generates a c-func with signature double(double, *some_tuple) of a Numba or Python function
Useful eg. for scipy.integrate.quad low level interface
"""
if not callable(func) and not isinstance(func,nb.core.registry.CPUDispatcher):
raise TypeError("The first argument func must be a function ")
if not isinstance(args, tuple):
raise TypeError("args must be a tuple")
if isinstance(func,nb.core.registry.CPUDispatcher):
func = func.py_func
nb_func = nb.njit(types.double(types.double,nb.typeof(args)),inline="always",cache = cache)(func)
@nb.cfunc(types.double(types.double,types.CPointer(nb.typeof(args))),error_model="numpy",cache = cache)
def c_func(x,args_ptr):
return nb_func(x,args_ptr[0])
sig = types.double(types.double,types.CPointer(nb.typeof(args)))
c_func_address = compile_inner(nb_func,sig,cache)
sys.modules[__name__].c_func_address = c_func_address #For caching
return c_func_address
scipy_integrate.py
- Wrapping a ctypes annotated function to be callable in a cache-able way from Numba
- Why something like this is not natively available in Numba (cache-able)?
import ctypes
from llvmlite import ir, binding
from numba.core import cgutils
from numba import types
from numba.extending import intrinsic
import numba as nb
from numba import typeof
from scipy import LowLevelCallable
from scipy.integrate import quad
from gen_cfunc import cast_tuple_to_voidptr
#Class to hold a reference on function pointer and register function pointer, while caching is possible
#Very Preliminary
# - Automatic instrinsic generation not working (args hardcoded, no generator possible)
# - More complicated arguments like structs (even complex numbers) need platform specific handling
class CFuncRegistry:
def __init__(self):
self._cfuncs = {}
self._cfuncs_ptr = {}
self._nb_sig = {}
def register(self, name,func,restype, argtypes):
cfunc_type = ctypes.PYFUNCTYPE(restype, *argtypes)
cfunc = cfunc_type(func)
self._cfuncs[name] = cfunc
self._cfuncs_ptr[name] = ctypes.cast(cfunc, ctypes.c_void_p).value
binding.add_symbol(name, self._cfuncs_ptr[name])
self._nb_sig[name] = typeof(self._cfuncs[name]).get_call_signatures()[0][0]
return self
registry = CFuncRegistry()
#Definition of Python function and intrinsic which calls it
############################################
def scipy_quad_with_args_simple(func_address, a, b, user_data_address):
"""
func -> address to func
args -> address to args
"""
#generate Cfunc
FUNC = ctypes.CFUNCTYPE(ctypes.c_double,ctypes.c_double,ctypes.c_void_p)
cfunc_ctypes = FUNC(func_address)
low_level=LowLevelCallable(cfunc_ctypes,ctypes.c_void_p(user_data_address))
return quad(low_level, a, b)[0]
name = registry.register("scipy_quad_with_args_simple", scipy_quad_with_args_simple,
ctypes.c_double,[ctypes.c_int64,ctypes.c_double, ctypes.c_double, ctypes.c_void_p])
@intrinsic(inline="always")
def nb_scipy_quad_with_args_simple(typingctx, func_address, a, b, user_data_address):
sig = types.float64(func_address,a, b,user_data_address)
def codegen(context, builder, signature, args):
arg_types = [context.get_value_type(t) for t in signature.args]
return_type = context.get_value_type(signature.return_type)
func_ty = ir.FunctionType(return_type, arg_types)
#func_ty = ir.FunctionType(ir.DoubleType(), [ir.IntType(64), ir.DoubleType(),ir.DoubleType(),ir.IntType(8).as_pointer()])
fn = cgutils.get_or_insert_function(builder.module, func_ty, name="scipy_quad_with_args_simple")
return builder.call(fn, args)
return sig, codegen
###########################################
#Wrapped simplified quad function
@nb.njit()
def integrate(func_address,a,b,args):
user_data_address = cast_tuple_to_voidptr(args)
res = nb_scipy_quad_with_args_simple(func_address, a, b, user_data_address)
return res