← retour aux snippets

Numba parallèle: prange et nopython

activer le parallélisme et le mode nopython pour vectoriser des boucles

Numba parallèle: prange et nopython

objectif

Expliquer et montrer comment activer le parallélisme et le mode nopython pour vectoriser des boucles.

code minimal

import numpy as np
from numba import njit, prange

@njit(parallel=True, fastmath=True)
def matvec(A, x):
    m, n = A.shape
    y = np.empty(m)
    for i in prange(m):
        s = 0.0
        for j in range(n):
            s += A[i,j] * x[j]
        y[i] = s
    return y

A = np.random.rand(2000, 512)
x = np.random.rand(512)
y = matvec(A, x)
y.shape

utilisation

# vérifier la somme pour sanity check
float(y.mean())

variante(s) utile(s)

# vectorisation via np.dot reste souvent plus rapide, comparez!

notes

  • fastmath permet des optimisations au prix d’un flou numérique.
  • prange parallélise la boucle externe si possible.