如何包装一个CFFI函数在Numba采取指针



这应该是一项简单的任务,但我找不到如何将标量值的指针传递到Numba函数中的CFFI函数。使用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)

这工作起来没有问题,但在不修改CFFI函数本身的情况下,如何修改Test函数以获取标量值b=5.

使用Numba通过引用传递标量值

为了获得有用的计时,我对包装函数进行了一些修改。该函数只是将标量(通过值传递(与标量b(通过引用传递(相加。

使用内部函数的方法的优点和缺点

  • 仅在nopyson模式下工作
  • 对于运行时间短的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的数组,这相当慢。

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)

最新更新