How can I compile a C function into a numpy ufunc and load it dynamically?

user_1119340 0
Viewed 507

I have some Python code that automatically generates a C function. This function takes some doubles as input and returns a double, calling various functions from the C standard library along the way.

One of the things I would like to do with this is compile it into a numpy ufunc and load it into the running Python process. I just want the function to run element-wise on its input numpy arrays, like numpy's minimum for example, at reasonable speed.

I was surprised that I couldn't find clear instructions or examples how to do this. Numpy has clear instructions on writing extensions, but it's not clear how I could load these into the current Python process. With ctypes I can compile my function and load it, no problem, but it's not clear how to make it a ufunc rather than a normal Python function. Cython can also do this, and if I use pyximport it will even build the shared library for me, which is ideal because then I can distribute it without worrying about how to build the C code on another system. But again it's not clear how to make a ufunc rather than a normal function.

TL;DR: how can I take a simple C function, compile it into a ufunc, and load it dynamically? The more foolproof and less boilerplate the better.

cython ctypes numpy python
1 Answers

One idea could be to use numba for creating ufuncs and cffi for compiling the c-code.

For example if we want to double the value of every element in a numpy-array, i.e. having the following C-function as a string:

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

a possible solution is the following prototype:

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

And now:

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.])

or if we want to multiply with 10.0:

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

Obviosly, while showing what is possible, this prototype needs some polishing. For example albeit compact, verify is deprecated and calling create_ufunc twice with the same code will lead to a warning.

Another issue: the version above does not compile in the nopython-mode, despite the fact that cffi-functions are supported by numba. Not sure what is going wrong here? See further below for a workaround: a more complicated version which builds in nopython mode.

However, this is probably still a good starting point.

It seems to be possible to compile numba in nopython-mode, if we use out-of-line (compile) instead of in-line (verify) API-mode:

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

Important details:

  • There is a new extension (so/pyd-file) for every code. We distinguish between them via hash-value of the passed code.
  • over the time there will be quite some myufuncXXXX.so-files around, one could think about implementing an infrastructure similar to one used by cffi.verify.
  • ffibuilder.compile(verbose=True) is just for debugging purposes, probably verbose=False makes more sense in release.
user_57694630
Related