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

Dit artikel beschrijft de GPU-versnelling van de Abinit-code voor plane-wave DFT-berekeningen, waarbij de prestaties van CPU- versus GPU-nodes worden vergeleken en twee diagonalisatie-algoritmen op hun GPU-efficiëntie worden geëvalueerd.

De kern

Stel je voor dat Abinit een gigantische, superkrachtige keuken is waar chefs (wetenschappers) proberen te voorspellen hoe nieuwe materialen zich gedragen. Om dit te doen, moeten ze een ingewikkeld recept volgen: ze moeten de "elektronische golf" van elk atoom berekenen. Dit is als proberen de exacte beweging van miljarden dansende mieren tegelijk te voorspellen.

Vroeger deden ze dit met gewone computers (CPU's), maar dat duurde eeuwen. In dit artikel vertellen de auteurs hoe ze deze keuken hebben omgebouwd om GPU's (de krachtige videokaarten die we kennen van gaming) te gebruiken. Het resultaat? De berekeningen gaan veel sneller en zuiniger.

Hier is hoe ze dat hebben gedaan, vertaald in alledaags taal:

1. Het probleem: Te veel werk voor één kok

De berekeningen zijn zo zwaar dat je duizenden "elektronische banen" (de dansende mieren) tegelijk moet bekijken.

• De oude manier (CPU): De chef pakt één mier, berekent zijn beweging, legt hem neer, pakt de volgende. Dit is traag.
• De nieuwe manier (GPU): Een GPU is als een leger van duizenden kleine koks die allemaal tegelijk werken. Maar om dit te laten werken, moet je de opdracht anders geven. Je kunt niet zeggen "doe dit voor mier 1, dan mier 2". Je moet zeggen: "Neem een bak vol mieren en doe dit allemaal tegelijk!"

2. De strategie: "Batch Processing" (De bak met mieren)

De auteurs hebben een slimme truc bedacht: Batch Processing.
Stel je voor dat je in plaats van één voor één brieven te posten, ze allemaal in één grote zak stopt en de postbode (de GPU) laat zien: "Hier, bezorg deze hele zak!"

• In Abinit hebben ze de data zo gerangschikt dat de GPU niet één elektronische golf per keer hoeft te berekenen, maar een hele groep (een "batch") tegelijk.
• Dit voorkomt dat de GPU tijd verliest met wachten op nieuwe instructies. Het is als het verschil tussen een trein die stopt bij elk station (traag) en een hogesnelheidstrein die alleen stopt bij grote knooppunten (snel).

3. De twee methoden: Twee manieren om de dans te leren

Om de golfbewegingen te vinden, gebruiken ze twee verschillende algoritmen (recepten). De auteurs hebben gekeken welke het beste werkt op de GPU:

• LOBPCG (De "Oefen-met-elkaar" methode):
  Deze methode is als een groep dansers die steeds weer naar elkaar toe lopen, elkaar corrigeren, en dan weer een stapje verder gaan. Ze moeten constant communiceren en controleren of ze nog in de pas lopen.

  • Nadeel: Omdat ze constant moeten "praten" (communiceren) en wachten op elkaar, is dit traag op een GPU. Het is te veel gedoe voor de snelle koks.

• Chebyshev Filtering (De "Herhaling" methode):
  Deze methode is als een danser die een moeilijke beweging duizend keer herhaalt om hem perfect te krijgen, zonder constant naar de anderen te kijken.

  • Voordeel: De GPU is fantastisch in het herhalen van dezelfde simpele taak heel vaak. Deze methode laat de GPU zijn kracht volledig gebruiken.
  • Resultaat: Chebyshev Filtering is de winnaar. Het is sneller, zuiniger en schaalbaarder.

4. Het geheugenprobleem: De koelkast en het aanrecht

Een groot probleem bij GPU's is het verplaatsen van data.

• CPU is als de hoofdkok (het aanrecht).
• GPU is als een super-snelle snijmachine (het snijbord).
• Als je de groenten (data) steeds heen en weer moet slepen tussen het aanrecht en de snijmachine, kost dat tijd.
• De oplossing: De auteurs hebben ervoor gezorgd dat de "groenten" (de elektronische golven) direct in de snijmachine worden bewaard. Ze slepen ze maar één keer naar binnen, laten de snijmachine al het zware werk doen, en halen ze pas weer terug als het echt nodig is. Dit bespaart enorm veel tijd.

5. De uitkomst: Snelheid en zuinigheid

Wat hebben ze gevonden?

• Snelheid: Op de nieuwe systemen (met NVIDIA en AMD GPU's) gaan de berekeningen tot wel 17 keer sneller dan met alleen gewone computers.
• Energie: Omdat de GPU's zo efficiënt zijn, verbruiken ze minder stroom voor dezelfde hoeveelheid werk. Het is alsof je met een elektrische auto dezelfde afstand rijdt als met een benzineauto, maar dan met een kwart van de brandstof.
• Schalen: Als je meer computers toevoegt, werkt de Chebyshev-methode veel beter dan de oude methode. De oude methode raakt in de war als je te veel koks toevoegt (ze praten te veel met elkaar), maar de nieuwe methode blijft soepel draaien.

Conclusie

Dit artikel laat zien dat je niet zomaar een oude software kunt "omtoveren" naar een GPU. Je moet het recept (de algoritmen) zelf herschrijven. Door te kiezen voor de juiste methode (Chebyshev Filtering) en slim met data om te gaan (alles in het geheugen van de GPU houden), kunnen wetenschappers nu materialen veel sneller ontwerpen. Dit helpt bij het ontwikkelen van betere batterijen, zonnepanelen en nieuwe medicijnen.

Kortom: Ze hebben de keuken van een traag restaurant omgebouwd tot een supersnelle, geautomatiseerde fabriek.

Technische samenvatting

Titel: GPU-versnelling van vlakke-golf DFT-berekeningen in Abinit

Auteurs: Ioanna-Maria Lygatsika et al. (CEA, Université Paris-Saclay, Alliance Service Plus)

---

1. Het Probleem

Grootschalige simulaties in de materiaalkunde op kwantumniveau, zoals die uitgevoerd worden met Density Functional Theory (DFT) op basis van vlakke golven (plane-wave), vereisen het oplossen van de Kohn-Sham-vergelijkingen voor duizenden elektronische toestanden. Dit leidt tot enorme rekenkundige lasten, voornamelijk door:

• Het diagonaliseren van grote Hamilton-matrices.
• Het uitvoeren van talloze lineaire algebra-operaties en Fast Fourier Transforms (FFTs).
• De beperkingen van traditionele CPU-only High-Performance Computing (HPC) systemen om deze problemen binnen een redelijke tijd op te lossen.

Hoewel GPU's (Graphics Processing Units) al geruime tijd beschikbaar zijn als acceleratoren, is het porten van complexe, iteratieve DFT-codes zoals Abinit naar GPU-architecturen uitdagend. Het gaat niet alleen om het gebruik van leveranciersbibliotheken, maar vereist fundamentele herzieningen van de algoritmen (zoals iteratieve diagonalisatie) om ze efficiënt te laten werken op GPU's, waarbij rekening moet worden gehouden met geheugenbandbreedte, communicatie-overhead en de aard van de wiskundige operaties.

2. Methodologie en Porting-Strategie

De auteurs hebben de Abinit-code volledig opnieuw geïmplementeerd voor GPU's, gebruikmakend van een hybride MPI + OpenMP offloading-model. De kern van de strategie omvat de volgende elementen:

• Batch-verwerking (Batch Processing): In plaats van operaties per elektronische band (rij) uit te voeren, worden data gegroepeerd in batches. Dit maximaliseert de parallelle verwerking en verhoogt de doorvoer. De golffunctie ($\Psi$) wordt behandeld als een batch van kolomvectoren, waardoor FFT's en matrix-matrix vermenigvuldigingen efficiënter kunnen worden uitgevoerd via één kernel-aanroep in plaats van vele kleine.
• GPU-resident Golffunctie: Om de kostbare overdracht van data tussen host (CPU) en device (GPU) te minimaliseren, wordt de golffunctie tijdens de iteratieve diagonalisatie (binnen een SCF-stap) volledig in het GPU-geheugen bewaard. Data-overdracht vindt alleen plaats aan het begin en einde van een Self-Consistent Field (SCF)-iteratie.
• Communicatie en MPI: De code gebruikt "GPU-aware MPI" om directe communicatie tussen GPU's mogelijk te maken. Er worden twee distributiemodellen gebruikt:
  • Rij-gedistribueerd: Voor lineaire algebra-oplossers.
  • Kolom-gedistribueerd (Waveform layout): Voor het toepassen van de Hamilton-operator en FFT's.
    Het schakelen tussen deze modellen vereist een "all-to-all" communicatie, die zo efficiënt mogelijk wordt uitgevoerd.
• Abstractie-laag: Een softwarelaag is ontwikkeld die de complexe interactie tussen de Abinit-datastructuren en de lage-niveau GPU-bibliotheken (zoals cuBLAS, rocBLAS, cuFFT, rocFFT, cuSOLVER, rocSOLVER) verbergt. Dit maakt de code portabel voor zowel NVIDIA als AMD GPU's.

3. Vergelijking van Diagonalisatie-algoritmen

Een centraal onderdeel van het onderzoek is de vergelijking van twee iteratieve diagonalisatie-algoritmen in termen van GPU-efficiëntie:

1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient): Een vector-georiënteerde methode die blokken van vectoren gebruikt. Deze methode vereist frequente orthogonisatie tussen blokken, wat leidt tot veel MPI-communicatie en een hogere frequentie van het Rayleigh-Ritz-procedure (een bottleneck).
2. Chebyshev-polynoomfiltering: Een spectrum-georiënteerde methode die de eigenwaarde-spectrum filtert met behulp van polynomen. Deze methode elimineert inter-blok afhankelijkheden en voert de Hamilton-operator veelvuldig toe met slechts twee globale communicatiestappen (aan het begin en einde van de filtering).

Theoretische analyse: De auteurs tonen aan dat Chebyshev-filtering een veel hogere arithmetic intensity (rekenintensiteit) heeft per MPI-proces dan LOBPCG, omdat het de Hamilton-operator $k$ keer toepast per data-transpositie, terwijl LOBPCG dit slechts één keer doet per iteratie.

4. Resultaten en Prestaties

De prestaties zijn getest op verschillende supercomputers (Jean Zay, Adastra, Topaze) met NVIDIA (A100, H100) en AMD (MI250X, MI300) GPU's.

• Snelheidswinst (Speedup):
  • NVIDIA GPU's tonen consistent hogere snelheidswinsten dan AMD GPU's.
  • Voor een systeem met 255 atomen en 4096 banden: 2 GPU-nodes presteren beter dan 8 CPU-nodes.
  • De Chebyshev-filterstap wordt enorm versneld (compute-bound), terwijl de Rayleigh-Ritz-stap (diagonalisatie) minder goed schaalt en een bottleneck blijft, vooral op AMD-architecturen waar de ROCm-bibliotheken minder goed presteren voor deze specifieke operaties.
• Energie-efficiëntie:
  • GPU's bieden aanzienlijke energiebesparingen. Op Jean Zay (NVIDIA) is de energiebesparing factor 11 tot 15 hoger dan op CPU-only nodes.
  • De energieconsumptie op AMD-systemen (Adastra) schaalt bijna lineair met het aantal nodes, terwijl deze op NVIDIA-systemen vrijwel constant blijft bij het toevoegen van meer nodes, wat wijst op een betere schaalbaarheid.
• Algoritme-vergelijking:
  • Het verhogen van de polynoomgraad bij Chebyshev-filtering verbetert de nauwkeurigheid van de eigenvectoren zonder de totale uitvoeringstijd te verslechteren op GPU's.
  • Bij LOBPCG leidt het verhogen van het aantal minimalisatiestappen (minimization lines) tot meer rekenwerk zonder evenredige winst in SCF-convergentie.
  • Conclusie: Chebyshev-filtering is beter geschikt voor GPU's omdat het de rekenintensieve operaties (Hamilton-toepassing) maximaliseert en de communicatie-intensieve operaties (Rayleigh-Ritz) minimaliseert.

5. Belangrijkheid en Conclusie

Dit werk markeert een mijlpaal in de modernisering van Abinit voor moderne HPC-architecturen. De belangrijkste bijdragen zijn:

• Validatie van een nieuwe porting-strategie: Het succesvol implementeren van een "GPU-resident" golffunctie-aanpak met batch-verwerking.
• Algoritmische optimalisatie: Het aantonen dat voor GPU's algoritmen die rekenintensief zijn (zoals Chebyshev-filtering) superieur zijn aan die die veel communicatie en geheugenbandbreedte vereisen (zoals LOBPCG).
• Hardware-onafhankelijkheid: De code werkt op zowel NVIDIA als AMD hardware, hoewel de prestaties op NVIDIA (Ampere/Hopper) momenteel superieur zijn.
• Toekomstperspectief: De auteurs suggereren dat voor nog grotere systemen methoden zoals "Spectrum Slicing" nodig zijn om de bottleneck van de Rayleigh-Ritz-diagonalisatie volledig te doorbreken.

Samenvattend toont dit artikel aan dat een zorgvuldige herontwerp van algoritmen, gecombineerd met het gebruik van leveranciersbibliotheken en batch-verwerking, essentieel is om de rekenkracht van GPU's volledig te benutten voor complexe kwantumsimulaties.