Mon article précédent (qui a été honnêtement créé pour tester le thème de ce site), présentait quelques extraits de code qui calculaient $N$ termes de la somme des inverses des carrés. J’ai écrit le code dans mes 4 langages préférés—Python, C, Rust, et Haskell—mais quand j’ai exécuté le code Python, il était embarrassamment lent. Comparé aux $\approx 950$ ms qu’a pris Rust séquentiel, Python a pris 70 secondes ! Donc, dans cet article, nous allons tenter d’obtenir des chiffres plus raisonnables pour Python.

La solution originale

def basel(N: int) -> float:
    return sum(x**(-2) for x in range(1,N))

C’est la manière la plus Pythonique d’accomplir la tâche. Une ligne de code lisible et facile à comprendre qui utilise des générateurs intégrés. Chronométrons-la, avec $N=10^8$ au lieu de $10^9$ comme dans le dernier article (car je ne veux pas attendre si longtemps) :

import time

# Chronométrer la fonction d'entrée
def time_function(func):
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        print(f"Function '{func.__name__}' executed in {execution_time:.6f} seconds.")
        return result
    return wrapper

>>> f = time_function(basel)
>>> f(100000000) # 10^8
Function 'basel' executed in 6.641589 seconds.
1.644934057834575

Essayons de la réécrire, mais de manière moins Pythonique, en utilisant une boucle for :

def basel_less_pythonic(N: int) -> float:
    s = 0.0
    for x in range(1, N):
        s += x**(-2)
    return s

>>> f(100000000) # 10^8
Function 'basel_less_pythonic' executed in 5.908466 seconds.
1.644934057834575

Intéressant. La manière moins idiomatique de l’écrire a en fait amélioré les performances. Pourquoi ? Investigons en utilisant dis, le module de désassemblage de Python.

Voici la version originale. J’ai ajouté quelques commentaires utiles :

# Charger les variables
00 LOAD_GLOBAL              0 (sum)
02 LOAD_CONST               1 (<code object <genexpr> at 0x104f42e40>)
04 LOAD_CONST               2 ('basel.<locals>.<genexpr>')

# Créer la fonction
06 MAKE_FUNCTION            0

# Créer l'objet `range`
08 LOAD_GLOBAL              1 (range)
10 LOAD_CONST               3 (1)
12 LOAD_FAST                0 (N)
14 CALL_FUNCTION            2

# Convertir en itérateur
16 GET_ITER

# Appeler le générateur
18 CALL_FUNCTION            1

# Appeler la fonction sum
20 CALL_FUNCTION            1

# Retourner
22 RETURN_VALUE

f <code object <genexpr> at 0x104f42e40>:
# Essentiellement une boucle `for` avec yield
00 GEN_START                0
02 LOAD_FAST                0 (.0)
04 FOR_ITER                 7 (to 20)
06 STORE_FAST               1 (x)
08 LOAD_FAST                1 (x)
10 LOAD_CONST               0 (-2)
12 BINARY_POWER
14 YIELD_VALUE
16 POP_TOP
18 JUMP_ABSOLUTE            2 (to 4)
20 LOAD_CONST               1 (None)
22 RETURN_VALUE

Voici la version avec boucle for :

# Charger les variables
00 LOAD_CONST               1 (0.0)
02 STORE_FAST               1 (s)

# Créer l'objet `range`
04 LOAD_GLOBAL              0 (range)
06 LOAD_CONST               2 (1)
08 LOAD_FAST                0 (N)
10 CALL_FUNCTION            2
12 GET_ITER

# Déclarer la boucle `for`
14 FOR_ITER                 8 (to 32)
16 STORE_FAST               2 (x)

# Puissance, addition, stockage
18 LOAD_FAST                1 (s)
20 LOAD_FAST                2 (x)
22 LOAD_CONST               3 (-2)
24 BINARY_POWER
26 INPLACE_ADD
28 STORE_FAST               1 (s)

# Retourner au début de la boucle `for`
30 JUMP_ABSOLUTE            7 (to 14)

# retourner `s`
32 LOAD_FAST                1 (s)
34 RETURN_VALUE

Dans ce cas spécifique, l’utilisation du générateur a ralenti le programme. Comme nous pouvons le voir, le code de la boucle du générateur est identique à la seconde version. La première version fait juste un travail supplémentaire avec l’objet générateur, ce qui nécessite des directives CALL_FUNCTION (lentes). Notez qu’en pratique, les générateurs sont généralement plus rapides grâce à leur nature paresseuse. C’est juste que dans ce cas, le surplus n’en vaut pas la peine.

Quoi qu’il en soit, la différence de performance entre les deux versions ( $\approx 1$s) est insignifiante comparée à la différence globale entre Python et C/Rust. Même avec la version plus rapide, le code Rust est $\approx 65$ fois plus rapide que Python.

Pourquoi ? Principalement parce que Python est un langage interprété et faiblement typé. Cela signifie que l’interpréteur Python (CPython) doit traduire le code Python en quelque chose de facilement exécutable par votre ordinateur, à la volée. Il fait cela en compilant rapidement le code Python en une forme généralisée d’assembleur, appelée bytecode Python, que nous venons d’examiner.

Celui-ci est ensuite exécuté par l’interpréteur. Cette étape de compilation est le premier coup porté aux performances que subit Python. Puisque les langages comme Rust ne nécessitent qu’une seule compilation du programme, ce temps n’est pas comptabilisé dans le temps d’exécution d’un programme. Mais le plus gros coup porté aux performances vient de la qualité médiocre du bytecode, agnostique à l’architecture plutôt qu’optimisé nativement. C’est simplement une partie de la nature des langages interprétés, puisqu’ils ne peuvent pas se permettre de passer autant de temps sur une compilation de haute qualité.

Alors, comment peut-on écrire du Python rapide ? Eh bien, on ne peut pas. Mais on peut appeler du code C rapide, via des bibliothèques hautement optimisées comme Numpy. Ces bibliothèques contiennent des fonctions C pré-compilées et vectorisées qui vous permettent de contourner complètement l’interpréteur Python.

Utilisons Numpy

import numpy as np

def basel_np(N: int) -> float:
    # [1, 1, ..., 1]
    ones = np.ones(N - 1)
    # [1, 2, ..., N]
    r = np.arange(1, N)
    # [1, 1/2, ..., 1/N]
    div = ones / r
    # [1, 1/4, ..., 1/N^2]
    inv_squares = np.square(div)
    # ~ pi^2/6
    return float(np.sum(inv_squares))

>>> f(100000000) # 10^8
Fonction 'basel_np' exécutée en 0.460317 secondes.
1.6449340568482196

Waouh, cela a tourné $\approx 13$ fois plus vite ! Regarder le bytecode dans ce cas ne nous apprendra pas grand-chose, car il ne contiendra que quelques CALL_FUNCTION vers numpy, qui fait le vrai travail. Voyons quelle ligne prend le plus de temps :

def basel_np(N: int) -> tuple[float, list[float]]:
    times = []

    # Chronométrer une seule étape
    start = time.perf_counter()
    ones = np.ones(N - 1)
    end = time.perf_counter()
    step_time = end - start
    times.append(step_time)

    # Code de chronométrage restant omis
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    ret = np.sum(square)

    return ret, times

Réfléchissons à ces chiffres une minute. Les étapes de création des uns/plage n’impliquent pas de calculs intensifs. Pourtant, elles prennent presque autant de temps que les étapes “plus difficiles”, comme la mise au carré et la division. Étonnamment, la somme du tableau final est l’étape la plus rapide ! Cela suggère que le goulot d’étranglement ici n’est pas le CPU, mais l’accès à la mémoire. Voyons comment les performances évoluent, en faisant varier $N$ de $10^7$ à $10^9$.

Dans ce graphique, le bord de la ligne supérieure représente le temps total pris, et l’espace entre deux lignes consécutives représente le temps pour chaque étape.

Nous pouvons voir que

1. Le temps total augmente de manière non linéaire, même si l’algorithme est en $O(n)$

2. L’étape divide prend le plus de temps, avec des augmentations brutales vers $\approx 3\times 10^8$ et $\approx 7\times 10^8$.

C’est logique, car contrairement aux autres opérations, np.divide doit accéder à deux grands tableaux simultanément, et écrire le résultat dans un nouveau tableau. Cela signifie de nombreux accès à la mémoire principale, et possiblement des lectures depuis le SSD, ce qui est extrêmement lent.

Diffusion (Broadcasting)

Il existe en fait une optimisation pour ce genre de problèmes intégrée directement dans Numpy, appelée diffusion (broadcasting). Celle-ci “projette” virtuellement des vecteurs de petite taille dans des tailles plus grandes pour les opérations vecteur-vecteur.

def basel_np_broadcast(N) -> float:
    ones = 1
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    # Comme si c'était [1, 1, ..., 1] / [1, 4, ..., N^2]
    return float(np.sum(square))

Cela nous permet d’économiser beaucoup d’espace précieux dans le cache. Exécutons les mêmes métriques que précédemment. Avec $N=10^8$ :

À partir de maintenant, je désignerai la version avec diffusion comme “la solution Numpy”.

Bien que ce soit une amélioration significative, nous observons toujours ces pics de latence. Essayons de résoudre cela.

Réflexion sur la mémoire

Pour améliorer l’utilisation de la mémoire de la fonction, nous pourrions traiter par blocs, en effectuant les mêmes opérations sur de petits segments de la plage à la fois, puis en additionnant le tout à la fin. Cela nous permettrait de ne garder qu’un seul segment en mémoire à la fois, et devrait conduire à une meilleure utilisation du cache tout en bénéficiant toujours du code vectorisé de Numpy.

Modifions la fonction pour prendre une plage de $N$ :

# [N1, N2], inclusif
def basel_np_range(N1: int, N2: int) -> float:
    # Code de chronométrage omis
    ones = 1
    r = np.arange(N1, N2 + 1)
    div = ones / r
    squares = np.square(div)
    return float(np.sum(squares))

Et écrivons une fonction pour additionner tous les blocs.

def basel_chunks(N: int, chunk_size: int) -> float:
    # Code de chronométrage omis
    s = 0.0
    num_chunks = N // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(i * chunk_size + 1, (i + 1) * chunk_size)
    return s

Avec $N=10^8$ :

La fonction 'basel_chunks' s'est exécutée en 0.108557 secondes.
1.6449340568482258

Parfait ! Elle s’exécute $\approx 3$ fois plus vite que la solution Numpy.

En prime, le résultat sera en réalité légèrement plus précis, en raison de la nature des flottants IEEE 754. Examinons comment les performances évoluent de la même manière que précédemment, en gardant la taille des blocs constante.

Tout d’abord, regardez l’axe des ordonnées. Nous travaillons maintenant à l’échelle de secondes, et non de minutes. Nous constatons également que le temps d’exécution augmente linéairement, ce qui est cohérent avec la complexité temporelle en $O(n)$ de l’algorithme. Nous pouvons théoriser qu’avec une chunk_size de 20000, toutes les informations peuvent tenir dans le cache, évitant ainsi les pics de temps d’exécution dus aux accès hors cache.

Essayons de faire varier la chunk_size dans la plage $[5\times 10^2,10^6]$, en gardant $N$ constant à $10^9$.

D’après la figure, nous observons des gains de performance jusqu’à une chunk_size d’environ $\approx 51000$, au-delà de laquelle des pics de latence apparaissent, probablement causés par des défauts de cache.

Du calcul sur un coin de table

Chaque float64 utilise 64 bits, soit 8 octets. Nous travaillons avec 3 tableaux de $N$ float64, donc l’utilisation mémoire pour un appel est de $51000\times 8\times 3=1224000\approx 1.2$ Mo.

Mon M1 Pro a les spécifications suivantes :

Cela suggère que l’espace cache maximum $L1+L2$ disponible pour la fonction est d’environ $\approx 1.2$ Mo, et si nous le dépassons, nous devons attendre le cache $L3$.

Multiprocessing

Essayons maintenant de faire en sorte que l’ordinateur place plus du tableau dans le cache. Une possibilité est d’exécuter la fonction sur tous les cœurs, afin de pouvoir utiliser plus de cache $L2$ pour notre fonction. Essayons cela.

D’abord, modifions basel_chunks pour qu’elle prenne une plage de $N$ en entrée.

# (N1, N2]
def basel_chunks_range(N1: int, N2: int, chunk_size: int):
    # Code de chronométrage omis
    s = 0.0
    num_chunks = (N2 - N1) // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(N1 + i * chunk_size + 1, N1 + (i + 1) * chunk_size)
    return s

Passons maintenant au vrai multiprocessing :

from multiprocessing import Pool

def basel_multicore(N: int, chunk_size: int):
    # 10 cœurs sur ma machine
    NUM_CORES = 10 
    N_PER_CORE = N // NUM_CORES
    Ns = [
        (i * N_PER_CORE, (i + 1) * N_PER_CORE, chunk_size)
        for i in range(NUM_CORES)
    ]
    # traite 10 lots en parallèle
    with Pool(NUM_CORES) as p:
        result = p.starmap(basel_chunks_range, Ns)
    return sum(result)

En arrière-plan, le module multiprocessing de Python « pickle » l’objet fonction et lance 9 nouvelles instances de python3.10, qui exécutent toutes la fonction avec $N=10^9$ et num_chunks = 50000.

Comparons les performances.

La fonction 'basel_multicore' s'est exécutée en 1.156558 seconde(s).
1.6449340658482263
La fonction 'basel_chunks' s'est exécutée en 1.027350 seconde(s).
1.6449340658482263

Oups, c’est pire. Cela est probablement dû au fait que le travail initial effectué par le multiprocessing est coûteux, il ne sera donc bénéfique que si la fonction prend relativement longtemps à s’exécuter.

En augmentant $N$ à $10^{10}$ :

La fonction 'basel_multicore' s'est exécutée en 2.314828 secondes.
1.6449340667482235
La fonction 'basel_chunks' s'est exécutée en 10.221904 secondes.
1.6449340667482235

Voilà ! C’est $\approx 5$ fois plus rapide. Essayons avec $10^{11}$

La fonction 'basel_multicore' s'est exécutée en 13.844876 secondes.
multicore 1.6449340668379773
La fonction 'basel_chunks' s'est exécutée en 102.480372 secondes.
chunks 1.6449340668379773

C’est $\approx 8$ fois plus rapide ! Lorsque $N$ augmente, le rapport monocœur sur 10 cœurs devrait approcher 10:1.

L’axe des x est sur une échelle logarithmique, c’est pourquoi les temps d’exécution qui augmentent linéairement semblent exponentiels

On peut voir que l’utilisation du multiprocessing a une surcharge d’environ $\approx 1$ seconde, il n’y a donc un gain de performance que lorsque le temps d’exécution par blocs est plus élevé. Cela se produit vers $\approx N=1.3\times 10^9$. Au-delà, cependant, la différence est considérable.

Par expérimentation (non montrée ici), j’ai trouvé qu’une chunk_size d’environ 50000 est également optimale pour le code de multiprocessing, ce qui nous indique que le système d’exploitation impose certaines restrictions sur l’utilisation du cache L2 pour un seul cœur.

Résumé des résultats

Comment se compare-t-il aux langages du dernier article ?

Extrêmement bien ! Le code Numpy par blocs est la solution séquentielle la plus rapide ! Et rappelez-vous, Python a tout un interpréteur qui fonctionne en arrière-plan pendant qu’il effectue les calculs.

Comparaison de Rust et Python multicœur :

Clairement, la méthode de parallélisation de Rust (via les threads du système d’exploitation) est bien plus efficace pour les petits $N$ que le parallélisme basé sur les processus de Python. Mais la différence entre les deux diminue à mesure que $N$ augmente.

Conclusion

Si ce n’était pas clair, cet article entier n’était qu’un exercice pour apprendre un peu comment Python fonctionne sous le capot. S’il vous plaît, n’essayez pas d’impressionner votre manager en hyper-optimisant votre code Python. Il existe presque toujours une meilleure solution, comme

from math import pi

def basel_smart() -> float:
    return pi*pi / 6

Ta-da ! Des ordres de grandeur plus rapide et infiniment plus simple que toutes les autres fonctions.

Quand on pense à la performance, il faut être prudent. Ne supposez pas que la performance est un problème avant de l’avoir mesuré et d’en être certain. Si c’est le cas, cherchez toujours s’il existe une manière plus intelligente de résoudre le problème avant de vous précipiter pour taper import multiprocessing et de brusquer votre chemin à travers un problème autrement simple.

Aussi, si votre application est vraiment si critique en performance, vous ne devriez probablement pas l’écrire en Python. Python a été créé comme un outil qui excelle dans le prototypage, pas dans la vitesse d’exécution. Mais comme vous pouvez le voir, il n’est pas impossible d’obtenir des performances excellentes.

Quoi qu’il en soit, j’espère que vous avez apprécié mon premier vrai article sur ce site. Si vous trouvez une « solution » Python plus rapide à ce « problème », envoyez-moi un email !

Journal des mises à jour