{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": true }, "outputs": [], "source": [ "from IPython.display import display, HTML" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Calling C Libraries from Numba Using CFFI\n", "\n", "Joshua L. Adelman\n", "11 February 2016\n", "\n", "**TL;DR** - The python CFFI library provides an easy and efficient way to call C code from within a function jitted (just-in-time compiled) by Numba. This makes it simple to produce fast code with functionality that is not yet available directly in Numba. As a simple demonstration, I wrap several statistical functions from the Rmath library. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Background and Motivation\n", "\n", "A large fraction of the code that I write has a performance requirement attached to it. Either I'm churning through a large amount of data in an analysis pipeline, or it is part of a real-time system and needs to complete a specific calculation in a constrained amount of time. Sometimes I can rely on numpy, pandas or existing Python packages that wrap C or Fortran code under the covers to get sufficient performance. Often times though, I'm dealing with algorithms that are difficult to implement efficiently using these tools.\n", "\n", "Since I started coding primarily in Python ~6 years ago, in those instances I'd typically reach for [Cython](http://cython.org/) to either wrap something I or others wrote in C/C++/Fortran or to provide sufficient type information to my code so that Cython could generate a performant C-extension that I could call from Python. Although Cython has been a pretty rock solid solution for me, the amount of boilerplate often required and some of the strange semantic of mixing python and low-level C code often feels less than ideal. I also collaborate with people who know Python, but don't have backgrounds in C and/or haven't had enough experience with Cython to understand how it all fits together. \n", "\n", "More and more frequently, I find myself using [Numba](http://numba.pydata.org/) in instances that I had traditionally used Cython. In short, through a simple decorator mechanism, Numba converts a subset of Python code into efficient machine code using LLVM. It uses type inference so you don't have to specify the type of every variable in a function like you do in Cython to generate fast code. This subset primarily deals with numerical code operating on scalars or Numpy arrays, but that covers 95% of the cases where I need efficient code so it does not feel that limiting. That said, the most common mistake I see people making with Numba is trying to use it as a general Python compiler and then being confused/disappointed when it doesn't speed up their code. The library has matured incredibly over the last 6-12 months to the point where at work we have it deployed in a couple of critical pieces of production code. When I first seriously prototyped it maybe a year and a half ago, it was super buggy and missing a number of key features (e.g. caching of jitted functions, memory management of numpy arrays, etc). But now it feels stable and I rarely run into problems, although I've written a very extensive unit test suite for every bit of code that it touches.\n", "\n", "One of the limitations that I do encounter semi-regularly though is when I need some specialized function that is available in Numpy or Scipy, but that function has not been re-implemented in the Numba core library so it can be called in the so-called [\"nopython\" mode](http://numba.pydata.org/numba-doc/0.23.1/glossary.html#term-nopython-mode). Basically this means that if you want to call one of these functions, you have to go through Numba's [object mode](http://numba.pydata.org/numba-doc/0.23.1/glossary.html#term-object-mode), which typically cannot generate nearly as efficient code.\n", "\n", "While there is a [proposal](http://numba.pydata.org/numba-doc/dev/proposals/extension-points.html) [under development](http://numba.pydata.org/numba-doc/dev/extending/index.html) that should allow external libraries to define an interface to make usable in nopython mode, it is not complete and will them require adoption within the larger Scipy/PyData communities. I'm looking forward to that day, but currently you have to choose a different option. The first is to re-implement a function yourself using Numba. This is often possible for functionality that is small and limited in scope, but for anything non-trivial this approach can rapidly become untenable. \n", "\n", "In the remainder of this notebook, I'm going to describe a second technique that involves using [CFFI](https://cffi.readthedocs.org) to call external C code directly from within Numba jitted code. This turns out to be a really great solution if the functionality you want has already been written either in C or a language with a C interface. It is mentioned [in the Numba docs](http://numba.pydata.org/numba-doc/0.23.1/reference/pysupported.html#cffi), but there aren't any examples that I have seen, and looking at the tests only helped a little. \n", "\n", "I had not used CFFI before integrating it with Numba for a recent project. I had largely overlooked it for two reasons: (1) Cython covered the basic usecase of exposing external C code to python and I was already very comfortable with Cython, and (2) I had the (incorrect) impression that CFFI was mostly useful in the PyPy ecosystem. Since PyPy is a non-starter for all of my projects, I largely just ignored its existence. I'm thankfully correcting that mistake now. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Rmath. It's not just for R\n", "\n", "Every once in a while I fire up R, usually through rpy2, to do something that I can't do using Statsmodel or Scikit-Learn. But for the most part I live squarely in the Python world, and my experience with R is rudimentary. So it wasn't totally surprising that I only recently discovered that the math library that underpins R, Rmath, can be built in a standalone mode without invoking R at all. In fact, the [Julia](http://julialang.org/) programming language uses Rmath for its probability distributions library and maintains a fork of the package called [Rmath-julia](https://github.com/JuliaLang/Rmath-julia).\n", "\n", "Discovering Rmath over the summer, led to the following tweet (apologies for the Jupyter Notebook input cell) and a horrific amalgamation of code that worked, but was pretty difficult to maintain and extend:" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false }, "outputs": [ { "data": { "text/html": [ "

Today's \n", "the sort of day that I wrote C code that used the Rmath-julia library and then called that from Python via Cython. \n", "Don't ask

— Joshua Adelman (@synapticarbors) \n", "July 11, 2015\n", "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "display(HTML('''

Today's \n", "the sort of day that I wrote C code that used the Rmath-julia library and then called that from Python via Cython. \n", "Don't ask

— Joshua Adelman (@synapticarbors) \n", "July 11, 2015\n", "
'''))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "As I began to introduce more and more Numba into various code bases at work, I recently decided to revisit this particular bit and see if I could re-implement the whole thing using Numba + CFFI + Rmath. This would cut out the C code that I wrote, the Cython wrapper that involved a bunch of boilerplate strewn across multiple .pyx and .pxd files, and hopefully would make the code easier to extend in the future by people who didn't know C or Cython, but could write some Python and apply the appropriate Numba jit decorator.\n", "\n", "So to begin with, I vendorized the whole Rmath-julia library into our project under `externals/Rmath-julia`. I'll do the same here in this example. Now the fun begins..." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Building the Rmath library using CFFI\n", "\n", "Since we are going to use what cffi calls the \"API-level, out-of-line\", we need to define a build script (`build_rmath.py`) that we will use to compile the Rmath source and produce an importable extension module. The notebook \"cell magic\", `%%file` will write the contents of the below cell to a file." ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting build_rmath.py\n" ] } ], "source": [ "%%file build_rmath.py\n", "\n", "import glob\n", "import os\n", "import platform\n", "\n", "from cffi import FFI\n", "\n", "\n", "include_dirs = [os.path.join('externals', 'Rmath-julia', 'src'),\n", " os.path.join('externals', 'Rmath-julia', 'include')]\n", "\n", "rmath_src = glob.glob(os.path.join('externals', 'Rmath-julia', 'src', '*.c'))\n", "\n", "# Take out dSFMT dependant files; Just use the basic rng\n", "rmath_src = [f for f in rmath_src if ('librandom.c' not in f) and ('randmtzig.c' not in f)]\n", "\n", "extra_compile_args = ['-DMATHLIB_STANDALONE']\n", "if platform.system() == 'Windows':\n", " extra_compile_args.append('-std=c99')\n", "\n", "ffi = FFI()\n", "ffi.set_source('_rmath_ffi', '#include ',\n", " include_dirs=include_dirs,\n", " sources=rmath_src,\n", " libraries=[],\n", " extra_compile_args=extra_compile_args)\n", "\n", "# This is an incomplete list of the available functions in Rmath\n", "# but these are sufficient for our example purposes and gives a sense of\n", "# the types of functions we can get\n", "ffi.cdef('''\\\n", "// Normal Distribution\n", "double dnorm(double, double, double, int);\n", "double pnorm(double, double, double, int, int);\n", "\n", "// Uniform Distribution\n", "double dunif(double, double, double, int);\n", "double punif(double, double, double, int, int);\n", "\n", "// Gamma Distribution\n", "double dgamma(double, double, double, int);\n", "double pgamma(double, double, double, int, int);\n", "''')\n", "\n", "if __name__ == '__main__':\n", " # Normally set verbose to `True`, but silence output\n", " # for reduced notebook noise\n", " ffi.compile(verbose=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Then we simply run the script as below assuming we have a properly configured C compiler on our system. For larger projects [integration with setuptools is supported](http://cffi.readthedocs.org/en/release-1.5/cdef.html#preparing-and-distributing-modules). The exclamation point tells the notebook to execute the following command in a system shell." ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "collapsed": false }, "outputs": [], "source": [ "!python build_rmath.py" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We should now have an extension module named `_rmath_ffi` that gives us access to the functions whose prototypes we enumerated in the `ffi.cdef(...)`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## An example of replicating scipy.stats with Numba\n", "\n", "Now that we have built our module wrapping Rmath using cffi, we can write Numba jit-able versions of scipy stats functions that we can call without additional overhead" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "import numba as nb\n", "\n", "import scipy.stats\n", "\n", "# Import our Rmath module\n", "import _rmath_ffi" ] }, { "cell_type": "markdown", "metadata": { "collapsed": true }, "source": [ "Now we can define a number of shorter aliases to the Rmath functions for use in the python namespace" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "collapsed": false }, "outputs": [], "source": [ "dnorm = _rmath_ffi.lib.dnorm\n", "pnorm = _rmath_ffi.lib.pnorm\n", "\n", "dunif = _rmath_ffi.lib.dunif\n", "punif = _rmath_ffi.lib.punif\n", "\n", "dgamma = _rmath_ffi.lib.dgamma\n", "pgamma = _rmath_ffi.lib.pgamma" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In order for us to use these methods from within a function that we've jit-ed with Numba, we need to import `cffi_support` and register the module: " ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "collapsed": true }, "outputs": [], "source": [ "from numba import cffi_support\n", "\n", "cffi_support.register_module(_rmath_ffi)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "I'll start off by writing a function that is equivalent to `scipy.stats.norm.cdf` using two different styles. In the first, `pnorm_nb`, I'll make the assumption that I'm going to be working on a 1d array, and in the second, I'll use `numba.vectorize` to lift that constraint and create a universal function that can operate on arbitrary dimensional numpy arrays:" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "collapsed": true }, "outputs": [], "source": [ "@nb.jit(nopython=True)\n", "def pnorm_nb(x):\n", " y = np.empty_like(x)\n", " for k in xrange(x.shape[0]):\n", " y[k] = pnorm(x[k], 0.0, 1.0, 1, 0)\n", " \n", " return y\n", "\n", "@nb.vectorize(nopython=True)\n", "def pnorm_nb_vec(x):\n", " return pnorm(x, 0.0, 1.0, 1, 0)" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "True\n", "True\n" ] } ], "source": [ "x = np.random.normal(size=(100,))\n", "\n", "y1 = scipy.stats.norm.cdf(x)\n", "y2 = pnorm_nb(x)\n", "y3 = pnorm_nb_vec(x)\n", "\n", "# Check that they all give the same results\n", "print np.allclose(y1, y2)\n", "print np.allclose(y1, y3)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "And now let's do the same calculation for 2D data, demonstrating that the vectorized form of the Numba function automatically create the appropriate universal function for the given dimensionality of the inputs:" ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "True\n" ] } ], "source": [ "x = np.random.normal(size=(100,100))\n", "\n", "y1 = scipy.stats.norm.cdf(x)\n", "y2 = pnorm_nb_vec(x)\n", "\n", "# Check that they all give the same results\n", "print np.allclose(y1, y2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Timing the scipy and numba versions:" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1000 loops, best of 3: 618 µs per loop\n", "1000 loops, best of 3: 336 µs per loop\n" ] } ], "source": [ "%timeit scipy.stats.norm.cdf(x)\n", "%timeit pnorm_nb_vec(x)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can see that our Numba version is almost 2x faster than the scipy version, with the added bonus that it can be called from within other Numba-ized methods without going through the python object layer, which can be quite slow.\n", "\n", "Just for kicks, lets also try to take advantage of multiple cores using the `target` argument:" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "collapsed": false }, "outputs": [], "source": [ "@nb.vectorize([nb.float64(nb.float64),], nopython=True, target='parallel')\n", "def pnorm_nb_vec_parallel(x):\n", " return pnorm(x, 0.0, 1.0, 1, 0)" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "True\n", "\n", "The slowest run took 17.27 times longer than the fastest. This could mean that an intermediate result is being cached \n", "1000 loops, best of 3: 131 µs per loop\n" ] } ], "source": [ "y3 = pnorm_nb_vec_parallel(x)\n", "print np.allclose(y1, y3)\n", "print\n", "\n", "%timeit pnorm_nb_vec_parallel(x)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "So on my laptop with 4 physical cores, we get a nice additional speed-up over the serial numba version and the `scipy.stats` function.\n", "\n", "Finally, I'm going to programmatically wrap the all of the Rmath functions I exposed and compare them to the equivalent scipy functions." ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "norm_pdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 545 µs per loop\n", "numba timing:\n", "1000 loops, best of 3: 212 µs per loop\n", "\n", "norm_cdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 634 µs per loop\n", "numba timing:\n", "1000 loops, best of 3: 328 µs per loop\n", "\n", "unif_pdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 436 µs per loop\n", "numba timing:\n", "10000 loops, best of 3: 68.4 µs per loop\n", "\n", "unif_cdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 495 µs per loop\n", "numba timing:\n", "10000 loops, best of 3: 128 µs per loop\n", "\n", "gamma_pdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 616 µs per loop\n", "numba timing:\n", "1000 loops, best of 3: 277 µs per loop\n", "\n", "gamma_cdf\n", "allclose: True\n", "scipy timing:\n", "1000 loops, best of 3: 1.13 ms per loop\n", "numba timing:\n", "1000 loops, best of 3: 1.11 ms per loop\n", "\n" ] } ], "source": [ "from collections import OrderedDict\n", "\n", "func_map = OrderedDict([\n", " ('norm_pdf', (scipy.stats.norm.pdf, dnorm)),\n", " ('norm_cdf', (scipy.stats.norm.cdf, pnorm)),\n", " ('unif_pdf', (scipy.stats.uniform.pdf, dunif)),\n", " ('unif_cdf', (scipy.stats.uniform.cdf, punif)),\n", " ('gamma_pdf', (scipy.stats.gamma.pdf, dgamma)),\n", " ('gamma_cdf', (scipy.stats.gamma.cdf, pgamma)),\n", " ])\n", "\n", "def generate_function(name, rmath_func):\n", " if 'norm' in name or 'unif' in name:\n", " def cdf_func(x):\n", " return rmath_func(x, 0.0, 1.0, 1, 0)\n", "\n", " def pdf_func(x):\n", " return rmath_func(x, 0.0, 1.0, 0)\n", " \n", " elif 'gamma' in name:\n", " def cdf_func(x, shape):\n", " return rmath_func(x, shape, 1.0, 1, 0)\n", "\n", " def pdf_func(x, shape):\n", " return rmath_func(x, shape, 1.0, 0)\n", " \n", " \n", " if 'cdf' in name:\n", " return cdf_func\n", " elif 'pdf' in name:\n", " return pdf_func\n", "\n", "x = np.random.normal(size=(100,100))\n", " \n", "for name, (scipy_func, rmath_func) in func_map.iteritems():\n", " nb_func = nb.vectorize(nopython=True)(generate_function(name, rmath_func))\n", " \n", " \n", " print name\n", " if 'norm' in name or 'unif' in name:\n", " y1 = scipy_func(x)\n", " y2 = nb_func(x)\n", " print 'allclose: ', np.allclose(y1, y2)\n", " \n", " print 'scipy timing:'\n", " %timeit scipy_func(x)\n", " print 'numba timing:'\n", " %timeit nb_func(x)\n", " elif 'gamma' in name:\n", " shape = 1.0\n", " y1 = scipy_func(x, shape)\n", " y2 = nb_func(x, shape)\n", " print 'allclose: ', np.allclose(y1, y2)\n", " \n", " print 'scipy timing:'\n", " %timeit scipy_func(x, shape)\n", " print 'numba timing:'\n", " %timeit nb_func(x, shape)\n", " print" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion\n", "\n", "To wrap up, CFFI + Numba provides a powerful and surprisingly simple way to generate fast python code and extend the currently limited repertoire of functionality that is baked into Numba. Pairing this approach with Rmath specifically has been particularly useful in my own work." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Appendix\n", "\n", "For completeness, I'll use Sebastian Raschka's watermark package to specify the libraries and hardware used to run these examples:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%install_ext https://raw.githubusercontent.com/rasbt/watermark/master/watermark.py" ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPython 2.7.11\n", "IPython 4.0.3\n", "\n", "numpy 1.10.2\n", "scipy 0.16.1\n", "numba 0.23.1\n", "cffi 1.5.0\n", "\n", "compiler : GCC 4.2.1 (Apple Inc. build 5577)\n", "system : Darwin\n", "release : 13.4.0\n", "machine : x86_64\n", "processor : i386\n", "CPU cores : 8\n", "interpreter: 64bit\n" ] } ], "source": [ "%load_ext watermark\n", "%watermark -vm -p numpy,scipy,numba,cffi" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 2", "language": "python", "name": "python2" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython2", "version": "2.7.11" } }, "nbformat": 4, "nbformat_minor": 0 }