Accelerating High-Order Finite Element Simulations at Extreme Scale with FP64 Tensor Cores

Deze paper demonstreert dat het gebruik van FP64-tensorcores op NVIDIA-GPU's, gecombineerd met kernelfusie-optimalisaties, de prestaties en energie-efficiëntie van high-order eindige-elementensimulaties aanzienlijk verbetert, met een schaalbaarheid tot bijna 10.000 GPU's op het exascale-systeem Alps.

De kern

Stel je voor dat je een gigantische, complexe puzzel moet oplossen. Deze puzzel is eigenlijk een computermodel van een tsunami die over de oceaan raast. Om te voorspellen waar het water naartoe gaat en hoe hoog de golven worden, moeten wetenschappers miljarden kleine rekensommen doen. Dit heet een "finite element simulatie".

Vroeger was dit als het proberen om die puzzel op te lossen met een potlood en papier: het kan, maar het duurt eeuwen.

Het probleem: De "Bureaublad" is te klein
In de wereld van supercomputers (zoals de enorme machines in Zwitserland waar dit onderzoek plaatsvond), werken de processors als een team van duizenden werknemers. Elke werknemer heeft een klein bureau (geheugen) om tijdelijke notities op te zetten.

De oude manier van werken was inefficiënt: elke werknemer moest zijn eigen kopie van de gegevens van een grote stapel papier halen, ze op zijn bureau leggen, een som maken, en de resultaten weer wegleggen. Omdat iedereen zijn eigen kopieën maakte, was hun bureau snel vol, en ze spendeerden 90% van hun tijd aan het heen en weer lopen van de stapel naar het bureau, in plaats van aan het rekenen zelf.

De oplossing: De "Super-Bureaus" (Tensor Cores)
NVIDIA heeft een nieuw type computerchip gemaakt met speciale "Tensor Cores". Je kunt deze zien als een team van 32 werknemers die een gemeenschappelijk, supersnel bureau delen.

In plaats dat elke werknemer zijn eigen kopie van de gegevens haalt, halen ze samen één keer de gegevens, en dan werken ze allemaal tegelijkertijd aan die ene set. Dit is wat de auteurs in dit paper hebben gedaan: ze hebben de software herschreven zodat de rekenkracht van deze "Super-Bureaus" (de Tensor Cores) wordt gebruikt, zelfs voor de zware, dubbel-precisie rekeningen die nodig zijn voor wetenschappelijke nauwkeurigheid.

De analogie van de "Geknipte Puzzelstukken"
De auteurs ontdekten dat de rekenopdrachten niet perfect pasten bij de standaard formaat van deze Super-Bureaus. Het was alsof je een puzzelstukje van 25x5 probeerde te passen in een vakje dat gemaakt is voor 8x8. Er bleef veel ruimte over die niet werd gebruikt.

Om dit op te lossen, hebben ze twee slimme trucs toegepast:

1. De "Geknipte" aanpak: Ze hebben de grote puzzelstukjes op een slimme manier opgedeeld en herschikt, zodat ze perfect in de vakjes van de Super-Bureaus passen. Hierdoor verdwijnen de "dode hoeken" en wordt de ruimte optimaal benut.
2. De "Samengevoegde" aanpak (Kernel Fusion): Vroeger moesten de werknemers na elke kleine som even pauzeren, hun resultaten opschrijven, en dan pas de volgende stap doen. De auteurs hebben de software zo aangepast dat ze meerdere stappen in één keer doen zonder te stoppen. Het is alsof je niet meer tussendoor je notitieboekje dichtdoet, maar gewoon doorgaat met schrijven tot de hele zin af is. Dit bespaart enorm veel tijd.

Het resultaat: Een tsunami in een seconde
Door deze verbeteringen is het resultaat verbluffend:

• Snelheid: De simulaties zijn tot 2 keer sneller geworden. Wat vroeger uren duurde, gaat nu in een flits.
• Energie: Omdat de computers minder tijd nodig hebben om te "lopen" en meer tijd besteden aan het "rekenen", verbruiken ze tot 83% minder energie. Het is alsof je dezelfde afstand rijdt, maar met een auto die veel zuiniger is.
• Schaalbaarheid: Ze hebben dit getest op een van de grootste computers ter wereld (de "Alps" in Zwitserland) met bijna 10.000 van deze krachtige chips. Het systeem werkte perfect samen, alsof het één enorme, super-efficiënte machine was.

Waarom is dit belangrijk?
Dit is niet alleen een theoretisch experiment. Dit wordt gebruikt voor een digitale tweeling van een tsunami. Als er een aardbeving is, kan dit systeem binnen één seconde voorspellen hoe hoog de tsunami zal worden en welke kustlijnen in gevaar zijn. Dit geeft mensen kostbare tijd om te evacueren.

Kortom: De auteurs hebben de "motor" van de computercode getuned, zodat hij niet meer vastloopt in de "verkeersdrukte" van het geheugen, maar soepel en razendsnel over de weg rijdt. Hierdoor kunnen we levens redden door snellere en zuinigere voorspellingen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Accelerating High-Order Finite Element Simulations at Extreme Scale with FP64 Tensor Cores", geschreven in het Nederlands.

Probleemstelling

Finite-elementen (FE) simulaties zijn cruciaal voor toepassingen variërend van auto-ontwerp tot tsunami-modellering en computationele elektromagnetica. Voor praktische toepassingen en wetenschappelijke inzichten zijn hoge resoluties vereist, wat leidt tot de noodzaak van hoog-orde methoden en supercomputing op extreme schaal.

Hoewel er al vooruitgang is geboekt bij het porten van FE-codes naar GPU-systemen, blijven er uitdagingen bestaan:

1. Precisie-eisen: Veel wetenschappelijke toepassingen, zoals inversieproblemen voor golfvoortplanting (bijv. tsunami-voorspelling), vereisen volledige dubbele precisie (FP64) om numerieke instabiliteit en fouten te voorkomen. Traditionele Tensor Cores op NVIDIA GPU's waren oorspronkelijk gericht op gemengde precisie (FP16/FP32).
2. Bottlenecks: Bestaande GPU-kernels voor hoog-orde FE-berekeningen worden vaak beperkt door de bandbreedte van het gedeeld geheugen (shared memory) en niet door de rekenkracht zelf. De verhouding tussen uitgevoerde floating-point operaties (FLOPs) en gelezen bytes is vaak te laag.
3. Efficiëntie: Er is een constante vraag naar verbeteringen in zowel rekensnelheid als energie-efficiëntie voor deze complexe simulaties op exascale-systemen.

Methodologie

De auteurs hebben hun onderzoek gefocust op de MFEM-bibliotheek (een open-source, hoog-presterende FE-bibliotheek) en een specifieke toepassing: een digitale tweeling voor tsunami-early warning (een inversieprobleem voor akoestisch-zwaartekrachtgolfvoortplanting).

De kern van de methodologie bestaat uit drie hoofdfasen:

1. Directe Programmering van FP64 Tensor Cores (DMMA):

   • In plaats van te vertrouwen op standaard CUDA-kernen of bibliotheekfuncties (zoals cuBLAS), hebben de auteurs de Double Precision Matrix-Multiply-Accumulate (DMMA) instructies direct geprogrammeerd in CUDA-kernen.
   • Ze hebben de tensorcontracties in de FE-operatoren ontbonden in kleine, dichte matrixvermenigvuldigingen (GEMMs) van de orde $O(10)$ (bijv. $25 \times 5 \times 4$).
   • Door gebruik te maken van de m8n8k4 DMMA-instructie op NVIDIA Ampere, Hopper en Blackwell-architecturen, wordt de data uit het gedeeld geheugen gedeeld door threads binnen een warp, wat de totale hoeveelheid gelezen data drastisch verlaagt.

2. Optimalisatie van Gedeeld Geheugen (Bank Conflicts):

   • Een groot probleem bij het gebruik van Tensor Cores is het vermijden van "shared memory bank conflicts", wat de prestaties kan vertragen.
   • De auteurs hebben optimale mappings ontwikkeld voor de rij- en kolomindices van de matrices naar de "lane indices" van de warp. Hierdoor wordt voorkomen dat meerdere threads tegelijkertijd naar dezelfde geheugenbank schrijven of lezen.
   • Ze hebben ook de volgorde van tensorindices herschikt (cyclic ordering) om te zorgen dat de index waarover wordt opgeteld, de snelst veranderende index is, wat conflicten verder minimaliseert.

3. Kernel Fusie (Loop Fusion):

   • De auteurs hebben verschillende operatoren in de FE-assemblage (zoals $G$, $B$, $D$) samengevoegd tot één geoptimaliseerde kernel ("Fused PA" en "Fused MF").
   • Dit elimineert de noodzaak om tussenstappen in het geheugen op te slaan, vermindert de databeweging en maximaliseert de rekenintensiteit.
   • Specifiek voor de "Matrix-Free" (MF) aanpak worden basisfuncties en coördinaten direct berekend in plaats van opgeslagen, wat het geheugengebruik verder verlaagt.

Belangrijkste Bijdragen

• Eerste toepassing: Dit is naar weten de eerste keer dat FP64 Tensor Cores direct geprogrammeerd zijn om grote-schaal wetenschappelijke toepassingen (PDE-gebaseerd) te versnellen, in plaats van alleen voor grote lineaire algebra-bibliotheken.
• Ontwerp voor onregelmatige matrices: Een gedetailleerd ontwerp en analyse van het programmeren van FP64 Tensor Cores voor onregelmatig gevormde matrixvermenigvuldigingen (kleine maten, niet-perfecte veelvouden van 8).
• Energie-efficiëntie analyse: Een unieke vergelijking van de energie-efficiëntie tussen de GH200 (Grace Hopper) en GB200 (Grace Blackwell) chips voor deze specifieke kleine matrix-operaties.
• Exascale-schaalbaarheid: Demonstratie van de prestaties op het Alps-systeem (Swiss National Supercomputing Centre) met bijna 10.000 GPU's.

Resultaten

De experimenten werden uitgevoerd op de NVIDIA GH200 en GB200 superchips, en geschaald over het Alps-systeem (tot 9.216 GPU's).

• Prestatieverbetering (Single GPU):
  • De DMMA-optimalisaties alleen leidden tot een snelheidswinst van 35% tot 59% voor de kernkernen ten opzichte van de originele CUDA-kernen.
  • Door combinatie met kernel fusie ("DMMA Fused PA") werd een totale snelheidswinst van 2x bereikt ten opzichte van de originele PA-kernel.
• Energie-efficiëntie:
  • Het gebruik van FP64 Tensor Cores verbeterde de prestaties per Watt met 27% op de GH200 en 18% op de GB200.
  • Met kernel fusie bedroeg de verbetering in energie-efficiëntie tot 83% op de GH200.
• Schaalbaarheid:
  • Weak Scaling: Ideaal (lineair) schaalgedrag (bijna 100% efficiëntie) over een 64x toename in het aantal knooppunten (van 36 tot 2.304 nodes).
  • Strong Scaling: Uitstekende strong scaling efficiëntie van 86% tot 91% over een 64x toename in het aantal knooppunten.
• Toepassing: De geoptimaliseerde kernels zijn direct toegepast in de 2025 Gordon Bell Prize-winnende toepassing voor real-time tsunami-voorspelling, wat de totale rekentijd aanzienlijk verkortte.

Betekenis en Conclusie

Dit paper toont aan dat FP64 Tensor Cores, die oorspronkelijk werden geïntroduceerd voor grote matrixvermenigvuldigingen, ook extreem effectief kunnen worden ingezet voor de kleine, onregelmatige matrixoperaties die kenmerkend zijn voor hoog-orde finite-elementen methoden.

De belangrijkste implicaties zijn:

1. Versnelling van wetenschappelijke ontdekking: Door de rekentijd voor complexe inversieproblemen (zoals tsunami-voorspelling) te halveren, worden real-time simulaties haalbaarder.
2. Energiebesparing: De aanzienlijke verbetering in energie-efficiëntie is cruciaal voor exascale-computing, waar energiekosten een beperkende factor zijn.
3. Open Source Impact: De optimalisaties worden geïntegreerd in de MFEM-bibliotheek, waardoor de gehele HPC-gemeenschap toegang krijgt tot deze prestatiewinsten zonder zelf de lage-niveau CUDA-programmering te hoeven doen.

Kortom, dit werk markeert een doorbraak in het benutten van moderne GPU-architecturen voor dubbele precisie wetenschappelijke berekeningen, waarbij hardware-specifieke optimalisaties leiden tot zowel snelheids- als energie-voordelen op exascale-schaal.