Comment faire en sorte que numba @jit utilise tous les cœurs de processeur (paralléliser numba @jit)

Question

J'utilise numbas @jit pour ajouter deux tableaux numpy en python. Les performances sont très élevées si j'utilise @jit par rapport à python .

Toutefois, il est ne pas utiliser tous les cœurs de l'unité centrale même si je passe en @numba.jit(nopython = True, parallel = True, nogil = True) .

Existe-t-il un moyen d'utiliser tous les cœurs de l'unité centrale avec le nombre de cœurs ? @jit .

Voici mon code :

import time                                                
import numpy as np                                         
import numba                                               

SIZE = 2147483648 * 6                                      

a = np.full(SIZE, 1, dtype = np.int32)                     

b = np.full(SIZE, 1, dtype = np.int32)                     

c = np.ndarray(SIZE, dtype = np.int32)                     

@numba.jit(nopython = True, parallel = True, nogil = True) 
def add(a, b, c):                                          
    for i in range(SIZE):                                  
        c[i] = a[i] + b[i]                                 

start = time.time()                                        
add(a, b, c)                                               
end = time.time()                                          

print(end - start)

Answer 1

Pour être complet, en 2018 (numba v 0.39), vous pouvez simplement faire

from numba import prange

et remplacer range con prange dans la définition originale de la fonction, c'est tout.

L'utilisation de l'unité centrale atteint immédiatement 100 % et, dans mon cas, le temps d'exécution passe de 2,9 à 1,7 secondes (pour SIZE = 2147483648 * 1, sur une machine dotée de 16 cœurs et de 32 threads).

Pour les noyaux plus complexes, il est souvent possible d'accélérer encore plus la vitesse en passant par le paramètre fastmath=True .