Jour 5 — Python sur CPU

26 juin 2026 · Intervenants : Alice Faure, Jean-Marc Colley, Sébastien Valat & Nabil Garroum · Auditorium Marcel Vivargent + sites satellites (dont la CINERI). Le TP est une suite de six notebooks Jupyter (GrayScott2026/day-5/CPU/tutorial/, solutions incluses) — le dossier GPU/ attendra le Jour 7.

Session du matin — mesurer, puis vectoriser

1. Profiler d'abord — temps et mémoire

Fidèle à la règle du Jour 1, le premier notebook (1_Optimization) n'optimise rien : il mesure. Chronométrage (timeit, cProfile), puis — c'est la spécialité du jour, les slides s'appellent gray-scott-python-mem — le profilage mémoire avec tracemalloc : en Python, chaque tableau temporaire est une allocation, et le Gray-Scott naïf en crée à chaque pas de temps.

2. Le GIL — pourquoi les threads ne sauvent pas Python

Le Global Interpreter Lock sérialise les threads Python purs. Le vrai parallélisme CPU passe par les bibliothèques qui libèrent le GIL pendant le calcul natif — NumPy, le prange de Numba, XLA sous JAX — ou par multiprocessing. Aucune accélération du jour ne combat le GIL : elles descendent toutes sous lui, dans du code compilé.

3. Des boucles aux tableaux

Notebooks 2_Numpy et 3_Python_Implementation : le Laplacien s'écrit en slices NumPy (u[:-2, 1:-1] + u[2:, 1:-1] + … − 4*u[1:-1, 1:-1]) — zéro boucle Python, les itérations partent dans les boucles C de NumPy. C'est la ligne de base du benchmark du jour.

Session de l'après-midi — compiler Python

4. Numba — compiler la boucle qu'on a déjà

Notebook 4_Numba_Implementation : on garde la boucle explicite, on ajoute @njit, et LLVM compile la fonction à la première exécution. Idéal quand l'algorithme est naturellement une boucle (stencils !) — et prange la parallélise.

5. JAX — tracer, puis laisser XLA fusionner

Notebook 5_JAX, le morceau de choix du jour. JAX impose une discipline en échange de la vitesse : tableaux immuables (u.at[i, j].set(v) au lieu de l'affectation), pas de vérification d'index (l'erreur silencieuse guette), et surtout des transformations composables — jit, vmap, grad — qui ne fonctionnent que sur des fonctions pures. Le mécanisme derrière jax.jit :

Le notebook détaille les contraintes : arguments statiques à déclarer (static_argnums), formes fixes (chaque nouvelle forme retrace), débogage via jax.debug, et les opérateurs de contrôle de flux (lax.cond, lax.fori_loop) qui remplacent if/for dans le code tracé.

6. Porter Gray-Scott en JAX

Notebook 6_JAX_Implementation : deux versions concurrentes — le stencil générique (convolution 3×3 quelconque) et le stencil 3×3 spécialisé (les neuf termes écrits à la main, que XLA fusionne en un seul noyau). Les deux colonnes JAX du benchmark, c'est elles.

7. Le verdict — un Gray-Scott, quatre vitesses

Chiffres officiels du dépôt (CPU/Benchmarks.md), 32×1000 itérations :

8. Le pont vers le GPU

Cette implémentation est la référence que le Jour 7 porte sur l'accélérateur : JAX rejoue le même code sur GPU, rejoint par CuPy et cuNumeric. C'est aussi exactement la démarche du projet SenLand — profiler PyTorch, porter en JAX, comparer honnêtement.

Le TP — GrayScott2026/day-5/CPU/

Trois façons de suivre, au choix :

# 1) en local, environnement épinglé par pixi
git clone https://gitlab.in2p3.fr/alice.faure/gray-scott-python.git
pixi run jupyter-lab           # ouvre les notebooks de tutorial/

# 2) sur le cluster MUST (LAPP) : https://jupyter.must-dc.cloud
#    → « Gray-Scott Revolutions » → « Python CPU »

# 3) en conteneur (apptainer / podman / docker) : image vscode du cours

tutorial/ pose les exercices, solutions/ les corrige, scripts/gray_scott_utils.py fournit l'I/O commun, et results/ contient une simulation témoin (simulation.h5 + vidéo).

En vidéo — le replay officiel

Sources & matériel officiel

• Le dépôt du cours (notebooks + solutions + benchmarks) : gitlab.in2p3.fr/alice.faure/gray-scott-python
• Les slides du jour (PDF, wiki GitLab de l'école) : 2026-06-gray-scott-python-mem.pdf
• La plateforme Jupyter du cluster MUST : jupyter.must-dc.cloud
• Le conteneur de base de l'école : GrayScottBaseContainer
• Les bibliothèques : numpy.org · numba.pydata.org · docs.jax.dev
• Replays vidéo (YouTube) : Gray Scott Thursdays
• Site de l'école : GrayScott2026