IPython + jit decorator

Let’s say that we have this kind of code:

import numpy as np

image = np.random.rand(1024, 1024)

def laplace(image):
    """Laplace operator in NumPy for 2D images."""
    laplacian = (
        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]
        - 4*image[1:-1, 1:-1]
    )
    thresh = np.abs(laplacian) > 0.05
    return thresh

We can define a jitted function in IPython (and Jupyter, which uses IPython). The only limitation is that everything needed for the function (even np) has to be (re)defined in the same cell!

import numpy as np
from transonic import jit

@jit
def laplace_pythran(image):
    """Laplace operator in NumPy for 2D images."""
    laplacian = (
        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]
        - 4*image[1:-1, 1:-1]
    )
    thresh = np.abs(laplacian) > 0.05
    return thresh

We first call the function to launch the compilation (“warmup the jit”):

result = laplace_pythran(image)

Since the compilation is quit long, we explicitelly wait for the extension to benchmark the compiled version.

from transonic import wait_for_all_extensions
wait_for_all_extensions()

Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__25970b035d85f309561968e246899d70/laplace_pythran.py
File /home/pierre/.transonic/pythran/__jit__/__ipython__25970b035d85f309561968e246899d70/laplace_pythran_1638c8c99ead3deefe4acf18f6c03934_8b9e99482f5d1fa6279790dd20e989ce.cpython-39-x86_64-linux-gnu.so created by pythran

Now, let’s benchmark the 2 functions:

from transonic.util import timeit_verbose

namespace = {"laplace": laplace, "laplace_pythran": laplace_pythran, "image": image}

norm = timeit_verbose("laplace(image)", globals=namespace)
time_pythran = timeit_verbose("laplace_pythran(image)", globals=namespace, norm=norm)
print(f"\nIt corresponds to a speedup of ~ {norm/time_pythran:.1f}!")

laplace                          :     1 * norm
norm = 0.00831 s
laplace_pythran                  : 0.131 * norm

It corresponds to a speedup of ~ 7.6!