如何将 C 函数编译为 numpy ufunc 并动态加载它？

Question

我有一些自动生成 C 函数的 Python 代码。该函数接受一些双精度数作为输入并返回一个双精度数，同时调用 C 标准库中的各种函数。

我想做的一件事是将其编译成 numpy ufunc 并将其加载到正在运行的 Python 进程中。我只希望该函数以合理的速度在其输入 numpy 数组（例如 numpy 的数组）上按元素运行minimum。

我很惊讶我找不到如何执行此操作的明确说明或示例。Numpy 有关于编写扩展的明确说明，但不清楚如何将它们加载到当前的 Python 进程中。使用 ctypes，我可以编译我的函数并加载它，没问题，但不清楚如何使其成为 ufunc 而不是普通的 Python 函数。Cython 也可以做到这一点，如果我使用pyximport它，甚至可以为我构建共享库，这是理想的，因为这样我就可以分发它，而不必担心如何在另一个系统上构建 C 代码。但同样不清楚如何制作 ufunc 而不是普通函数。

TL;DR：如何将一个简单的 C 函数编译成 ufunc，并动态加载它？越简单、样板越少越好。

Answer 1

一种想法可能是使用 numba创建 ufunc，使用cffi编译 C 代码。

例如，如果我们想要将 numpy 数组中每个元素的值加倍，即将以下 C 函数作为字符串：

double f(double a){
    return 2.0*a;
}

一个可能的解决方案是以下原型：

import numba as nb
import cffi

def create_ufunc(code):
    # 1. step: compile the C-code and load the resulting extension
    ffibuilder = cffi.FFI()
    ffibuilder.cdef("double f(double a);", override=True)
    built_module=ffibuilder.verify(source=code)
    fun = built_module.f

    # 2. step: create an ufunc out of the compiled C-function
    @nb.vectorize([nb.float64(nb.float64)])
    def f(x):
      return fun(x)
    return f

现在：

import numpy as np
a=np.arange(6).astype(np.float64)
my_f1=create_ufunc("double f(double a){return 2.0*a;}")
my_f1(a)
# array([  0.,   2.,   4.,   6.,   8.,  10.])

或者如果我们想乘以10.0：

my_f2=create_ufunc("double f(double a){return 10.0*a;}")
# array([  0.,  10.,  20.,  30.,  40.,  50.])

显然，在展示可能性的同时，这个原型还需要一些改进。例如，虽然紧凑，但verify已被弃用，并且使用相同的代码调用create_ufunc两次将导致警告。

另一个问题：尽管numba 支持cffi 函数，但上面的版本无法在 nopython 模式下编译。不知道这里出了什么问题？请参阅下面的解决方法：一个以 nopython 模式构建的更复杂的版本。

然而，这可能仍然是一个很好的起点。

如果我们使用out-of-line ( compile) 而不是in-line ( verify) API 模式，似乎可以在 nopython 模式下编译 numba ：

import numba as nb
import cffi
import zlib
import importlib
import numba.cffi_support as nbcffi

def create_ufunc(code):
    # 1. step: compile the C-code and load the resulting extension
    # create a different so/dll for different codes
    # and load it
    module_name="myufunc"+str(zlib.adler32(code.encode('ascii')))
    ffibuilder = cffi.FFI()
    ffibuilder.cdef("double f(double a);", override=True)
    ffibuilder.set_source(module_name=module_name,source=code)
    ffibuilder.compile(verbose=True)
    loaded = importlib.import_module(module_name)


    # 2. step: create an ufunc out of the compiled C-function
    # out-of-line modules must be registered in numba:      
    nbcffi.register_module(loaded)
    fun = loaded.lib.f

    @nb.vectorize([nb.float64(nb.float64)], nopython=True)
    def f(x):
      return fun(x)
    return f

重要细节：

• 每个code. 我们通过传递的哈希值来区分它们code。
• 随着时间的推移，将会出现相当多的myufuncXXXX.so- 文件，人们可以考虑实现一个类似于cffi.verify.
• ffibuilder.compile(verbose=True)仅用于调试目的，可能verbose=False在发布时更有意义。