Designing quantum chemistry algorithms with just-in-time compilation

Dit paper introduceert just-in-time compilatie voor integraalkernen van Gauss-type orbitalen, wat leidt tot aanzienlijke snelheidswinsten bij de berekening van elektronenrepulsie-integralen en JK-matrices op GPU's, met prestaties die tot viermaal beter zijn dan bestaande methoden.

De kern

De Kern: Een "Op Maat Gemaakte" Rekenmachine voor Moleculen

Stel je voor dat chemici proberen te voorspellen hoe moleculen zich gedragen. Dit is als het proberen te voorspellen hoe een enorme groep mensen in een drukke stad met elkaar omgaat. De wiskunde hierachter (het berekenen van hoe elektronen elkaar afstoten) is enorm complex en vereist miljarden berekeningen.

Voor decennia hebben supercomputers deze berekeningen gedaan met een statische aanpak. Dit is alsof je een enorme, universele handleiding hebt voor elke mogelijke situatie in de stad. Of het nu regent, of er een festival is, of dat er maar drie mensen zijn: je gebruikt altijd dezelfde, zware handleiding. Dit werkt, maar het is traag en inefficiënt omdat je veel tijd besteedt aan regels die op dat moment niet nodig zijn.

JoltQC (de nieuwe software in dit paper) introduceert een revolutionaire nieuwe manier: Just-in-Time (JIT) compilatie.

De Metafoor: De Kookplaat vs. De Voorverwarmde Oven

Om te begrijpen wat JIT doet, laten we kijken naar koken:

1. De Oude Methode (AOT - Ahead-of-Time):
   Stel je voor dat je een gigantische oven hebt die je altijd voorverwarmt op de hoogste stand, ongeacht wat je gaat bakken. Of je nu een klein koekje (een simpel molecuul) of een enorme taart (een complex molecuul) wilt maken, de oven staat altijd op vol vermogen. Je verliest veel energie en tijd omdat de oven niet slim is. Dit is hoe oude quantum-chemie-software werkte: het codeerde alles vooraf in één groot, stijf programma.

2. De Nieuwe Methode (JIT - Just-in-Time):
   JoltQC is als een slimme, aanpasbare kookplaat.

   • Zodra je zegt: "Ik ga vandaag een taart met aardbeien bakken" (een specifiek molecuul met bepaalde eigenschappen), past de kookplaat zich direct aan.
   • Hij schakelt alleen de branders in die je nodig hebt.
   • Hij verwarmt precies op de temperatuur die voor die taart nodig is.
   • Hij doet dit in een splitseconde, terwijl je kookt.

In de taal van de computer betekent dit: in plaats van één groot, zwaar programma te gebruiken voor alles, genereert JoltQC op het moment zelf (tijdens het rekenen) een klein, perfect geoptimaliseerd stukje code dat exact past bij het molecuul dat je onderzoekt.

Waarom is dit zo snel? (De "Grote" en "Kleine" Problemen)

De onderzoekers van ByteDance Seed hebben twee grote problemen opgelost:

1. Het probleem van de "Kleine" Moleculen (Lage hoekmomenten)
Bij simpele moleculen (zoals water of kleine ketens) zijn de berekeningen klein, maar er zijn er ontzettend veel.

• De oude manier: De computer moest steeds controleren: "Moet ik nu dit doen? Of dat? Is de temperatuur goed?" Dit kostte tijd (zoals een chef die elke seconde de receptenboeken moet raadplegen).
• De JoltQC-oplossing: Omdat de code op maat wordt gemaakt, zijn die vragen weg. De computer weet al precies wat hij moet doen. Het is alsof de chef het recept uit zijn hoofd kent en direct begint met bakken.
• Resultaat: Op de nieuwste grafische kaarten (GPUs) is dit 2 keer sneller dan de beste bestaande software.

2. Het probleem van de "Grote" Moleculen (Hoge hoekmomenten)
Bij complexe moleculen (met zware atomen of grote basisgroepen) worden de berekeningen zo groot dat ze niet meer in het snelle geheugen van de computer passen.

• De oude manier: De computer moest constant heen en weer springen tussen het snelle geheugen (registers) en het langzamere geheugen (RAM). Dit is als een kok die constant de keuken uit moet rennen om ingrediënten uit de kelder te halen.
• De JoltQC-oplossing: Ze hebben een nieuw "fragmentatie"-systeem bedacht. In plaats van één kok die alles doet, delen ze het werk op in kleine stukjes die door een team van koks (threads) tegelijk worden gedaan. Ze gebruiken een slimme "opslagkast" (shared memory) zodat niemand hoeft te rennen.
• Resultaat: Voor deze zware moleculen is het 4 keer sneller.

Het Magische Trucje: Enkele Precisie (FP32)

Normaal gesproken gebruiken wetenschappers "dubbele precisie" (FP64) voor berekeningen, wat betekent dat ze tot 15 decimalen nauwkeurig rekenen. Dit is als het meten van een afstand tot op een atoom nauwkeurig. Maar voor de meeste chemische berekeningen is "enkele precisie" (FP32) – met 7 decimalen – al meer dan genoeg.

• De Analogie: Stel je hebt een vrachtwagen die 1000 dozen kan vervoeren (dubbele precisie). Maar je hebt maar 500 dozen nodig. Je gebruikt dus een halfvolle vrachtwagen.
• JoltQC's truc: Ze gebruiken een kleinere, lichtere bestelbus (enkele precisie) die twee keer zo snel rijdt en minder brandstof verbruikt.
• Omdat moderne grafische kaarten (zoals die in gaming-computers of AI-servers) veel sneller zijn in het verwerken van deze "lichtere" berekeningen, wordt de software 3 keer sneller op deze hardware, zonder dat de nauwkeurigheid voor de meeste toepassingen vermindert.

Waarom is dit belangrijk voor de toekomst?

1. Minder Code, Meer Snelheid: De kern van dit nieuwe systeem is slechts ongeveer 1.000 regels code (vergeleken met 20.000 regels in de oude software). Het is lichter, makkelijker te onderhouden en makkelijker aan te passen aan nieuwe computers.
2. Open Source: De onderzoekers hebben hun code ("JoltQC") gratis beschikbaar gesteld. Dit betekent dat elke wetenschapper deze "slimme kookplaat" nu kan gebruiken.
3. Toekomstbestendig: Omdat de software zich aanpast aan de hardware, blijft het snel zijn, zelfs als er in de toekomst nog krachtigere computers komen.

Kortom: JoltQC maakt het rekenen van complexe chemische reacties niet alleen sneller, maar ook slimmer. Het stopt met het gebruiken van één groot, zwaar gereedschap voor alles en begint met het maken van perfect op maat gemaakte tools op het moment dat ze nodig zijn. Hierdoor kunnen wetenschappers grotere moleculen bestuderen en nieuwe medicijnen of materialen veel sneller ontdekken.

Technische samenvatting

Titel: Ontwerp van kwantumchemie-algoritmes met just-in-time compilatie (JIT)

Auteurs: Xiaojie Wu, Qiming Sun, Yuanheng Wang (ByteDance Seed)
Project: JoltQC (Open Source Library)

---

1. Het Probleem

Traditionele kwantumchemie-software (zoals GPU4PySCF) maakt gebruik van Ahead-of-Time (AOT) compilatie en monolithische codegeneratie in Fortran of C++. Deze benadering heeft ernstige beperkingen in moderne High-Performance Computing (HPC) omgevingen, vooral op GPU's:

• Statische compilatie: De code moet conservatief alle mogelijke conditionele takken, parametercombinaties en numerieke drempels ondersteunen voor een vast aantal gebruiksscenario's.
• Inefficiëntie: Dit leidt tot "opgeblazen" binaries, excessieve controleflow-divergentie, slecht cache- en registergebruik, en inefficiënt gebruik van memory-bandwidth.
• GPU-beperkingen: Op GPU's worden deze inefficiënties versterkt door strikte beperkingen op registerdruk, warp-divergentie en hardware-bezettingsgraad.
• Hoge hoekimpuls: Hoewel lage hoekimpuls-integralen (s, p) goed presteren, worstelen bestaande pakketten met de prestaties van integrals met hoge hoekimpuls (d, f, g), die last hebben van al bovenstaande problemen.

2. Methodologie

De auteurs introduceren Just-in-Time (JIT) compilatie voor de integral-kernels van Gauss-type orbitalen (GTO's). In plaats van vooraf te compileren, wordt code gegenereerd tijdens de runtime op basis van de daadwerkelijke invoer.

Kernontwerpbeslissingen:

• Statische vs. Dynamische parameters:
  • Statisch (Compile-time): Kwantumchemische methode, basisset-type, hoekimpuls-patronen, aantal primitieven en GPU-architectuur. Dit maakt het mogelijk om loops volledig te "unrollen" (ontrollen), registers statisch toe te wijzen en alle vertakkingen te elimineren.
  • Dynamisch (Runtime): Moleculaire informatie zoals atoomcoördinaten, exponenten en contractiecoëfficiënten. Deze worden als invoer doorgegeven.
• Implementatie: De kern is gebouwd met NVRTC (CUDA runtime compilation library) via CuPy. De totale implementatie is compact (~1.000 regels code) vergeleken met ~20.000 regels in GPU4PySCF v1.4.
• Mixed Precision: Dankzij de Schwarz-ongelijkheid kunnen de meeste schelp-kwartetten (shell quartets) in enkele precisie (FP32) worden berekend in plaats van dubbele precisie (FP64). JIT maakt dit eenvoudig: dezelfde broncode genereert zowel FP32- als FP64-kernels zonder code-duplicatie. FP32 biedt voordelen zoals hogere rekenkracht op consumentengpu's, minder geheugengebruik en versnelde hardware-instructies voor functies zoals exp en erf.

Twee Specifieke Algoritmen:

1. 1q1t (One Quartet One Thread):
   • Toewijzing van één thread per schelp-kwartet.
   • Ideaal voor lage hoekimpuls (s, p).
   • Profiteert van agressieve loop-unrolling omdat alle grenzen compile-time bekend zijn.
2. 1qnt (Quartet Fragmentation):
   • Ontworpen voor hoge hoekimpuls (d, f, g) waar registers tekortschieten.
   • Gebruikt een dynamische fragmentatie-strategie: een groep threads (2 tot 256) werkt samen aan één schelp-kwartet.
   • Multi-level reduction: Tussenresultaten worden in shared memory opgeslagen en via een hiërarchisch proces (thread-niveau en block-niveau) gereduceerd om memory-bandwidth-bottlenecks te verminderen.
   • De fragmentatiegrootte wordt geoptimaliseerd via een grid search voor specifieke GPU's en precisies.

3. Belangrijkste Bijdragen

1. Introductie van JIT in Kwantumchemie: De auteurs brengen JIT naar het domein en optimaliseren de twee-elektron integral-kernels. Het framework integreert naadloos met GPU4PySCF en is beschikbaar als de open-source bibliotheek JoltQC.
2. Nieuw Fragmentatie-algoritme: Een innovatief algoritme voor integrals met hoge hoekimpuls dat data-localiteit verbetert en memory-bandwidth-bottlenecks verlicht door middel van multi-level reductie.
3. Prestatiebewijzen: Benchmarking van single-precision kernels onder het JIT-framework toont aanzienlijke prestatiewinsten, vooral in combinatie met de nieuwe algoritmes.

4. Resultaten

De prestaties zijn getest op NVIDIA A100-80G (FP64) en A10-24G (FP32) GPU's.

• Snelheidswinst vs. GPU4PySCF v1.4:
  • Kleine basissets (6-31G):* 2x snelheidswinst op A100 (FP64).
  • Grote basissets (def2-TZVPP): Tot 4x snelheidswinst op A100.
  • Single Precision (FP32): Op de A10-24G bereikt de 1q1t-algoritme een 30x snelheidswinst ten opzichte van FP64, en de 1qnt-algoritme voor hoge hoekimpuls is 10-20x sneller dan de bestaande methoden.
• Vergelijking met TeraChem:
  • Voor grote basissets (def2-TZVPP) is JoltQC ongeveer 4x sneller dan TeraChem v1.9 (FP64).
  • Voor basissets met g-functies (cc-pVQZ), die TeraChem niet ondersteunt, behaalt JoltQC een 7-8x snelheidswinst ten opzichte van GPU4PySCF.
  • In FP32 is JoltQC ongeveer 3x sneller dan TeraChem voor de JK-kernels.
• Compilatiekosten: De JIT-compilatie introduceert een eenmalige overhead (bijv. ~200s voor 750 kernels bij def2-TZVPP), maar deze wordt geamortiseerd door caching op schijf en in het geheugen. Na de eerste compilatie zijn laadtijden <1s.
• Nauwkeurigheid: In FP64 is JoltQC numeriek identiek aan GPU4PySCF. In FP32 worden de resultaten geaccumuleerd in FP64 om nauwkeurigheid te behouden; de energieafwijkingen zijn sub-milliHartree (<1 mHa), wat acceptabel is voor veel toepassingen.

5. Betekenis en Conclusie

Dit paper demonstreert dat JIT-compilatie een transformatieve paradigmaverschuiving is voor kwantumchemie op GPU's.

• Prestaties: Door loops te unrollen en registers optimaal te benutten, worden de beperkingen van statische compilatie opgelost, vooral voor complexe, hoge-hoekimpuls berekeningen.
• Ontwikkelervaring: De code is aanzienlijk compacter (1.000 vs 20.000 regels), makkelijker te onderhouden en uit te breiden. Het stelt ontwikkelaars in staat om snel nieuwe algoritmen te prototypen zonder het hele systeem opnieuw te hoeven compileren.
• Toekomst: De auteurs zien een toekomst waarin JIT een standaardtool wordt in de kwantumchemie-stack, met potentie voor verdere integratie met machine learning pipelines (via frameworks zoals Triton of JAX) en geavanceerde mixed-precision strategieën.

De JoltQC bibliotheek is beschikbaar op GitHub en biedt een praktische basis voor het toepassen van JIT in kwantumchemische workflows.