Optimisation des performances pas si anodine en Python

Mon article précédent (qui a été créé honnêtement pour tester le thème de ce site), présentait quelques extraits de code qui calculaient $N$ termes de la somme des carrés inverses. 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 en séquentiel, Python a mis 70 secondes ! Donc, dans cet article, nous allons tenter d’obtenir des chiffres un peu 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, 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 aussi 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 à celui de la seconde version. La première version fait juste un travail supplémentaire pour gérer 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 surcoût 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 aux performances vient de la production d’un bytecode de qualité médiocre, agnostique de l’architecture, par opposition à un bytecode optimisé nativement. Cela fait simplement partie de la nature des langages interprétés, car ils ne peuvent pas se permettre de passer tant 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 vectorisées pré-compilées qui 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, puisqu’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 ones/de la plage n’impliquent pas de calculs lourds. Pourtant, elles prennent presque le même 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 de chaque étape.

Nous pouvons voir que

1. Le temps total augmente de manière non linéaire, bien que l’algorithme soit en $O(n)$

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

Cela 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 beaucoup d’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 type de problèmes que NumPy intègre nativement, appelée diffusion (broadcasting). Elle “projette” virtuellement des vecteurs de petite taille dans des dimensions 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 mesures que précédemment. Avec $N=10^8$ :

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

Bien qu’il s’agisse d’une amélioration significative, nous observons toujours ces pics de latence. Essayons de régler cela.

Réflexion sur la mémoire

Pour améliorer l’utilisation de la mémoire de la fonction, nous pourrions travailler bloc par bloc, en effectuant les mêmes opérations sur de petits segments de l’intervalle à 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 un intervalle 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 sommer tous les segments.

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$ :

Function 'basel_chunks' executed in 0.108557 seconds.
1.6449340568482258

Excellent ! 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 qu’avant, en gardant la taille des segments 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 $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 l’intervalle $[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 de cache maximum $L1+L2$ disponible pour la fonction est d’environ $1.2$ Mo, et si nous dépassons cette limite, nous devons attendre le cache $L3$.

Multiprocessing

Essayons maintenant de faire en sorte que l’ordinateur utilise davantage le cache pour le tableau. 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.

Premièrement, modifions basel_chunks pour qu’elle prenne un intervalle 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 véritable 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)
    ]
    # traiter 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 toutes exécutent 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 secondes.
1.6449340658482263
La fonction 'basel_chunks' s'est exécutée en 1.027350 secondes.
1.6449340658482263

Oups, c’est pire. Cela est probablement dû au fait que le travail d’initialisation du 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 environ $\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 environ $\approx 8$ fois plus rapide ! À mesure que $N$ augmente, le rapport entre le temps mono-cœur et 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 supérieur. Cela se produit à partir d’environ $\approx N=1.3\times 10^9$. En revanche, après ce point, 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 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 cela se compare-t-il aux langages du dernier article ?

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

Comparaison Rust et Python multicœur :

Il est clair que 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 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.

Lorsque vous pensez aux performances, il faut être prudent. Ne supposez pas que la performance est un problème avant de mesurer 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 forcer votre chemin à travers un problème autrement simple.

De plus, si votre application est vraiment si critique en termes de 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 impressionnantes.

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