Jour 7 — Python sur GPU

30 juin 2026 · Intervenants : Alice Faure, Jean-Marc Colley, Sébastien Valat & Nabil Garroum · quatre sessions Python-GPU s'enchaînent (CuPy → cuPyNumeric → JAX à 14 h → synthèse) · Auditorium Marcel Vivargent + satellites (dont la CINERI). Le TP vit dans GrayScott2026/day-5/GPU/ — trois tutoriels + solutions, et les benchmarks A100 officiels.

1. Ce qui change vraiment : le péage PCIe

Le code array-first du Jour 5 se transpose presque tel quel — ce qui change, c'est la géographie mémoire. Le GPU calcule à ~2 To/s mais se nourrit par un tuyau de ~32 Go/s :

Toute la journée applique la même règle : cp.asarray / jax.device_put une fois au départ, toute la boucle de temps sur le device, asnumpy une fois à l'arrivée.

2. Trois routes vers le même GPU

Session CuPy — NumPy sur CUDA, sans réécriture

Tutoriel 3_Python_GPU_Cupy.md : CuPy reflète l'API NumPy sur CUDA — remplacer numpy par cupy exécute le même stencil sur le GPU.

import cupy as cp
u = cp.asarray(u_host)      # hôte → device, UNE fois
# … mêmes expressions de stencil qu'en NumPy …
u_host = cp.asnumpy(u)      # device → hôte, seulement au besoin

Bonus du TP : CuPy est la seule version où l'I/O HDF5 parallèle est implémentée — écrire les 1000 images pendant que le GPU calcule économise ~4 s (12 s au total sur A100).

Session cuPyNumeric — NumPy distribué

Tutoriel 2_Python_GPU_cuPyNumeric.md : cuPyNumeric (NVIDIA, moteur Legate) exécute du code NumPy sur plusieurs GPU et plusieurs nœuds sans MPI et sans réécriture — le même script, une machine plus grande. Le prix de la généralité se lit dans le benchmark : la version convolve générique est la plus lente du plateau (128 s).

Session JAX (14 h) — le code du Jour 5, re-jitté

Tutoriel 1_Python_GPU_JAX.md : le Gray-Scott JAX du Jour 5 rejoue sans changement — XLA compile le stencil tracé en noyau CUDA, et JAX place les tableaux sur le device par défaut. Les solutions du TP montrent les trois outils qui font la différence :

C'est exactement la boîte à outils du port JAX de SenLand (lax.fori_loop pour fusionner les steps, batch résident sur device).

Le verdict — chiffres A100 officiels

GPU/Benchmarks.md du dépôt : 32×1000 itérations, grille 1920×1080 en float32 :

La boucle se boucle : les 377 s du meilleur CPU du Jour 5 tombent à 18 s sur A100 — ×21, toujours en Python. Et le classement rappelle les leçons de la semaine : la spécialisation du stencil (Jour 5) et la résidence des données (aujourd'hui) pèsent plus que le choix de la bibliothèque.

Le TP — GrayScott2026/day-5/GPU/

# En local (NVIDIA, CUDA ≥ 12, Python 3.10-3.12)
git clone https://gitlab.in2p3.fr/alice.faure/gray-scott-python.git
python -m venv gpu-env && source gpu-env/bin/activate
pip install h5py opencv-python numpy matplotlib scipy \
            "jax[cuda12]" cupy-cuda12x nvidia-cupynumeric

Alternatives officielles : image Docker du cours, ou apptainer sur le cluster MUST (Install_satellite_sites.md pour les sites comme la CINERI). AMD : CuPy et JAX ont des routes ROCm expérimentales — cuPyNumeric non. Sur une petite carte locale (GTX 1650, 4 Go), on réduit la grille : la leçon de résidence des données reste identique.

En vidéo — le replay officiel

Sources & matériel officiel

• Le dépôt du cours (tutoriels GPU/tutorial/, solutions, benchmarks A100) : gitlab.in2p3.fr/alice.faure/gray-scott-python
• Les bibliothèques : docs.cupy.dev · docs.nvidia.com/cupynumeric · docs.jax.dev
• La plateforme MUST : jupyter.must-dc.cloud
• Replays vidéo (YouTube) : Gray Scott Thursdays
• Site de l'école : GrayScott2026