Mon précédent article (qui a été honnêtement créé pour tester le thème de ce site), fournissait 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 essayer 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 de faire 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 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"La fonction '{func.__name__}' a été exécutée en {execution_time:.6f} secondes.")
        return result
    return wrapper

>>> f = time_function(basel)
>>> f(100000000) # 10^8
La fonction 'basel' a été exécutée en 6.641589 secondes.
1.644934057834575

Essayons de la réécrire, mais 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
La fonction 'basel_less_pythonic' a été exécutée en 5.908466 secondes.
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)

# Aller en haut 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 deuxième 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 (lentes) CALL_FUNCTION. Notez qu’en pratique, les générateurs sont généralement plus rapides en raison de leur nature paresseuse. C’est juste que dans ce cas, le surcoût supplémentaire 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 par rapport à 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 de regarder.

Ce bytecode est ensuite exécuté par l’interpréteur. Cette étape de compilation est le premier coup porté aux performances que Python subit. Comme 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 porté aux performances vient de la qualité médiocre du bytecode, qui est agnostique à l’architecture, par opposition à un bytecode nativement optimisé. Cela fait simplement partie de la nature des langages interprétés, car 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 fortement 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.

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
La fonction 'basel_np' a été exécutée en 0.460317 secondes.
1.6449340568482196

Wow, cela a tourné $\approx 13$ fois plus vite ! Regarder le bytecode dans ce cas ne nous dira pas grand-chose, car il ne consistera qu’en 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 pendant une minute. Les étapes de création des uns/plage n’impliquent pas de travail de calcul intensif. Cependant, elles prennent presque le même temps que les étapes “plus difficiles”, comme la mise au carré et la division. Étonnamment, la somme sur le 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 la plus haute représente le temps total pris, et l’espace entre les 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 $O(n)$

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

Cela a du sens, 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 y a en fait une optimisation pour ce genre de problèmes que Numpy intègre appelée diffusion, qui “projette” virtuellement des vecteurs de plus 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 qu’avant. Avec $N=10^8$ :

À partir de maintenant, je ferai référence à la version diffusée comme “la solution Numpy”.

Bien qu’il y ait une amélioration significative, nous voyons toujours ces pics de latence. Essayons de les corriger.

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 morceaux de la plage à la fois, et en additionnant tout à la fin. Cela nous permettrait de ne garder qu’un seul morceau en mémoire à la fois, et espérons-le, de mieux utiliser le cache tout en bénéficiant 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 morceaux.

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' a été exécutée en 0.108557 secondes.
1.6449340568482258

Super ! Cela tourne $\approx 3$ fois plus vite que la solution Numpy.

En prime, la réponse sera en fait légèrement plus précise, en raison de la nature des flottants IEEE 754. Regardons comment les performances évoluent de la même manière que nous l’avons fait avant, en gardant la taille des morceaux constante.

Tout d’abord, regardez l’axe des y. Nous travaillons maintenant à l’échelle des secondes, pas des minutes. Nous voyons é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 être stockées dans le cache, donc il n’y a pas de pics de temps d’exécution pour les accès hors cache.

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

D’après la figure, nous voyons des gains de performance jusqu’à une chunk_size de $\approx 51000$, après quoi il y a des pics de latence probablement causés par des défauts de cache.

Un peu de calcul rapide

Chaque float64 utilise 64 bits, soit 8 octets. Nous travaillons avec 3 tableaux de $N$ float64, donc l’utilisation de la mémoire pour un appel est $51000(8)(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 $\approx 1.2$ Mo, et si nous dépassons cela, nous devons attendre le cache $L3$.

Multiprocessing

Maintenant, essayons de faire en sorte que l’ordinateur stocke plus du tableau dans le cache. Une possibilité est d’essayer d’exécuter la fonction sur tous les cœurs, afin que nous puissions utiliser plus de $L2$ pour notre fonction, alors essayons cela.

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