如何在 Numba 中使用指针包装 CFFI 函数

Question

这应该是一项简单的任务，但我找不到如何将标量值的指针传递给 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)

这可以正常工作，但如何修改函数Test以获取标量值b=5.而不修改 CFFI 函数本身？

Answer 1

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

\n

为了获得有用的计时，我对包装函数进行了一些修改。该函数只是将标量（按值传递）添加到标量 b（按引用传递）。

\n

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

\n

\n
• 仅在 nopython 模式下工作
\n
• C 或 Fortran 函数速度更快，运行时间短（真实示例）
\n

\n

示例函数

\n

import cffi\n\nffi = cffi.FFI()\ndefs="void foo_f(double a,double *b);"\nffi.cdef(defs, override=True)\nsource="""\nvoid foo_f(double a,double *b){\n  b[0]+=a;\n  }\n"""\nffi.set_source(module_name="foo",source=source)\nffi.compile()\n

\n

使用临时数组的包装器

\n

这非常简单，但需要分配一个大小为 1 的数组，这非常慢。

\n

import numpy as np\nimport numba as nb\nfrom numba import cffi_support\nimport cffi\nffi = cffi.FFI()\nimport foo\n\nnb.cffi_support.register_module(foo)\nfoo_f = foo.lib.foo_f\n\n@nb.njit("float64(float64,float64)")\ndef method_using_arrays(a,b):\n    b_arr=np.empty(1,dtype=np.float64)\n    b_arr[0]=b\n    b_arr_ptr=b_wrap=ffi.from_buffer(b_arr)\n    foo_f(a,b_arr_ptr)\n    return b_arr[0]\n

\n

使用内在函数的包装器

\n

from numba import types\nfrom numba.extending import intrinsic\nfrom numba import cgutils\n\n@intrinsic\ndef ptr_from_val(typingctx, data):\n    def impl(context, builder, signature, args):\n        ptr = cgutils.alloca_once_value(builder,args[0])\n        return ptr\n    sig = types.CPointer(data)(data)\n    return sig, impl\n\n@intrinsic\ndef val_from_ptr(typingctx, data):\n    def impl(context, builder, signature, args):\n        val = builder.load(args[0])\n        return val\n    sig = data.dtype(data)\n    return sig, impl\n\n@nb.njit("float64(float64,float64)")\ndef method_using_intrinsics(a,b):\n    b_ptr=ptr_from_val(b)\n    foo_f(a,b_ptr)\n    return val_from_ptr(b_ptr)\n

\n

时间安排

\n

#Just call the wrapped function a few times\n@nb.njit()\ndef timing_method_using_intrinsics(a,b):\n    for i in range(1000):\n        b=method_using_intrinsics(a,b)\n    return b\n\n#Just call the wrapped function a few times\n@nb.njit()\ndef timing_method_using_arrays(a,b):\n    for i in range(1000):\n        b=method_using_arrays(a,b)\n    return b\n\na=1.\nb=1.\n\n%timeit timing_method_using_intrinsics(a,b)\n#5.15 \xc2\xb5s \xc2\xb1 33.9 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n%timeit timing_method_using_arrays(a,b)\n#121 \xc2\xb5s \xc2\xb1 601 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n

\n