CFFI-ize!

Ok, so it doesn’t have as nice a ring to it as “Cythonize!” but the Python C-Foreign Function Interface blew my mind away. I had previously used Cython to write code that was a mind-altering mix of Python and C-like static type defs and so on. It was impressive how much speed you could gain by just telling the interpreter the types of a few of your variables.

Cython, however, leads to cluttered code, losing some of the charm of Python. You also have these funny .pyx and .pyd files starting to lie around. With CFFI you can write straight C code, compile it separately into a library (if you so wish: you don’t need to compile it separately if you don’t want), use it with a pure C project, have proper syntax highlighting, run it through a debugger – whatever. And then, with CFFI, you can very, very easily integrate it with your Python code, write unit tests for it and so on and so forth. The documentation page also has a simple recipe for how to write setup scripts with the c-sources.

Here’s our Mandelbrot set example again, this time rewritten via CFFI.

The plain Python version

def f(x, y, max_iter=1000):
  """
  z <- z^2 + c
  """
  c = x + y*1j
  z = (0 +0j)
  for n in range(max_iter):
    z = z**2 + c
    if abs(z) > 2:
      return 1.0 - float(n)/max_iter
  else:
    return 0.0

def mandelbrot(x0=-2.0, y0=-1.5, x1=1.0, y1=1.5, grid=50, max_iter=1000):
  dx, dy = float(x1 - x0)/grid, float(y1 - y0)/grid
  return [[f(x0 + n * dx, y0 + m * dy, max_iter) for n in range(grid)] for m in range(grid)]

def quick_plot(grid=100, max_iter=10):
  import pylab
  pylab.imshow(mandelbrot(grid=grid, max_iter=max_iter), cmap=pylab.cm.gray)
  pylab.show()

The CFFI-zed version

import numpy
from cffi import FFI

ffi = FFI()
ffi.cdef("""
float f(float, float, int);
void mandelbrot(float, float, float, float, int, int, float*);
""")
lib = ffi.verify("""#include <fmandel.h>""", include_dirs=['./'], sources=['fmandel.c'])

def mandelbrot(x0=-2.0, y0=-1.5, x1=1.0, y1=1.5, grid=50, max_iter=1000):
  g = numpy.empty((grid, grid), dtype=numpy.float32)
  lib.mandelbrot(x0, y0, x1, y1, grid, max_iter, ffi.cast("float*", g.ctypes.data))
  return g

def quick_plot(x0=-2.0, y0=-1.5, x1=1.0, y1=1.5, grid=100, max_iter=10):
  import pylab
  pylab.imshow(mandelbrot(x0, y0, x1, y1, grid=grid, max_iter=max_iter), cmap=pylab.cm.gray)
  pylab.show()

Where fmandel.h is:

#include <stdlib.h>

float f(float, float, int);
void mandelbrot(float, float, float, float, int, int, float*);

and fmandel.c is:

#include <stdlib.h>

float f(float x, float y, int max_iter) {
  float cr = x, ci = y, zr = 0, zi = 0;
  int n;
  for(n = 0; n < max_iter; n++) {
    float tzr = zr, tzi = zi;
    zr = tzr*tzr - tzi*tzi + cr;
    zi = 2*tzr*tzi + ci;
    if(zr*zr + zi*zi > 4.0) break;
  }
  return 1.0 - (float)n/(float)max_iter;
}

void mandelbrot(float x0, float y0, float x1, float y1, int grid, int max_iter, float* mv) {
  float dx = (x1 - x0)/(float)grid, dy = (y1 - y0)/(float)grid;
  for(int n=0; n < grid; n++)
    for(int m=0; m < grid; m++)
      mv[n * grid + m] = f(x0 + (float)m * dx, y0 + (float)n * dy, max_iter);
}

The only tricky thing here, which is not well explained in the CFFI docs, is passing the array data between Python and C.

In the first version of the code I had a malloc allocate a 2-D block of memory in the C-function, and then return this 2D pointer. CFFI has an excellent interface that understands 2D pointers, however this is not compatible with matplotlib. I was first copying over the CFFI Cdata pointer element by element into a numpy array, but that was expensive. In addition the malloc with no corresponding free of the memory is a bug.

I finally decided to let numpy handle the memory management: I first allocate memory using numpy.empty which is fast (and does not waste time initializing). I then use numpy.ctypes.data to access the underlying C-pointer to this data, cast it to a float* using ffi.cast.

The one annoyance with CFFI is that it caches the compiled c-code but does not know when the c-sources have been altered. Only if you inline the C-source will CFFI recompile the inlined code. I had to resort to repeatedly deleting the __pycache__ directory each time I change the C-source.

mandel_whole