GPU acceleration of plane-wave density functional theory calculations in Abinit

Ce rapport présente le portage du code Abinit sur des architectures multi-GPU pour les calculs DFT en ondes planes, en détaillant les révisions algorithmiques nécessaires et en comparant les performances de deux méthodes de diagonalisation (LOBPCG et filtrage par polynômes de Chebyshev) sur des nœuds hétérogènes CPU-GPU.

L'essentiel

🚀 Accélérer la simulation de la matière : Le grand saut de l'ordinateur vers la "carte graphique"

Imaginez que vous êtes un architecte qui doit construire des gratte-ciels atomiques. Pour cela, vous devez calculer comment chaque atome interagit avec ses voisins. C'est ce qu'on appelle la Théorie de la Fonctionnelle de la Densité (DFT). C'est un travail de titan qui demande des calculs mathématiques gigantesques.

Pendant des années, les scientifiques ont utilisé des processeurs classiques (CPU), qui sont comme des chefs d'orchestre très intelligents mais qui ne peuvent jouer qu'un instrument à la fois. Ils sont excellents pour gérer la logique complexe, mais ils deviennent vite fatigués quand il faut jouer des milliers de partitions en même temps.

Récemment, les ordinateurs ont ajouté des cartes graphiques (GPU). Ces GPU sont comme des armées de milliers de petits musiciens. Ils ne sont pas très intelligents individuellement, mais s'ils doivent tous jouer la même note en même temps, ils sont imbattables.

Le but de ce papier est de raconter comment l'équipe derrière le logiciel Abinit (un outil célèbre pour simuler la matière) a réussi à faire passer ses calculs du "chef d'orchestre" (CPU) à l'"armée de musiciens" (GPU).

---

🎻 Le problème : La musique ne sonne pas bien sur le GPU

Au début, essayer de faire tourner Abinit sur un GPU était comme essayer de faire jouer une symphonie complexe à une armée de petits musiciens en leur donnant des partitions différentes toutes les 2 secondes.

• Le problème : Les GPU détestent attendre. Si le chef d'orchestre (le CPU) doit leur donner une nouvelle note toutes les 2 secondes, les musiciens passent 99% de leur temps à attendre. C'est inefficace.
• La solution : Il faut leur donner une énorme pile de partitions à jouer d'un coup, sans interruption. C'est ce qu'on appelle le "traitement par lots" (Batch processing).

🛠️ Comment ils ont réorganisé le travail ?

Les auteurs ont dû réinventer la façon dont le logiciel fonctionne pour qu'il corresponde à la force des GPU. Voici les trois clés de leur réussite :

1. Le "Batching" : La méthode du camion de déménagement

Au lieu de transporter un seul meuble à la fois (un seul calcul à la fois), le logiciel a appris à charger tout un camion de meubles (des milliers de calculs) et à les déplacer d'un seul coup.

• Analogie : Imaginez que vous devez déplacer des livres d'une étagère à une autre.
  • Ancienne méthode (CPU) : Vous prenez un livre, vous marchez, vous le posez, vous revenez. Répétez 10 000 fois.
  • Nouvelle méthode (GPU) : Vous remplissez un chariot de 100 livres, vous faites un seul voyage, et vous les posez tous. C'est beaucoup plus rapide !

2. Le "Séjour sur place" : Ne pas faire de navette

Dans les calculs précédents, les données (les "ondes électroniques") devaient faire des allers-retours constants entre la mémoire du processeur et celle de la carte graphique. C'est comme si le chef d'orchestre devait courir jusqu'à la scène pour donner une note, puis revenir à son pupitre, puis courir à nouveau.

• La nouvelle astuce : Une fois les données arrivées sur la carte graphique (le GPU), elles y restent. Tout le travail de calcul se fait là-bas. On ne les renvoie au processeur que lorsque le travail est fini. Cela économise un temps fou.

3. Choisir la bonne stratégie de calcul : Le filtre vs. L'escalier

Pour trouver les solutions mathématiques (les états électroniques), il existe deux grandes méthodes. Les auteurs ont comparé deux stratégies :

• LOBPCG (La méthode de l'escalier) : C'est comme monter un escalier marche par marche. À chaque marche, vous devez vérifier votre équilibre (une communication complexe entre les processeurs). C'est précis, mais lent sur un GPU car il y a trop de vérifications.
• Filtrage de Chebyshev (Le filtre à café) : C'est comme passer un grand tamis. Vous lancez tout le contenu d'un coup, et le filtre sépare ce qui est important de ce qui ne l'est pas. Cette méthode demande beaucoup de calculs simples (parfaits pour le GPU) mais très peu de vérifications.
• Résultat : Sur les GPU, la méthode du filtre (Chebyshev) est bien plus rapide et consomme moins d'énergie que la méthode de l'escalier.

---

📊 Les résultats : Qui gagne ?

L'équipe a testé leur nouveau logiciel sur deux types de supercalculateurs :

1. NVIDIA (les cartes graphiques vertes) : C'est un succès total. Le logiciel est 17 fois plus rapide que sur les vieux processeurs classiques pour certains calculs. De plus, cela consomme beaucoup moins d'électricité pour faire le même travail.
2. AMD (les cartes graphiques rouges) : C'est aussi une réussite, mais un peu moins spectaculaire que sur NVIDIA. Les cartes graphiques AMD sont bonnes, mais le logiciel doit encore être un peu plus optimisé pour elles.

Le point clé : Le plus gros goulot d'étranglement (le moment où ça ralentit) n'est plus le calcul lui-même, mais la partie où l'on doit "redresser" les résultats (l'étape de Rayleigh-Ritz). C'est comme si l'armée de musiciens jouait à toute vitesse, mais que le chef d'orchestre prenait trop de temps pour corriger la partition entre deux mesures.

🌍 Pourquoi c'est important pour nous ?

Grâce à ce travail, les scientifiques peuvent maintenant :

• Simuler des matériaux plus complexes (pour créer de meilleures batteries, des panneaux solaires plus efficaces, ou de nouveaux médicaments).
• Le faire beaucoup plus vite.
• Le faire en consommant moins d'énergie, ce qui est crucial pour l'environnement.

En résumé

Ce papier raconte l'histoire d'une équipe qui a appris à parler le langage des cartes graphiques. Au lieu de forcer un vieux logiciel à courir sur des jambes neuves, ils ont réorganisé toute la course pour qu'elle corresponde à la vitesse et à la puissance de ces nouvelles jambes. Résultat : la science des matériaux entre dans une nouvelle ère de rapidité et d'efficacité énergétique.

Résumé technique

1. Problématique

Les simulations de grande envergure en science des matériaux, basées sur la théorie de la fonctionnelle de la densité (DFT) et l'approche des ondes planes, nécessitent le calcul de milliers d'états électroniques. Cela implique de résoudre le problème de Kohn-Sham, qui se traduit par la diagonalisation itérative de matrices Hamiltoniennes de très grande taille.

• Défi principal : Exploiter efficacement les architectures de calcul haute performance (HPC) modernes, en particulier les nœuds hétérogènes combinant CPU et GPU.
• Goulot d'étranglement : Les algorithmes traditionnels de diagonalisation (comme la méthode LOBPCG) comportent des étapes fortement dépendantes de la mémoire (orthogonalisation, procédure de Rayleigh-Ritz) et des communications MPI fréquentes, ce qui limite leur efficacité sur les GPU. De plus, le transfert de données entre l'hôte (CPU) et le dispositif (GPU) peut annuler les gains de performance si le code n'est pas optimisé pour résider entièrement sur le GPU.

2. Méthodologie

Les auteurs ont porté le code Abinit (version 10.6+) sur des architectures multi-GPU en adoptant une approche hybride et algorithmique.

• Modèle de programmation : Utilisation du modèle MPI + OpenMP avec le offloading (délégation de tâches) vers le GPU via les directives OpenMP 5.0+. Cela permet une portabilité entre les architectures NVIDIA (CUDA) et AMD (ROCm/HIP).
• Gestion de la mémoire et des données :
  • Résidence GPU : La fonction d'onde ($\Psi$) est transférée une seule fois au début de l'itération de champ auto-cohérent (SCF) et reste en mémoire GPU tout au long du processus de diagonalisation.
  • Traitement par lots (Batching) : Les opérations sont regroupées pour traiter plusieurs bandes électroniques simultanément. Cela remplace les boucles séquentielles par des appels de noyaux (kernels) massivement parallèles, optimisant l'utilisation de la bande passante mémoire et réduisant la surcharge des appels de noyaux.
  • Layouts de données : Gestion intelligente des transpositions de matrices (distribution par lignes vs par colonnes) pour adapter les données aux besoins des FFT (Transformées de Fourier Rapides) et des solveurs linéaires, en utilisant des communications MPI GPU-aware.
• Bibliothèques et abstractions :
  • Utilisation de bibliothèques bas niveau des fournisseurs (cuBLAS, cuSOLVER, rocBLAS, rocSOLVER, cuFFT, rocFFT) encapsulées dans une couche d'abstraction Fortran.
  • Remplacement des opérations matrice-vecteur par des opérations matrice-matrice (niveau 3 BLAS) pour maximiser l'intensité arithmétique.
• Comparaison d'algorithmes : L'étude compare deux stratégies de diagonalisation itérative :
  1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient) : Basé sur l'orthogonalisation par blocs et le gradient conjugué.
  2. Filtrage polynomial de Chebyshev : Basé sur l'application itérative de l'opérateur Hamiltonien pour filtrer le spectre, minimisant les communications et les opérations d'orthogonalisation.

3. Contributions Clés

• Portage complet d'Abinit : Mise en œuvre d'une version d'Abinit capable de fonctionner sur des nœuds CPU-GPU hétérogènes, en maintenant la fonction d'onde en mémoire GPU.
• Analyse théorique des coûts : Développement de modèles théoriques pour estimer le coût en FLOPS, les communications MPI et l'intensité arithmétique des algorithmes LOBPCG et Chebyshev.
• Optimisation des algorithmes : Démonstration que le filtrage de Chebyshev est intrinsèquement plus adapté aux GPU car il maximise l'intensité arithmétique (plus d'applications de l'Hamiltonien par communication MPI) et réduit la fréquence des goulots d'étranglement liés à la procédure de Rayleigh-Ritz.
• Nouveaux métriques : Introduction d'indicateurs de performance basés sur l'intensité arithmétique et la bande passante mémoire pour évaluer l'efficacité des noyaux GPU.

4. Résultats

Les benchmarks ont été réalisés sur des supercalculateurs français (Jean Zay avec NVIDIA A100/H100, Adastra avec AMD MI250X, Topaze).

• Accélération et Performance :
  • Les GPU NVIDIA offrent des accélérations supérieures aux GPU AMD, principalement en raison de la maturité des bibliothèques ROCm pour la procédure de Rayleigh-Ritz (qui reste un goulot d'étranglement sur AMD).
  • Sur les architectures NVIDIA, l'utilisation de 2 nœuds GPU surpasse 8 nœuds CPU, permettant des économies de ressources significatives.
  • Le filtrage de Chebyshev bénéficie davantage de l'accélération GPU que LOBPCG, car les étapes de filtrage sont compute-bound (limitées par le calcul), tandis que les étapes d'orthogonalisation de LOBPCG sont memory-bound (limitées par la mémoire).
• Efficacité Énergétique :
  • Les nœuds GPU consomment moins d'énergie par unité de travail que les nœuds CPU. Sur Jean Zay (NVIDIA), les facteurs d'économie d'énergie atteignent jusqu'à 15x par rapport aux CPU seuls pour certaines configurations.
  • L'énergie consommée par les nœuds GPU reste stable lors du passage à plusieurs nœuds, contrairement aux CPU où elle augmente linéairement.
• Comparaison des Algorithmes :
  • Chebyshev : Augmenter le degré du polynôme améliore la précision des vecteurs propres sans pénalité de temps excessive sur GPU, permettant une convergence SCF plus rapide (moins d'itérations globales).
  • LOBPCG : Augmenter le nombre de lignes de minimisation n'apporte pas de gains significatifs de vitesse et peut même être contre-productif car il augmente la fréquence des communications et des opérations d'orthogonalisation coûteuses.

5. Signification et Perspectives

• Impact Scientifique : Ce travail valide que la DFT à grande échelle peut être efficacement exécutée sur des architectures hétérogènes modernes. Il démontre que le simple portage de code n'est pas suffisant ; une réécriture algorithmique est nécessaire pour exploiter la puissance des GPU.
• Optimisation Future : L'étude conclut que les algorithmes basés sur le filtrage spectral (comme Chebyshev) sont préférables aux méthodes de gradient conjugué par blocs pour les GPU, car ils favorisent les opérations intensives en calcul.
• Évolutivité : Pour les systèmes très grands, les auteurs suggèrent l'adoption de méthodes de « Spectrum Slicing » (découpage du spectre) afin de réduire la taille des sous-espaces diagonalisés dans la procédure de Rayleigh-Ritz, évitant ainsi les communications globales coûteuses.
• Disponibilité : Ces améliorations sont intégrées dans la version 10.6 d'Abinit, rendant le code accessible à la communauté pour des simulations plus rapides et plus économes en énergie.

En résumé, cet article fournit une feuille de route technique pour le portage efficace de codes de chimie quantique sur GPU, en soulignant l'importance cruciale de l'adaptation des algorithmes de diagonalisation pour tirer parti de l'intensité arithmétique élevée des accélérateurs modernes.