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

时间:2018-10-23 13:16:20

标签: python numpy cython ctypes

我有一些Python代码可以自动生成C函数。此函数将一些double作为输入,并返回double,并从中调用C标准库中的各种函数。

我要执行的操作之一是将其编译为numpy ufunc并将其加载到正在运行的Python进程中。我只希望函数以合理的速度在其输入的numpy数组(例如numpy的minimum)上按元素运行。

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

TL; DR:如何使用简单的C函数,将其编译为ufunc,然后动态加载?越简单,样板越少越好。

1 个答案:

答案 0 :(得分:3)

一个想法可能是将numba用于creating ufuncscffi来编译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会导致警告。

另一个问题:尽管cffi-functions are supported by numba确实存在,但以上版本无法在nopython模式下编译。不知道这里出了什么问题?解决方法请参见以下内容:在nopython模式下构建的更复杂的版本。

但是,这可能仍然是一个很好的起点。


如果我们使用out-of-linecompile)而不是in-lineverify)API模式,则似乎可以在nopython模式下编译numba: / p>

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都有一个新的扩展名(so / pyd文件)。我们通过传递的code的哈希值来区分它们。
  • 随着时间的流逝,将会有很多myufuncXXXX.so文件,人们可以考虑实现类似于cffi.verify所使用的基础结构。
  • ffibuilder.compile(verbose=True)仅用于调试目的,可能verbose=False在发行中更有意义。