Accelerating finite-element-based projector augmented-wave density functional theory calculations with scalable GPU-centric computational methods

Ce papier présente une méthode d'ondes augmentées par projecteur à éléments finis (PAW-FE) centrée sur le GPU et évolutive, qui exploite des innovations algorithmiques telles que l'arithmétique de précision mixte et l'itération de sous-espace filtrée par Chebyshev pour obtenir des accélérations significatives et des performances prêtes pour l'exascale dans le cadre de simulations de grande envergure de la théorie de la fonctionnelle de la densité chimiquement précises.

L'essentiel

Imaginez que vous essayez de prédire le comportement d'une machine complexe, comme un moteur de voiture ou un nouveau type de batterie. Pour le faire avec précision, vous devez comprendre le comportement de chaque électron individuel à l'intérieur des matériaux qui composent la machine. C'est la tâche d'un domaine appelé Théorie de la Fonctionnelle de la Densité (DFT). C'est comme essayer de simuler une immense et complexe piste de danse où des milliards d'électrons se déplacent en synchronisation.

Pendant longtemps, les scientifiques ont eu un problème : simuler ces danses pour de petits groupes d'atomes est facile, mais dès que vous essayez de simuler un système grand et complexe (comme une nanoparticule métallique minuscule ou une feuille de matériau tordue), l'ordinateur est submergé. C'est comme essayer de diriger une danse pour 100 000 personnes en utilisant une méthode conçue pour 100 ; les instructions s'emmêlent, la mémoire se remplit, et la simulation prend une éternité pour se terminer.

Cet article présente une nouvelle méthode ultra-rapide pour exécuter ces simulations, spécifiquement conçue pour les ordinateurs modernes et puissants utilisant des GPU (les mêmes puces qui alimentent les jeux vidéo haut de gamme et l'IA). Voici comment ils ont procédé, décomposé en concepts simples :

1. L'Ancienne Méthode vs La Nouvelle Carte

• L'Ancienne Méthode (Ondes planes) : Imaginez essayer de cartographier une ville en utilisant une grille gigantesque et uniforme où chaque pouce carré a la même taille. Si vous voulez voir un détail minuscule (comme une brique unique sur un bâtiment), vous devez rendre l'ensemble de la grille incroyablement fine, même pour le ciel vide au-dessus de la ville. Cela gaspille une quantité massive de puissance informatique. C'est ainsi que fonctionne la plupart des logiciels actuels.
• La Nouvelle Méthode (Éléments finis) : Les auteurs utilisent une approche de « carte intelligente ». Imaginez une carte qui zoome uniquement là où c'est nécessaire (comme le centre-ville animé) et dézoome là où c'est vide (comme le ciel). C'est ce qu'on appelle la discrétisation par Éléments Finis (EF). Cela leur permet de concentrer leur puissance de calcul exactement là où les électrons font des choses intéressantes, économisant d'énormes quantités de temps et de mémoire.

2. L'Astuce « PAW » : Le Costume Magique

Pour rendre les mathématiques encore plus faciles, ils utilisent une méthode appelée Onde Augmentée par Projecteurs (PAW).

• Le Problème : Les électrons près du centre d'un atome (le noyau) ondulent et vibrent sauvagement, ce qui les rend difficiles à calculer.
• La Solution : La PAW est comme mettre un « costume lisse » sur les électrons. Elle fait semblant que les électrons sont lisses et faciles à manipuler pour la majeure partie du calcul, mais elle conserve un « tour de magie » secret qui lui permet de révéler instantanément le véritable comportement sauvage des électrons juste au moment où elle doit vérifier les détails près du noyau. Cela leur permet d'utiliser une carte beaucoup plus grossière (plus simple) sans perdre en précision.

3. L'Accélération GPU : La Chaîne de Montage

Les auteurs n'ont pas seulement changé la carte ; ils ont modifié la façon dont l'ordinateur fait les mathématiques pour s'adapter aux GPU modernes.

• Le Goulot d'Étranglement : Habituellement, les ordinateurs passent beaucoup de temps à attendre que les données se déplacent entre la mémoire et le processeur.
• La Correction : Ils ont redessiné les mathématiques afin que l'ordinateur puisse effectuer de nombreux calculs à la fois (comme une chaîne de montage) plutôt que les uns après les autres. Ils ont également utilisé une technique ingénieuse appelée Filtrage de Tchebychev, qui est comme un tamis qui sépare rapidement les électrons « importants » des électrons « non importants », afin que l'ordinateur ne perde pas de temps avec ceux dont il n'a pas besoin.

4. Les Raccourcis « Assez Bien » (Précision Mixte)

C'est peut-être la partie la plus créative.

• L'Analogie : Imaginez que vous peignez une immense fresque murale. Pour le ciel en arrière-plan, vous n'avez pas besoin de mélanger la peinture avec une précision microscopique ; un mélange « assez bon » fonctionne bien et est beaucoup plus rapide. Vous n'avez besoin d'une précision extrême que pour les détails minuscules d'un visage.
• L'Application : Les auteurs ont réalisé que pour les parties du calcul qui ont juste besoin de bien définir la forme générale, ils peuvent utiliser des mathématiques de précision inférieure (comme utiliser une règle avec moins de graduations). C'est beaucoup plus rapide sur les puces modernes. Ils ne passent aux mathématiques « ultra-précises » que pour les étapes finales et critiques.
• Le Résultat : En mélangeant des mathématiques de haute précision et de faible précision, et en superposant les transferts de données aux calculs (faire deux choses à la fois), ils ont fait fonctionner la simulation 8 à 20 fois plus vite qu'auparavant.

5. Ce Qu'ils Ont Réellement Accompli

L'article affirme qu'avec ces nouvelles méthodes :

• Vitesse : Ils peuvent maintenant simuler des systèmes contenant 10 000 à 130 000 électrons dans un temps pratique (de quelques minutes à quelques heures) sur des superordinateurs.
• Comparaison : Leur méthode est environ 8 fois plus rapide que le logiciel standard de pointe (Quantum ESPRESSO) pour des systèmes de cette taille.
• Échelle : Ils ont exécuté avec succès une simulation d'un matériau « bicouche torsadée » (deux feuilles d'atomes torsadées ensemble) contenant 130 000 électrons. C'est une taille qui était auparavant impossible à simuler avec ce niveau de précision en utilisant des méthodes standard.

Résumé

En bref, les auteurs ont construit un nouveau moteur hautement efficace pour simuler des matériaux. Ils ont combiné une « carte intelligente » qui zoome uniquement là où c'est nécessaire, une astuce de « costume magique » pour simplifier les mathématiques, et un mode « avance rapide » qui utilise une précision inférieure pour les étapes non critiques. Le résultat est un outil capable de modéliser des matériaux massifs et complexes sur des superordinateurs modernes en une fraction du temps qu'il fallait auparavant, ouvrant la voie à la conception de nouveaux matériaux pour les batteries, l'électronique et les catalyseurs beaucoup plus rapidement.

Résumé technique

1. Énoncé du problème

La modélisation précise de systèmes matériels complexes (par exemple, interfaces, défauts, nanoclusters, hétérostructures torsadées) nécessite des simulations de théorie de la fonctionnelle de la densité (DFT) impliquant de $10^4$ à $10^5$ électrons. Cependant, les implémentations existantes rencontrent des goulots d'étranglement significatifs :

• Limites des ondes planes (PW) : Les codes PW-PAW largement utilisés (par exemple, VASP, Quantum ESPRESSO) reposent sur des transformées de Fourier rapides (FFT), qui entraînent une communication massive de type « tout-à-tout ». Cela les rend limités par la bande passante et peu évolutifs sur les architectures GPU exascale modernes.
• Limites de l'espace réel : Les méthodes d'espace réel existantes utilisent souvent des pseudopotentiels à conservation de norme (ONCV), qui nécessitent un grand nombre de fonctions de base (degrés de liberté élevés) pour atteindre une précision chimique, en particulier pour les métaux de transition, entraînant une empreinte mémoire et des coûts de calcul élevés.
• Inadéquation matérielle : Les solveurs de valeurs propres actuels sont souvent mal adaptés aux fortes intensités arithmétiques et aux exigences de faible latence des clusters GPU modernes, ne parvenant pas à exploiter les capacités de précision mixte ou à superposer efficacement la communication et le calcul.

2. Méthodologie

Les auteurs présentent PAW-FE, une formulation d'ondes projectrices augmentées (PAW) discrétisée par éléments finis (FE) optimisée pour les architectures GPU multi-nœuds.

A. Formulation mathématique

• PAW-GHEP : La méthode résout le problème généralisé de valeurs propres hermitiennes (GHEP) : $H\tilde{\Psi} = S\tilde{\Psi}\Lambda$, où $H$ est le hamiltonien et $S$ la matrice de recouvrement PAW.
• Formalisme de spin collinéaire : Les équations sont dérivées dans un cadre de spin collinéaire pour traiter les systèmes magnétiques.
• Conditions aux limites : La discrétisation FE intègre naturellement les conditions aux limites périodiques, semi-périodiques (surfaces) et non périodiques (nanoclusters) sans artefacts de périodicité artificielle.

B. Innovations computationnelles

Pour résoudre efficacement le PAW-GHEP sur GPU, les auteurs ont développé plusieurs stratégies algorithmiques clés :

1. Itération de sous-espace filtrée par Chebyshev basée sur le résidu (R-ChFSI) :

   • Au lieu du ChFSI standard, ils utilisent une formulation basée sur le résidu ($R = H\tilde{\Psi} - S\tilde{\Psi}\Lambda$).
   • Avantage clé : Cette formulation tolère des produits matrice-vecteur inexacts, permettant l'utilisation d'inverses approximatifs pour la matrice de recouvrement PAW ($S^{-1}$) et d'arithmétique de précision réduite sans sacrifier la convergence.

2. Matrice de recouvrement inverse approximative :

   • Au lieu d'inverser explicitement la matrice creuse $S$, ils utilisent une approximation peu coûteuse ($\tilde{S}^{-1}$) dérivée d'une approximation diagonale de la matrice de masse et de corrections localisées. Cela évite les inversions globales coûteuses.

3. Quadrature multi-résolution :

   • Pour traiter les intégrales centrées sur les atomes (densités de pseudo-cœur) sur des grilles FE grossières, un schéma de quadrature multi-résolution est employé. Il utilise des règles de quadrature affinées uniquement à l'intérieur des sphères d'augmentation des atomes tout en utilisant des règles plus grossières ailleurs, garantissant la précision sans raffinement de maillage.

4. Précision mixte et communication basse précision :

   • Calcul : Utilise l'arithmétique FP32/TF32 pour les étapes de filtrage Chebyshev (multiplications matrice-matrice) et BF16 pour la communication entre voisins.
   • Robustesse : La dépendance de l'algorithme R-ChFSI aux résidus garantit que ces réductions de précision ne dégradent pas la précision finale en double précision de l'état fondamental.

5. Superposition calcul-communication :

   • Le filtrage Chebyshev est effectué par blocs. Tandis qu'un bloc de fonctions d'onde subit une communication inter-processeurs (MPI), un autre bloc effectue un calcul (opérations GEMM) sur le GPU. Cela masque la latence de communication, un facteur critique pour l'évolutivité exascale.

6. Opérations denses au niveau de la cellule :

   • Plutôt que de construire des matrices globales creuses, la méthode reformule les opérations comme des multiplications matrice-matrice denses au niveau de la cellule. Cela maximise l'intensité arithmétique et la localité du cache, ce qui est idéal pour l'exécution GPU.

3. Contributions clés

• Formulation PAW-FE : Une nouvelle formulation PAW discrétisée par éléments finis dans un formalisme de spin collinéaire supportant des conditions aux limites génériques.
• Solveur de valeurs propres R-ChFSI : Une extension de l'itération de sous-espace filtrée par Chebyshev basée sur le résidu pour résoudre le PAW-GHEP, permettant l'utilisation d'inverses approximatifs et de précision mixte.
• Implémentation GPU évolutive : Une stratégie d'implémentation complète intégrant l'arithmétique de précision mixte (FP32/TF32/BF16), la superposition calcul-communication et l'algèbre linéaire dense au niveau de la cellule.
• Intégration multi-résolution : Un schéma de quadrature permettant une évaluation précise des intégrales PAW centrées sur les atomes sur des maillages grossiers.

4. Résultats et performances

La méthode a été évaluée sur des supercalculateurs de classe de pointe (OLCF Frontier, ALCF Aurora, ALCF Polaris) utilisant des GPU AMD, Intel et NVIDIA.

• Précision : Validée contre des codes d'ondes planes (Abinit, Quantum ESPRESSO) pour des molécules ($O_2$, $NO_2$) et des cristaux (Cr CC). Les erreurs en énergie et en forces sont d'un ordre de grandeur inférieur aux exigences de précision chimique ($O(10^{-12})$ Ha/atome pour l'énergie, $O(10^{-6})$ Ha/bohr pour les forces).
• Accélération CPU-GPU : Un gain de vitesse de 8 à 20 fois a été réalisé sur les GPU par rapport aux exécutions uniquement CPU (architectures Intel et AMD).
• Comparaison avec les ondes planes (QE) :
  • Pour des systèmes d'environ 10 000 électrons, PAW-FE réalise une réduction de 8 fois du temps d'exécution minimal par rapport à Quantum ESPRESSO.
  • Pour des systèmes plus grands (>10 000 électrons), l'accélération augmente davantage en raison de la localité des fonctions de base FE par rapport à la communication globale des méthodes PW.
• Comparaison avec ONCV-FE :
  • PAW-FE nécessite environ 6 fois moins de ressources de calcul et atteint un temps de résolution environ 5 fois plus rapide que les méthodes FE utilisant des pseudopotentiels à conservation de norme (ONCV), principalement grâce aux degrés de liberté réduits permis par la méthode PAW.
• Évolutivité exascale :
  • Une démonstration réussie d'un calcul d'état fondamental pour un système de bicouche torsadée WTe2 comprenant 130 000 électrons (11 000 atomes).
  • Un temps de résolution de ~2 minutes par itération SCF a été atteint sur 400 nœuds, prouvant la viabilité de la méthode pour des simulations réalistes à grande échelle.

5. Importance

Ce travail établit PAW-FE comme une méthode prête pour l'exascale pour les simulations de premiers principes. En comblant le fossé entre la haute précision de la méthode PAW et l'efficacité des éléments finis d'espace réel, elle surmonte les goulots d'étranglement de communication des méthodes d'ondes planes. Le déploiement réussi de stratégies de précision mixte et de superposition sur diverses architectures GPU démontre une voie vers la simulation régulière de systèmes matériels complexes avec $10^5$ électrons, permettant de nouvelles découvertes en catalyse, interfaces de batteries et matériaux quantiques qui étaient auparavant hors de portée computationnelle.