IPython + jit decorator

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

[1]:
import numpy as np
[2]:
image = np.random.rand(1024, 1024)
[3]:
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!

[4]:
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”):

[5]:
result = laplace_pythran(image)

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

[6]:
from transonic import wait_for_all_extensions
wait_for_all_extensions()

Now, let’s benchmark the 2 functions:

[7]:
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.000 * norm
norm = 1.37e-02 s
laplace_pythran                  : 0.119 * norm

It corresponds to a speedup of ~ 8.4!