<jupyter_start><jupyter_text>CPU JIT Numba (http://numba.pydata.org) is a package created by Anaconda, Inc (http://www.anaconda.com). Numba takes pure Python code and translates it automatically (JIT) into optimized machine code. In practice, this means that we can write a non-vectorized function in pure Python, using for loops, and have this function vectorized automatically by using a single decorator. Performance speedups when compared to pure Python code can reach several orders of magnitude and may even outmatch manually-vectorized NumPy code.<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
size = 400
iterations = 100
def mandelbrot_python(size, iterations):
    m = np.zeros((size, size))
    for i in range(size):
        for j in range(size):
            c = (-2 + 3. / size * j + 1j * (1.5 - 3. / size * i))
            z = 0
            for n in range(iterations):
                if np.abs(z) <= 10:
                    z = z * z + c
                    m[i, j] = n
                else:
                    break
    return m
m = mandelbrot_python(size, iterations)
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
ax.imshow(np.log(m), cmap=plt.cm.hot)
ax.set_axis_off()
%timeit mandelbrot_python(size, iterations)<jupyter_output>4.75 s ± 34.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Let's try to accelerate this function using Numba. We add the @jit decorator right above the function definition, without changing a single line of code in the body of the function:<jupyter_code>from numba import jit
@jit
def mandelbrot_numba(size, iterations):
    m = np.zeros((size, size))
    for i in range(size):
        for j in range(size):
            c = (-2 + 3. / size * j + 1j * (1.5 - 3. / size * i))
            z = 0
            for n in range(iterations):
                if np.abs(z) <= 10:
                    z = z * z + c
                    m[i, j] = n
                else:
                    break
    return m
mandelbrot_numba(size, iterations)
%timeit mandelbrot_numba(size, iterations)<jupyter_output>26 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)<jupyter_text>GPU JIT<jupyter_code>#!conda install cudatoolkit
import math
import numpy as np
from numba import cuda
import matplotlib.pyplot as plt
%matplotlib inline
len(cuda.gpus)
print(list(gpu.name for gpu in cuda.gpus))
@cuda.jit
def mandelbrot_numba(m, iterations):
    # Matrix index.
    i, j = cuda.grid(2)
    size = m.shape[0]
    # Skip threads outside the matrix.
    if i >= size or j >= size:
        return
    # Run the simulation.
    c = (-2 + 3. / size * j + 1j * (1.5 - 3. / size * i))
    z = 0
    for n in range(iterations):
        if abs(z) <= 10:
            z = z * z + c
            m[i, j] = n
        else:
            break
size = 400
iterations = 100
m = np.zeros((size, size))
# 16x16 threads per block.
bs = 16
# Number of blocks in the grid.
bpg = math.ceil(size / bs)
# We prepare the GPU function.
f = mandelbrot_numba[(bpg, bpg), (bs, bs)]
f(m, iterations)
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
ax.imshow(np.log(m), cmap=plt.cm.hot)
ax.set_axis_off()
%timeit -n10 -r100 f(m, iterations)<jupyter_output><empty_output><jupyter_text>That's about 10 times faster than the CPU version obtained with Numba in the first recipe in this chapter, and 1,800 times faster than the pure Python version! But we can do even better. Numba takes care of transferring arrays automatically between the host machine (CPU) and the device (GPU). These data transfers are slow, sometimes even more than the actual on-device computation. Numba provides facilities to deal with these transfers manually, which can be interesting in some use cases. Let's estimate the time of the data transfers and the computation on the GPU.<jupyter_code># First, we send the NumPy array to the GPU with the cuda.to_device() function
%timeit -n10 -r100 cuda.to_device(m)
# Second, we run the computation on the GPU:
%%timeit -n10 -r100 m_gpu = cuda.to_device(m)
# Third, we copy the modified array from the GPU to the CPU.
m_gpu = cuda.to_device(m)
%timeit -n10 -r100 m_gpu.copy_to_host()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0989010989010989, 'fraction_numerical': 0.0241280458671763, 'mean_word_length': 2.8412844036697247, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '237613', 'license': 'MIT', 'n_tokens_mistral': 1472, 'n_tokens_neox': 1368, 'n_words': 530, 'path': 'Tutorial/JIT Compilation.ipynb', 'repo_name': 'FallenGameR/LearningPython', 'size': '4186'}