FLEXPART sur GPU : WebGPU pour aller partout, CUDA pour aller vite

Peut-on accélérer un modèle de dispersion atmosphérique de référence sur carte graphique, sans enfermer l’équipe dans l’écosystème NVIDIA ? C’est la question qui nous a menés à FLEXPART-GPU : un portage open source de FLEXPART en Rust et WebGPU, publié sur GitHub. Le projet n’est pas affilié à l’équipe officielle FLEXPART ; il s’en inspire, s’en dérive partiellement, et vise surtout la validation numérique et l’exploration de performances sur GPU grand public.

Le pari : exploiter WebGPU (via wgpu) pour que le même binaire tourne sur NVIDIA, AMD, Intel ou Apple Silicon, plutôt que de viser CUDA d’emblée. Le débit est nettement supérieur au Fortran de référence. Côté physique, l’advection et la diffusion verticale sont bien alignées sur les cas de test publiés ; la turbulence horizontale reste, elle, un point ouvert. Et pour le calcul intensif pur, CUDA reste difficile à battre sur une carte NVIDIA.

WebGPU nous a donné la portabilité. CUDA nous donne encore la souplesse quand chaque cycle compte.

FLEXPART, en bref

FLEXPART est un modèle lagrangien de dispersion atmosphérique : au lieu de résoudre une équation d’advection-diffusion sur une grille fixe, il suit un grand nombre de particules computationnelles transportées par le vent moyen et perturbées par la turbulence. Les concentrations se reconstruisent en comptant les masses de particules dans des cellules de sortie.

Cette approche convient particulièrement aux sources ponctuelles (accidents industriels, sites Seveso), aux terrains complexes et aux matrices source-récepteur en mode rétro (backward). Le modèle est maintenu par une communauté de recherche active ; la version de référence est historiquement en Fortran (aujourd’hui v10/v11).

flexpart-gpu en reprend les composantes cœur du mode forward : advection lagrangienne de Petterssen sur le vent résolu, paramétrisation Hanna (1982) de la couche limite, modèle stochastique de Lagrange (Ornstein-Uhlenbeck) avec condition de well-mixed de Thomson (1987), sous-pas vertical IFINE, dépôts secs (probabilité de survie à partir de v_d) et dépôt humide (scavenging). Ne sont pas encore portés : convection profonde (convmix, schéma d’Emanuel), grilles de sortie imbriquées, ni une validation sur cas réel type ETEX. Le mode backward FLEXPART existe dans le dépôt à titre expérimental, mais n’est pas au centre des rapports publiés.

Ce que nous avons fait

flexpart-gpu est un binaire Rust autonome, pas un plugin Fortran. Le Fortran sert d’oracle de validation : mêmes entrées, comparaison des sorties. Le cœur de calcul vit côté GPU en WGSL (WebGPU Shading Language), orchestré par wgpu.

À chaque pas de temps, le pipeline de production enchaîne un diagnostic de couche limite puis quatre dispatches GPU de physique particulaire, le tout dans un seul encodeur de commandes :

• Diagnostics PBL : friction u*, vitesse convective w*, longueur d’Obukhov L, hauteur de mélange h par cellule de grille (équivalent Fortran : calcpar / Richardson).
• Advection : vent résolu uniquement, schéma prédicteur-correcteur de Petterssen, interpolation trilinéaire en espace et linéaire entre deux bracket météo (wind_t0, wind_t1).
• Hanna + Langevin fusionnés : σ et tempscales de Lagrange (Hanna), puis mise à jour stochastique des vitesses turbulentes u′, v′, w′ avec sous-pas vertical, recalcul hanna_short entre sous-pas et réflexion à la hauteur de mélange (Thomson).
• Dépôt sec : perte de masse probabiliste dans la couche de surface (2·h_ref).
• Dépôt humide : scavenging par les précipitations.

C’est la même séquence logique que dans advance.f90 (Fortran), découpée en kernels pour le GPU : le vent moyen transporte les particules, la turbulence les disperse.

Le CPU gère l’orchestration : lecture des fichiers météo (GRIB2, NetCDF), injection des particules, boucle temporelle, écriture des sorties. Le GPU prend le relais sur tout ce qui est « par particule » : physique, compaction des particules actives, gridding des concentrations.

Un mode validation (FLEXPART_GPU_VALIDATION=1) sépare Hanna et Langevin en deux kernels distincts (cinq dispatches de physique au lieu de quatre), pour faciliter le débogage et la comparaison shader par shader. La physique est identique au chemin fusionné ; le rapport de validation Fortran/GPU repose sur le même moteur, avec readback activé pour exporter les positions particulaires.

Pourquoi WebGPU plutôt que CUDA ?

Notre objectif n’était pas seulement « aller plus vite ». Nous voulions quelque chose de portable : lancer une simulation sur le GPU disponible, qu’il soit NVIDIA, AMD, Intel ou Apple Silicon, sans recompiler pour chaque vendeur ni dépendre d’un cluster MPI pour les campagnes du quotidien.

WebGPU, via wgpu, abstrait Vulkan (Linux/Windows), Metal (macOS) et DX12. En pratique, sur une station de travail Linux avec carte NVIDIA, c’est souvent Vulkan qui tourne sous le capot. Mais le même code peut aussi s’exécuter sur un Mac M-series ou une carte AMD, ce qui change la donne pour la diffusion opérationnelle et les tests terrain.

Au-delà de la vitesse brute, le dépôt apporte une surface d’ingénierie plus sûre (typage Rust, tests intégrés, protocoles de benchmark versionnés) et un chemin de migration progressif : comparer côte à côte avec le Fortran avant de basculer.

Pourquoi le méga-kernel a été abandonné

Notre première intuition était d’embarquer toute la physique d’un pas de temps dans un méga-kernel unique (particle_step.wgsl) : advection, Hanna, Langevin, dépôts secs et humides en une seule passe GPU, sans barrière inter-kernels ni buffer intermédiaire.

Nous l’avons tenté, puis abandonné. La cause documentée dans le dépôt : une pression registres sévère sur le GPU cible : le shader devient trop lourd pour tenir en registres et le débit s’effondre. Ce n’est pas un problème abstrait de « shader trop gros » : c’est une contrainte matérielle concrète, comparable à ce qu’on rencontre aussi en CUDA quand on fusionne trop de physique dans un seul kernel.

Une seconde contrainte WebGPU joue aussi : le méga-kernel requiert 9 storage buffers par stage shader (maxStorageBuffersPerShaderStage), là où le minimum WebGPU est 8. Les GPU discrets récents dépassent ce seuil (souvent ≥16), mais ce n’est pas garanti partout, et ce n’est qu’un des freins, pas la cause principale d’abandon.

La solution retenue : découper en plusieurs kernels, chacun focalisé sur une étape physique. En mode production, les quatre passes (advection → Hanna+Langevin fusionnés → dépôts) partagent les mêmes buffers GPU et sont encodées dans un seul command encoder : pas de transfert CPU↔GPU entre passes tant qu’on reste en mode « fire-and-forget ». Les allers-retours hôte↔device n’interviennent qu’aux instants de readback (validation, sorties), pas entre chaque kernel.

Le chemin fusionné (langevin_fused.wgsl) va plus loin : au lieu de séparer Hanna puis Langevin (deux dispatches, un buffer HannaParams intermédiaire d’environ 64 Mo pour 1 M particules), Hanna est recalculé en ligne dans le kernel Langevin. Même physique, une barrière en moins.

En CUDA, on aurait peut-être poussé plus loin la fusion. En WebGPU/WGSL, la découpe répond à des contraintes réelles (registres, bindings), autant qu’à un choix de design.

WebGPU et ses limites (dont la mémoire workgroup)

WebGPU impose aussi des plafonds qui comptent pour d’autres kernels, même si ce n’est pas ce qui a tué le méga-kernel FLEXPART :

• Mémoire workgroup (var<workgroup>) : plafond minimal de 16 Ko par kernel (maxComputeWorkgroupStorageSize). Limite les accumulations locales (histogrammes, tuiles) : pas le cas principal de nos kernels particulaires, qui s’appuient surtout sur des storage buffers.
• Espace private en WGSL : plafond de 8 Ko de variables locales par fonction, à surveiller sur les shaders très profonds.
• Bindings et layouts figés à la compilation du pipeline ; moins de liberté qu’en CUDA pour réorganiser les données à la volée.
• Limites de buffers (maxStorageBufferBindingSize) variables selon le matériel, parfois sévères sur GPU intégrés.

WebGPU face à CUDA pour le calcul scientifique

WebGPU n’est pas mauvais. Pour notre cas d’usage (millions de particules indépendantes, arithmétique flottante régulière), il tient la route et les gains sont réels. Mais la gestion de l’information y est nettement moins souple que CUDA : pipelines et bindings rigides, plafonds matériels variables, et un écosystème de profiling moins mature (Nsight Compute vs les outils WebGPU encore jeunes ; le dépôt expose toutefois un profiling interne via FLEXPART_GPU_PROFILE).

Pour une équipe qui vit dans l’écosystème NVIDIA (clusters HPC, cuBLAS, NCCL, écosystème PyTorch/JAX), CUDA reste le standard de facto. Sur une RTX bien configurée, avec la liberté de structurer la mémoire et les kernels comme on l’entend, il est difficile de faire mieux pour du calcul scientifique intensif.

Notre conclusion pragmatique : WebGPU pour la portabilité et l’accessibilité (une station, n’importe quelle carte récente). CUDA pour la performance maximale et la flexibilité algorithmique quand on peut se permettre de viser NVIDIA.

Résultats mesurés

Les benchmarks du dépôt comparent le binaire GPU (mode production, sans readback à chaque pas) à FLEXPART 10.4 Fortran sérial (mono-thread, sans MPI), sur un scénario synthétique identique. Il s’agit d’un gain wall-clock, mesuré sur une RTX 4070 Laptop (Vulkan) dans les rapports publiés, pas d’extrapolation à d’autres machines ou cas opérationnels.

Validation scientifique

La validation repose sur des scénarios synthétiques documentés dans le dépôt : vent uniforme, source ponctuelle, traceur inerte, dépôts et convection désactivés (les kernels de dépôt existent, mais ne sont pas couverts par ces rapports), sous-pas vertical IFINE=4 des deux côtés. Référence Fortran : FLEXPART v10.4 (mono-thread). La métrique retenue est la comparaison des positions particulaires brutes (Fortran partposit vs readback GPU), pas la grille de concentration : l’opérateur Fortran conccalc ≠ le gridding GPU direct.

Sur un cas à 10 000 particules et 6 h de simulation (rapport de validation) :

• Advection horizontale : centre de masse à 5,5 km du Fortran (seuil interne : 10 km), soit environ 5 % d’écart relatif sur un déplacement total d’environ 120 km en 6 h (vent moyen ~5,8 m/s).
• Diffusion verticale : hauteur moyenne à +22 m (seuil : 200 m), ratio d’écart-type vertical 0,94 : la chaîne Langevin + Hanna + réflexion PBL tient la route.
• Diffusion horizontale : écart principal. L’étalement transversal GPU est environ ×2 à ×2,5 celui du Fortran (σ_lon ×2,0 ; σ_lat ×2,5). Cause identifiée dans le dépôt : écart de paramétrisation σ_v / planchers de turbulence et de u* (prescrit côté GPU vs calculé côté Fortran sur le cas test), pas un artefact de gridding. Prochaine étape : aligner sur hanna1.f90 Fortran.

À 1 M particules sur le même scénario physique, les ordres de grandeur se confirment : centre de masse à 5,4 km, Δz moyen +19 m, ratio σ_z 0,92, mais l’écart horizontal persiste (rapport benchmark-scientific-validation). Les 224 tests automatisés (parité CPU/GPU, conservation de masse, invariants) passent au vert.

D’autres différences sont documentées et assumées : générateur aléatoire Philox vs Fortran, précision f32 côté GPU (justifiée par la convergence Monte Carlo 1/√N), plafond de couche limite dur côté GPU vs comportement plus souple en Fortran (~7 % de particules au-dessus du mélange côté Fortran), hauteur de mélange prescrite vs calcul Richardson. Ce ne sont pas des bugs silencieux ; ce sont des écarts de modèle qu’il faut connaître avant un usage opérationnel.

Open source et attribution

Le projet est publié sous GPL-3.0-or-later, comme FLEXPART en amont. Le crédit scientifique revient à l’équipe FLEXPART et aux publications de référence (Stohl et al. 2005 ; Bakels et al. 2024 pour la v11). Notre portage est explicitement indépendant et non officiel ; les notices d’attribution sont dans le dépôt (NOTICE.md).

La documentation couvre l’architecture, les fondements scientifiques, les protocoles de benchmark et un guide de démarrage rapide. Docker est fourni pour des builds reproductibles avec accès GPU (Vulkan).

Ce que nous en retenons

• FLEXPART-GPU démontre qu’un modèle de dispersion mature peut être porté sur GPU portable, avec advection et diffusion verticale alignées sur Fortran dans les cas de test publiés, et des gains mesurés autour de ×8 à ×11 à 1 M particules.
• La turbulence horizontale reste à calibrer (σ_lon ~×2, σ_lat ~×2,5 vs Fortran) ; la convection et les grilles imbriquées ne sont pas encore implémentées.
• Le choix WebGPU était le bon pour la portabilité multi-vendeur ; il impose en revanche de fragmenter les kernels (pression registres, contraintes de bindings) plutôt que de tout fusionner en une passe.
• Pour du calcul scientifique intensif sur NVIDIA, CUDA reste la référence : plus de souplesse mémoire, moins de friction entre étapes de calcul.
• Le dépôt est ouvert, documenté et reproductible : un point de départ pour explorer la dispersion atmosphérique accélérée, pas un remplacement officiel de FLEXPART.

Nous poursuivons l’exploration (convection d’Emanuel, grilles imbriquées, comparaisons CUDA), dans la continuité de nos autres travaux sur le GPU appliqué à la modélisation atmosphérique (notre retour d’expérience à quatre voix). Si vous travaillez sur FLEXPART, la dispersion lagrangienne ou le portage GPU de codes legacy, le dépôt est là : github.com/Improba/flexpart-gpu.