Optimisation des performances pas si anodine en Python

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 facile à lire 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 la 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 en gérant 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 de complexité 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 de Python. Puisque les langages comme Rust ne nécessitent qu’une seule compilation, 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 mauvaise qualité du bytecode, agnostique à l’architecture par opposition à un code natif optimisé. Ceci fait simplement 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 pouvez-vous écrire du Python rapide ? Eh bien, vous ne pouvez pas. Mais, vous pouvez 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 l’interpréteur Python entier.

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 est la plus coûteuse :

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/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 écoulé, et l’espace entre deux lignes consécutives représente le temps de chaque étape.

Nous pouvons observer 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 brusques 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 type 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 d’obtenir potentiellement 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 qu’auparavant, 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 correspond à 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 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.

Quelques calculs 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(8)(3)=1224000\approx 1.2$ Mo.

Mon M1 Pro a les spécifications suivantes :

Ceci suggère que l’espace de 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 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.

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.
1.6449340658482263
La fonction 'basel_chunks' s'est exécutée en 1,027350 seconde.
1.6449340658482263

Oups, les performances sont moins bonnes. Cela est probablement dû au fait que le travail initial effectué par le multiprocessing est coûteux, donc il ne sera 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 entre le monocœur et les 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, donc il n’y a 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$. En revanche, après cela, 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 cela se compare-t-il aux langages du dernier article ?

Extrêmement bon ! Le code Numpy segmenté est la solution séquentielle la plus rapide ! Et rappelez-vous, Python a un interpréteur entier 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 dans les coulisses. 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.

Quand vous pensez aux performances, soyez prudent. Ne supposez pas que les performances sont un problème avant d’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 termes de performances, 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