Lunch Time Python

Lunch 6: numba

numba is a just-in-time (JIT) compiler for Python. With a few simple annotations, array-oriented and math-heavy Python code can be just-in-time optimized to performance similar as C, C++ and Fortran, without having to switch languages or Python interpreters.

Press Spacebar to go to the next slide (or ? to see all navigation shortcuts)

Lunch Time Python, Scientific Software Center, Heidelberg University

Motivation

• Many reasons to use Python, but performance not one of them
• What to do when a Python function is too slow?
• Ideally, find a library (e.g. numpy) with an equivalent function
• Otherwise:
  • use PyPy instead of CPython (if all your libraries are available)
  • write a fortan function and compile with f2py or fortranmagic
  • write a C function and compile with Cython
  • write a C++ function and compile using pybind11 or ipybind
  • magically make your slow Python function faster (numba)

numba installation

• Conda: conda install numba
• Pip: python -m pip install numba

Vector reduction example

Toy example: implement a vector reduction operation:

r(x,y) = $ \sum_i \cos(x_i) \sin(y_i) $

Some random vectors to benchmark our functions:

import numpy as np

x = np.random.uniform(low=-1, high=1, size=5000000)
y = np.random.uniform(low=-1, high=1, size=5000000)

Python

import math


def r_python(x_vec, y_vec):
    s = 0
    for x, y in zip(x_vec, y_vec):
        s += math.cos(x) * math.sin(y)
    return s

r_python(x, y)

-31.753388977618677

%timeit r_python(x,y)

1.65 s ± 20.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

numpy

def r_numpy(x_vec, y_vec):
    return np.dot(np.cos(x_vec), np.sin(y_vec))

r_numpy(x, y)

-31.753388977676764

%timeit r_numpy(x,y)

37.8 ms ± 225 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Cython

# pip install cython
%load_ext cython

%%cython

import math

def r_cython(x_vec, y_vec):
    s = 0
    for x,y in zip(x_vec, y_vec):
        s += math.cos(x) * math.sin(y)
    return s

r_cython(x, y)

-31.753388977618677

%timeit r_cython(x,y)

1.3 s ± 3.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%cython

import math
# use C math functions
from libc.math cimport sin, cos

# use C types instead of Python types
def r_cython(double[:] x_vec, double[:] y_vec):
    cdef double s = 0
    cdef int i
    for i in range(len(x_vec)):
        s += cos(x_vec[i])*sin(y_vec[i])
    return s

r_cython(x, y)

-31.753388977618677

%timeit r_cython(x,y)

114 ms ± 81.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Fortran

if "google.colab" in str(get_ipython()):
    !pip install fortran-magic -qqq
%load_ext fortranmagic

%%fortran

subroutine r_fortran(x_vec, y_vec, res)
    real, intent(in) :: x_vec(:), y_vec(:)
    real, intent(out) :: res
    integer :: i, n
    n = size(x_vec)
    res = 0
    do i=1,n
        res = res + cos(x_vec(i))*sin(y_vec(i))
    enddo
endsubroutine r_fortran

r_fortran(x, y)

-31.762666702270508

%timeit r_fortran(x,y)

90.2 ms ± 121 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

C++ / pybind11

if "google.colab" in str(get_ipython()):
    !pip install git+https://github.com/aldanor/ipybind.git -qqq
%load_ext ipybind

%%pybind11

#include <pybind11/numpy.h>
#include <math.h>
PYBIND11_PLUGIN(example) {
    py::module m("example");
    m.def("r_pybind", [](const py::array_t<double>& x, const py::array_t<double>& y) {
        double sum{0};
        auto rx{x.unchecked<1>()};
        auto ry{y.unchecked<1>()};
        for (py::ssize_t i = 0; i < rx.shape(0); i++){
            sum += std::cos(rx[i])*std::sin(ry[i]);
        }
        return sum;
    });
    return m.ptr();
}

r_pybind(x, y)

-31.753388977618677

%timeit r_pybind(x, y)

108 ms ± 47.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

numba

from numba import jit


@jit
def r_numba(x_vec, y_vec):
    s = 0
    for x, y in zip(x_vec, y_vec):
        s += math.cos(x) * math.sin(y)
    return s

r_numba(x, y)

-31.753388977618677

# pure python with numba JIT
%timeit r_numba(x,y)

110 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numba compilation

Two compilation modes

• nopython mode (default)
  • Fast because it doesn't access the Python C API
  • Needs to be able to infer the native (C) types of all values
• object mode (fallback)
  • Slow because it uses Python objects and the Python C API
  • Only used if nopython mode is not possible
  • To raise an error instead of falling back to this, set nopython=True or use @njit

Numba function signatures

You can optionally explicitly specify the function signature. Use cases:

• you want the function to be compiled when it is defined rather than when it is first called
• you need fine-grained control over types (e.g. if you want 32-bit floats)

from numba import float32


@jit(float32(float32, float32))
def sum(a, b):
    return a + b

sum(1, 0.99999999)

2.0

Numba options

• nopython=True disable Object mode fallback
• nogil=True release the Python Global Interpreter Lock (GIL)
• cache=True cache the compiled funtions on disk
• parallel=True enable automatic parallelization

Parallelization

• set parallel=True option to enable
• use prange to explicitly parallelize a loop over a range

from numba import jit, prange


@jit(parallel=True)
def r_numba(x_vec, y_vec):
    s = 0
    for i in prange(len(x_vec)):
        s += math.cos(x[i]) * math.sin(y[i])
    return s

r_numba(x, y)

-31.753388977717123

%timeit r_numba(x,y)

56.2 ms ± 784 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

NumPy universal functions

• a numpy ufunc is a function that operates on scalars
• can create one using @numba.vectorize and use it like built-in numpy ufuncs

from numba import vectorize, float64


@vectorize([float64(float64, float64)], target="parallel")
def r(x, y):
    return np.cos(x) * np.sin(y)

r(2, 3)

-0.05872664492762098

r(x, y)

array([-0.63636374,  0.00639962,  0.00882265, ..., -0.77644569,
        0.30776024,  0.52483202])

np.sum(r(x, y))

-31.753388977680412

%timeit np.sum(r(x,y))

87 ms ± 506 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Advanced features

• Ahead of Time (AoT) compilation
  • the compiled module only depends on NumPy
• Flexible specializations
  • @generated_jit decorator for compile-time logic, e.g. type specializations
• Stencil
  • @stencil decorator for creating a stencil to apply to an array
• C callbacks
  • @cfunc decorator to generate a C-callback (e.g. to pass to scipy.integrate)
• CUDA support
  • compile CUDA kernels to run on a GPU
• see numba.readthedocs.io for more