如何使用指针在 Numba 中包装 CFFI 函数 [英] How to wrap a CFFI function in Numba taking Pointers
问题描述
这应该是一项简单的任务,但我找不到如何在 Numba 函数中将标量值的指针传递给 CFFI 函数的方法.使用 ffi.from_buffer
传递指向数组的指针没有问题.
It should be a easy task, but I can't find a way how to pass a pointer of a scalar value to a CFFI function within a Numba function. Passing a pointer to an array works without problems using ffi.from_buffer
.
示例函数
import cffi
ffi = cffi.FFI()
defs="void foo_f(int a,double *b);"
ffi.cdef(defs, override=True)
source="""
#include <stdio.h>;
void foo_f(int a,double *b){
printf("%i",a);
printf(" ");
printf("%f",b[0]);
}
"""
ffi.set_source(module_name="foo",source=source)
ffi.compile()
传递指向数组的指针
import numpy as np
import numba as nb
import cffi
ffi = cffi.FFI()
import numpy as np
import ctypes
import foo
nb.cffi_support.register_module(foo)
foo_f = foo.lib.foo_f
@nb.njit()
def Test(a,b):
a_wrap=np.int32(a)
#This works for an array
b_wrap=ffi.from_buffer(b.astype(np.float64))
foo_f(a_wrap,b_wrap)
a=64.
b=np.ones(5)
Test(a,b)
这没有问题,但如何修改 Test
函数以获取标量值 b=5.
而不修改 CFFI 函数本身?
This works without problems, but how can I modify the Test
function to take a scalar value b=5.
without modifying the CFFI-function itself?
推荐答案
使用 Numba 通过引用传递标量值
为了获得有用的时间,我稍微修改了包装函数.该函数只是将一个标量(通过值传递)添加到一个标量 b(通过引用传递).
Pass scalar values by reference using Numba
To get useful timings I have modified the wrapped function a bit. The function simply adds a scalar (passed by value) to a scalar b (passed by reference).
使用内在函数的方法的优缺点
- 仅在 nopython 模式下工作
- 运行时间较短的 C 或 Fortran 函数更快(真实示例)
示例函数
import cffi
ffi = cffi.FFI()
defs="void foo_f(double a,double *b);"
ffi.cdef(defs, override=True)
source="""
void foo_f(double a,double *b){
b[0]+=a;
}
"""
ffi.set_source(module_name="foo",source=source)
ffi.compile()
使用临时数组的包装
这很简单,但需要分配一个大小为 1 的数组,这很慢.
This is quite straight forward, but requires to allocate an array of size one, which is quite slow.
import numpy as np
import numba as nb
from numba import cffi_support
import cffi
ffi = cffi.FFI()
import foo
nb.cffi_support.register_module(foo)
foo_f = foo.lib.foo_f
@nb.njit("float64(float64,float64)")
def method_using_arrays(a,b):
b_arr=np.empty(1,dtype=np.float64)
b_arr[0]=b
b_arr_ptr=b_wrap=ffi.from_buffer(b_arr)
foo_f(a,b_arr_ptr)
return b_arr[0]
使用内在函数的包装器
from numba import types
from numba.extending import intrinsic
from numba import cgutils
@intrinsic
def ptr_from_val(typingctx, data):
def impl(context, builder, signature, args):
ptr = cgutils.alloca_once_value(builder,args[0])
return ptr
sig = types.CPointer(data)(data)
return sig, impl
@intrinsic
def val_from_ptr(typingctx, data):
def impl(context, builder, signature, args):
val = builder.load(args[0])
return val
sig = data.dtype(data)
return sig, impl
@nb.njit("float64(float64,float64)")
def method_using_intrinsics(a,b):
b_ptr=ptr_from_val(b)
foo_f(a,b_ptr)
return val_from_ptr(b_ptr)
时间
#Just call the wrapped function a few times
@nb.njit()
def timing_method_using_intrinsics(a,b):
for i in range(1000):
b=method_using_intrinsics(a,b)
return b
#Just call the wrapped function a few times
@nb.njit()
def timing_method_using_arrays(a,b):
for i in range(1000):
b=method_using_arrays(a,b)
return b
a=1.
b=1.
%timeit timing_method_using_intrinsics(a,b)
#5.15 µs ± 33.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit timing_method_using_arrays(a,b)
#121 µs ± 601 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
这篇关于如何使用指针在 Numba 中包装 CFFI 函数的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!