Probando numba: compilador para Python basado en LLVM

Introducción

Hace unos días Travis E. Oliphant, creador de NumPy e importante contribuidor de SciPy entre otras muchas cosas, anunciaba en su Twitter y en su blog la liberación de numba 0.1, un proyecto que pretende ser el mejor compilador orientado a arrays del mundo, como puedes ver en la presentación que dio en la conferencia SciPy 2012 celebrada en Austin, Texas (vídeo y diapositivas).

Aunque el proyecto está en una fase bastante precaria todavía y hay unos cuantos fallos pendientes de solucionar todavía, hemos hecho algunas pruebas y los resultados son impresionantes. Vamos a hablar un poco de numba y a explicar cómo puedes probarlo tú mismo.

Motivación

Ya hemos hablado en dos ocasiones sobre la lentitud de Python y hemos dado algunas soluciones para evitarla, bien utilizando bibliotecas que aceleren el código escrito en Python o teniendo en cuenta pequeños trucos a la hora de escribir nuestros programas. En la charla que hemos mencionado antes, Oliphant menciona algunas de ellas, concretamente PyPy y Cython (David ya habló de las dos en su artículo). Estas dos soluciones tienen sus inconvenientes:

  • PyPy no funciona con CPython, la implementación principal de Python. Esto significa que las extensiones escritas para Python no se pueden utilizar fácilmente en PyPy, y hay una que es fundamental para nosotros: NumPy.
  • Cython supone otra sintaxis que hay que aprender, y si se quiere optimizar mucho se va perdiendo la legibilidad de Python poco a poco.

Con numba, tenemos una forma de acelerar código escrito en Python que:

  • Está especialmente pensada y optimizada para códigos numéricos,
  • Entiende los tipos de NumPy, su API y sus estructuras de datos,
  • No requiere aprender una nueva sintaxis, conservando así la belleza y legibilidad de Python, y
  • Produce programas con velocidades cercanas a las de C.

¿Y cómo se consigue esto? numba genera código LLVM y crea un wrapper para llamarlo desde Python. De esta forma conseguimos un compilador JIT para Python, en la línea de otros proyectos como el abandonado Unladen Swallow.

Interesante, ¿no? Si quieres probar numba en tu propio ordenador, sigue leyendo.

Instalación

Para no complicarnos la vida, vamos a dar las instrucciones de instalación para Linux. Parece ser que Windows no está soportado en Windows hay un par de problemas (#27, #28), y en cuanto a Mac OS X no tengo ni idea.

En la web de numba tienes los pasos que hay que seguir para instalarlo. Los requisitos son:

  1. llvm 3.1 compilado con las opciones --enable-pic--disable-libffi. La segunda opción soluciona un bug que afecta a algunos, incluido yo.
  2. llvm-py. Con la versión del repositorio ha funcionado.

En primer lugar, descarga y extrae las fuentes de llvm 3.1 y ejecuta los siguientes comandos para compilar la biblioteca:

$ ./configure --enable-pic --disable-libffi
$ make
$ make install

A continuación, clona el repositorio de llvm-py e instala la biblioteca:

$ git clone https://github.com/llvmpy/llvmpy.git
$ cd llvmpy
$ python setup.py install

Por último, ya puedes instalar numba directamente del repositorio también:

$ git clone https://github.com/numba/numba.git
$ cd numba
$ git submodule init
$ git submodule update
$ python setup.py install

En principio estas instrucciones funcionan en el caso más general, pero si tienes algún problema concreto durante el proceso o más adelante probando los ejemplos, dínoslo en los comentarios.

Ejemplos

Las fuentes de numba contienen algunos ejemplos y benchmarks que puedes probar para evaluar las capacidades de la biblioteca. Además, Alex Wiltschko ha estado haciendo pruebas e informando de algunos fallos y nos ha dejado que hagamos uso de los programas que ha escrito (thanks Alex!).

Las capacidades de numba se aprecian mejor si se aplican con programas con muchos bucles anidados, como algoritmos de procesamiento de imágenes. De todos modos, para empezar vamos a probar un ejemplo sencillo, que puedes encontrar en https://github.com/numba/numba/blob/master/examples/sum.py.

Nota: En el momento de escribir el artículo se imprime todo el código LLVM creado por numba cuando se aplica a una función. Para evitar este comportamiento, puedes ejecutar los ejemplos añadiendo la opción -O, que activa optimizaciones básicas.

El código del ejemplo es el siguiente:

from numba import d
from numba.decorators import jit as jit
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i,j]
    return result
csum2d = jit(ret_type=d, arg_types=[d[:,:]])(sum2d)
from numpy import random
arr = random.randn(100,100)
import time
start = time.time()
res = sum2d(arr)
duration = time.time() - start
print "Result from python is %s in %s (msec)" % (res, duration*1000)
start = time.time()
res = csum2d(arr)
duration2 = time.time() - start
print "Result from compiled is %s in %s (msec)" % (res, duration2*1000)
print "Speed up is %s" % (duration / duration2)

Como se puede ver, en la línea 12 se aplica la función jit (que también funciona como decorador, como ya veremos) al pequeño programa que hemos escrito para sumar los elementos de una matriz. El argumento arg_types indica que a la función sum2d se le va a pasar un array de NumPy de dos dimensiones. Esto viene a ser como una declaración de tipos.

Comprobemos los resultados:

$ python -O sum.py
Result from python is -72.8742506632 in 4.33301925659 (msec)
Result from compiled is -72.8742506632 in 0.028133392334 (msec)
Speed up is 154.016949153

Conseguimos un aumento de 150x, solamente añadiendo una línea.

Probemos otro de los ejemplos, https://github.com/numba/numba/blob/master/examples/example.py. En este caso se trata de comparar numba con una función de SciPy, y podemos ver cómo se utiliza el decorador:

@jit(arg_types=[int32[:,:], int32[:,:]], ret_type=int32[:,:])
def filter2d(image, filt):
    M, N = image.shape
    Mf, Nf = filt.shape
    Mf2 = Mf // 2
    Nf2 = Nf // 2
    result = numpy.zeros_like(image)
    for i in range(Mf2, M - Mf2):
        for j in range(Nf2, N - Nf2):
            num = 0.0
            for ii in range(Mf):
                for jj in range(Nf):
                    num += (filt[Mf-1-ii, Nf-1-jj] * image[i-Mf2+ii, j-Nf2+jj])
            result[i, j] = num
    return result

Ahora estamos indicando que la función requiere dos arrays de enteros. Nótese que tenemos cuatro (4) bucles anidados. Si ahora probamos el ejemplo:

$ python -O example.py
Time for LLVM code = 0.028516
Time for convolve = 0.041848

Aquí hemos conseguido una ganancia de 1.47x. Nótese que la versión de SciPy está escrita en C.

Por último, vamos a ver uno de los experimentos llevados a cabo por Alex Wiltschko creado a raíz de un artículo sobre Cython. Podemos ver los resultados del experimento en el IPython Notebook Viewer, una web que utiliza la nueva interfaz web de IPython en la que podemos compartir notebooks.

El notebook del experimento está en http://nbviewer.ipython.org/3362653, y aunque según los comentarios hay resultados dispares, en el ordenador de Wiltschko la versión con numba fue 2 veces más rápida que la mejr versión de Cython. Si ahora contamos el número de líneas:

  • Python: 17 líneas.
  • Cython + memviews (sin slicing): 39 líneas.
  • numba: 13 líneas.

Aun considerando que las velocidades sean del mismo orden, es una mejora espectacular.

Comentarios finales

Estos resultados son bastante impresionantes, pero como hemos dicho más arriba la biblioteca está todavía en pruebas y aún no tiene documentación. Por ejemplo, dentro de las funciones que compilemos con jit no funcionan la mayoría de las funciones de NumPy, como array.zeros.

Sin embargo, numba es un notable paso hacia adelante y, si no se queda estancada, puede que abra un nuevo camino lleno de posibilidades en la popularización de Python en entornos científicos.

Nosotros estamos muy entusiasmados con numba, ¿y tú? 🙂

Juan Luis Cano

Estudiante de ingeniería aeronáutica y con pasión por la programación y el software libre. Obsesionado con mejorar los pequeños detalles y con ganas de cambiar el mundo. Divulgando Python en español a través de Pybonacci y la asociación Python España.

More Posts - Website

Follow Me:
TwitterLinkedIn

12 thoughts on “Probando numba: compilador para Python basado en LLVM

  1. Pingback: Bitacoras.com
  2. He probado el primer ejemplo en un MacBook con Mac Os X 10.6 con llvm instalado con MacPorts, y aparte de un pequeño cambio en el script de llvmpy para que me lo pillara, ha funcionado a la primera. Eso si, cada vez que lo uso, obtengo un monton de texto por pantalla sobre bloques y llvm al importar.

    Para el primer ejemplo:

    Result from python is -13.5089689484 in -14.6307945251 (msec)
    Result from compiled is -13.5089689484 in -0.068187713623 (msec)
    Speed up is 214.566433566

    En realidad, eso es poco fiable porque es muy corto. Generando una matriz de 10 000 x 10 000, el resultado es:

    Result from python is 12008.571806 in -110127.022028 (msec)
    Result from compiled is 12008.571806 in -1454.92005348 (msec)
    Speed up is 75.6928339563

    Menos, pero igualmente considerable. Pero como siempre es bueno correr las cosas varias veces, al repetir:

    Result from python is -4215.08982708 in -92592.3008919 (msec)
    Result from compiled is -4215.08982708 in -313.112974167 (msec)
    Speed up is 295.715312143

    Comparando con el metodo .sum de NumPy la cosa no es tan espectacular:

    Result from python is -6189.47858681 in -94576.695919 (msec)
    Result from compiled is -6189.47858681 in -278.85890007 (msec)
    Result from NumPy is -6189.47858681 in -297.269105911 (msec)
    Speed up is 339.15609613
    Speed up is 1.06601978935 with NumPy

    Esencialmente, el mismo tiempo.

    1. Gracias por repetir las pruebas, parece que se confirman los resultados. Para desactivar la salida de código LLVM, a mí me ha funcionado invocando python -O.

      El hecho de que sea esencialmente el mismo tiempo que el sum de NumPy tiene mucho sentido, pero la gracia está en que uno está en C y otro en Python puro. Le veo mucho futuro a esto, además es muy pythonico.

      1. Sí, yo también le veo futuro (aunque sea el del comentario dispar ;)).

        Falta por ver cómo funcionará implementando algoritmos más complejos (ahí me da la sensación de que Cython seguirá siendo superior), pero para código puramente numérico el resultado es mucho más efectivo a todos los niveles: cantidad de código, coneptos/sintáxis…

        Por cierto, ¿alguien que sepa Fortran tiene alguna opinión sobre este comentario en el artículo sobre Cython? “””My own conclusion: I just implement the math directly in Fortran, in the most straightforward way and I usually get faster version than *I* can get after spending hours in Cython. If it is not fast enough, I can try to rewrite the loops in a few ways to help Fortran optimize it better.”””

      2. Ondřej Čertík se pasó de C++ a Fortran hace algo más de un año y ha experimentado bastante con él. De hecho creó http://fortran90.org/, donde ha ido recopilando información valiosísima sobre cómo escribir Fortran moderno y legible.

        Yo mismo he usado Fortran también, pero no a este nivel ni mucho menos, ni he hecho comparaciones con otros lenguajes. Pero si Čertík dice eso, me fío completamente 🙂

        Fortran está diseñado para que escribir matemáticas con él sea sencillo (FORmula TRANslation), y los compiladores están optimizados para manejo de arrays. Así que seguro que cualquier código científico que se base en matrices correrá más rápido en Fortran sin mucho esfuerzo.

  3. uysss. Me parece a mi que ya tengo tarea para el finde (mi señora me mata).
    Normalmente soy ateo, pero para el caso de llvm me considero agnóstico!!! La gente me maravilla.

  4. Ayer estuve probándolo en windows y se ‘rompía’ todo el rato. No conseguía pasar arrays, solo constantes, con errores de ‘translate’ todo el rato… Lo probaré en casa con Linux a ver si consigo algo diferente 🙁

      1. Prefiero no comentar nada en el issue tracker o en la lista de usuarios de Numba de momento puesto que usé el binario para windows descargado de http://www.lfd.uci.edu/~gohlke/pythonlibs/ y no es algo oficial. Probaré en linux en casa cuando tenga un rato con el código ‘oficial’.

        Por cierto, que gran recurso la página de Christoph Gohlke si eres usuario de windows.

Leave a Reply

Your email address will not be published. Required fields are marked *