Cascade Pipeline for Leading-Order Matrix Element Evaluation on AMD Versal AI Engine Arrays

Dit artikel presenteert een vijfstaps cascade-pijplijnarchitectuur geïmplementeerd op AMD Versal AI Engine-arrays om leidende-orde matrixelementen voor het $\gamma\gamma \to t\bar{t}g$-proces efficiënt te evalueren, met een projecteerde doorvoer van $1.0\times10^6$ evaluaties per seconde, een $34\times$ snelheidswinst en een $7.7\times$ verbetering in energie-efficiëntie ten opzichte van één CPU-kern, terwijl numerieke nauwkeurigheid op het niveau van delen per miljoen behouden blijft.

De kern

Stel je voor dat je probeert de uitkomst te voorspellen van een enorme, chaotische botsing tussen twee kleine deeltjes (zoals protonen) binnen een gigantische deeltjesversneller. Om dit te doen, gebruiken natuurkundigen een complexe wiskundige recept genaamd een "matrixelement". Het berekenen van dit recept is als het oplossen van een gigantisch meerstapsraadsel. Het probleem is dat ze, om een betrouwbaar antwoord te krijgen, ditzelfde raadsel miljoenen keren moeten oplossen, elke keer met lichtjes verschillende startvoorwaarden.

Op dit moment is het doen hiervan op standaard computerprocessors (CPU's) als het proberen om deze raadsels één voor één op te lossen met één persoon. Het is accuraat, maar het is ongelooflijk traag en verbruikt veel energie, vooral naarmate de deeltjesversneller krachtiger wordt.

Dit artikel presenteert een nieuwe manier om deze raadsels op te lossen met een speciaal type computerchip genaamd de AMD Versal AI Engine. In plaats van dat één persoon het hele raadsel oplost, hebben de auteurs een fabrieksassemblagelijn gebouwd, direct binnen de chip.

Hier is hoe hun oplossing werkt, opgesplitst in eenvoudige concepten:

1. Het "Assemblagelijn"-probleem

Het wiskundige recept voor deze specifieke deeltjesbotsing (twee gluonen die veranderen in een top-quark, een anti-top-quark en nog een gluon) is te groot om in het geheugen van één kleine processor op de chip te passen. Denk hieraan als het proberen om een 38 pagina's tellende handleiding in een zak te proppen die slechts 16 pagina's kan bevatten.

De Oplossing: De auteurs hebben de handleiding opgesplitst in vijf hoofdstukken. Ze hebben een vijfstaps assemblagelijn gecreëerd.

• Stap 1: Leest de ruwe ingrediënten (de botsingsdata) en bereidt de eerste paar stappen voor.
• Stap 2 & 3: Geven het werk langs de lijn door, waarbij meer stappen aan de berekening worden toegevoegd.
• Stap 4 & 5: Voltooien de laatste berekeningen en spugen het antwoord uit.

2. De "Transportband" (Cascadepijpleiding)

Deze vijf stappen zijn verbonden door een supersnelle, speciale transportband genaamd een cascade-interface.

• Stel je een fabriek voor waar werknemers niet stoppen om te praten of wachten op toestemming om een doos door te geven aan de volgende persoon. Ze schuiven de doos gewoon direct via een goot naar de volgende persoon.
• In deze chip zijn de "dozen" stukken data die tokens worden genoemd.
• De auteurs hebben een strikt regelboek ontworpen (een "deterministisch contract") om ervoor te zorgen dat de werknemers nooit vastlopen terwijl ze op elkaar wachten. Elke werknemer weet precies wanneer hij een doos moet doorgeven en wanneer hij er één moet ontvangen, zodat de lijn nooit verstopt raakt.

3. De "Superfabriek" (80 Lijnen Tegelijk)

De chip die ze gebruikten (de VCK190) is als een enorm magazijn met 400 kleine werknemers (die tegels worden genoemd).

• In plaats van slechts één assemblagelijn te bouwen, bouwden ze 80 identieke assemblagelijnen naast elkaar.
• Elke lijn heeft 5 werknemers. $80 \text{ lijnen} \times 5 \text{ werknemers} = 400 \text{ werknemers}$.
• Ze werken allemaal tegelijk, waarbij ze 80 verschillende raadsels tegelijkertijd oplossen.

4. De Resultaten: Snelheid en Efficiëntie

De auteurs hebben deze "fabriek" getest tegen twee andere methoden: een standaard computerprocessor (CPU) en een high-end grafische kaart (GPU).

• Snelheid: Hun fabriek met 80 lijnen is 34 keer sneller dan een enkele standaard computerkern.
  • Opmerking: Een topklasse grafische kaart (GPU) is nog steeds overall sneller (ongeveer 22 keer sneller dan hun chip), maar de GPU is een veel grotere en duurdere machine.
• Energie: Hier blinkt hun methode uit. Omdat de assemblagelijn zo efficiënt en gespecialiseerd is, verbruikt het zeer weinig elektriciteit.
  • Om één raadsel op te lossen, verbruikt hun chip 7,7 keer minder energie dan een standaard computerprocessor.
  • Het is minder energie-efficiënt dan de gigantische GPU, maar de GPU verbruikt een enorme hoeveelheid vermogen om dit te doen. De methode van de chip is een "sweet spot" voor situaties waarbij je snelheid nodig hebt, maar geen enorme stroomverslindende machine kunt aansluiten.

5. Nauwkeurigheidscontrole

Ze zorgden ervoor dat hun "assemblagelijn" geen fouten maakte. Ze vergeleken de antwoorden van hun chip met een "gouden standaard" berekening met dubbele precisie.

• De resultaten kwamen bijna perfect overeen. Het verschil was zo klein (ongeveer 1 op een miljoen) dat het als verwaarloosbaar wordt beschouwd voor de natuurkundeberekeningen die ze uitvoeren.

Samenvatting

Kortom, de auteurs namen een complexe natuurkundeberekening die te groot was voor één computerchip, hakten deze in vijf beheersbare stukken en bouwden 80 parallelle assemblagelijnen om ze allemaal tegelijk op te lossen. Deze aanpak creëert een "sweet spot" van hoge snelheid en laag energieverbruik, en biedt een krachtig alternatief voor het uitvoeren van de simulaties die nodig zijn om het universum te begrijpen bij de Large Hadron Collider.

Technische samenvatting

1. Probleemstelling

Moderne High Energy Physics (HEP) eventgeneratoren, zoals MadGraph5_aMC@NLO (MG5aMC), staan voor een aanzienlijke computatiefase bij het evalueren van het matrixelement ($|M|^2$) voor deeltjesbotsingen. Naarmate de Large Hadron Collider (LHC) de High-Luminosity-fase ingaat, groeit de vraag naar deze berekeningen niet-lineair, terwijl de schaalbaarheid van CPU's beperkt blijft.

• De Bottleneck: De evaluatie van matrixelementen staat voor 30–40% van de totale tijd voor eventgeneratie, met name voor multi-jet-processen die extra reële emissies omvatten.
• De Uitdaging: Bestaande oplossingen met GPU's (bijv. CUDACPP) bieden een hoge doorvoer maar verbruiken aanzienlijk vermogen. Field-Programmable Gate Arrays (FPGA's) bieden energie-efficiëntie, maar kampen met de beperking van 16 kB programmageheugen (PM) per tile op moderne AI Engine-arrays. Een monolithische implementatie van complexe processen (zoals $gg \to t\bar{t}g$) overschrijdt deze geheugenlimiet, waardoor directe toewijzing aan een enkele tile onmogelijk is.

2. Methodologie

De auteurs stellen een cascade-pipeline-architectuur voor op het AMD Versal AI Engine (AIE)-array (specifiek het VCK190-platform) om geheugenbeperkingen te overwinnen en parallelisme te maximaliseren.

A. Doelplatform & Architectuur

• Hardware: AMD Versal XCVC1902 ACAC met 400 AI Engine-tiles gerangschikt in een $50 \times 8$-rooster, geklokt op 1,25 GHz.
• Pipeline-decompositie: Het proces $gg \to t\bar{t}g$ (met 16 Feynmandiagrammen en 10 verschillende HELAS-functies) wordt opgesplitst in een vijftraps-pipeline.
  • Stap 1: Golfunctiegeneratie (externe spinoren/vectoren) en initialisatie van tokens.
  • Stap 2 & 3: Fermion-vector-vertexevaluaties (opsplitsing van 12 diagrammen om het geheugen in evenwicht te brengen).
  • Stap 4: Triple-gluon-vertexevaluaties (inclusief een vertraagde generator voor off-shell-bosonen).
  • Stap 5: Vier-gluon-contacttermen en reductie van de kleurmatrijs.
• Geheugenbeheer: Om binnen de limiet van 16 kB te blijven, hebben de auteurs partitionering van programmageheugen en vertraagde evaluatie toegepast. Bijvoorbeeld, de off-shell-bosongenerator ($FFV1P0\_3$) is verplaatst van Stap 1 naar Stap 4, waardoor het geheugengebruik van Stap 1 werd gereduceerd van 17,8 kB naar 15,5 kB.

B. Communicatie tussen Tiles (Cascade-protocol)

• Mechanisme: De stappen communiceren via de 384-bit unidirectionele cascade-interface (bandbreedte van 60 GB/s).
• Token-protocol: Er wordt een deterministisch, deadlock-vrij protocol gebruikt waarbij stappen gestructureerde "tokens" uitwisselen die golfuncties en partiële amplitude bevatten.
  • Uitgebreid Token (Stap 1–4): Draagt 5 externe golfuncties, 3 vooraf berekende propagatoren en 6 kleurstroom-amplitudes (18 beats/heliciteit).
  • Gereduceerd Token (Stap 4–5): Draagt na lokale evaluatie van de triple-gluon-vertex alleen nog 5 golfuncties en amplitudes (12 beats/heliciteit).
• Determinisme: Het systeem dwingt een "cascade-contract" af met identieke lusstructuren, voorwaardelijke schrijfbewerkingen en statisch afgestemde tokenaantallen om een werking zonder overhead en stroomregeling te garanderen.

C. Software-aanpassingen

• Porting van HELAS-bibliotheek: De standaard MG5aMC HELAS-bibliotheek (oorspronkelijk scalair dubbel-precisie C++) is gepporteerd naar AI Engine-vector-intrinsieken (enkel-precisie float32).
• Optimalisaties:
  • Vectorisatie: Golfuncties gemapt naar 8-brede SIMD-vectoren.
  • Complexe deling: Smith's methode (2 delingen) vervangen door één hardware-reciproque-instructie.
  • Heliciteit-caching: Vooraf berekende 10 golfuncties voor 32 heliciteitsconfiguraties, geselecteerd via bit-geïndexeerde opzoektabel om evaluaties met 16× te verminderen.
  • Kleurreductie: Delingen voor kleurnormalisatie samengevoegd tot compile-time constanten.

D. Systeemimplementatie

• Schaal: 80 onafhankelijke pipelines zijn toegewezen aan de 400 beschikbare tiles (5 tiles per pipeline).
• I/O: Een pakketgeschakelde architectuur in de Programmable Logic (PL) verdeelt fase-ruimtepunten over de pipelines en verzamelt resultaten.

3. Belangrijkste Bijdragen

1. Geheugengedreven Pipeline-architectuur: Introductie van een nieuwe 5-traps cascade-pipeline die een complexe matrixelementberekening met meerdere diagrammen succesvol partitioneert over meerdere AI Engine-tiles, waardoor de 16 kB PM-beperking wordt overwonnen.
2. Deterministisch Cascade-contract: Ontwikkeling van een deadlock-vrij communicatieprotocol met golfunctietokens en statische lusstructuren, waardoor complexe hardware voor stroomregeling overbodig wordt.
3. Volledige HELAS-port: Succesvolle porting van de volledige HELAS-amplitudelibrarie naar AI Engine-vector-intrinsieken, met inbegrip van complexe optimalisaties zoals binair-geïndexeerde heliciteit-caching en gereduceerde complexe deling.
4. Schalbare Implementatie: Demonstratie van een theoretische implementatie van 80 pipelines die 100% van de AI Engine-berekeningsbronnen van de VCK190 benutten.

4. Resultaten

• Doorvoer: De geprojecteerde doorvoer voor het array van 80 pipelines is $1.0 \times 10^6$ matrixelementevaluaties per seconde (ME/s).
  • Dit vertegenwoordigt een 34× snelheidswinst ten opzichte van een single-threaded CPU-core (Intel i5-10600).
  • Hoewel lager dan een NVIDIA A100 GPU ($2.18 \times 10^7$ ME/s), is de AI Engine-oplossing aanzienlijk energie-efficiënter.
• Energie-efficiëntie:
  • AI Engine: 54,8 µJ/ME (bij 54,8 W vermogen in het AIE-domein).
  • CPU: 422 µJ/ME.
  • GPU: 7,3 µJ/ME (maar bij 159 W vermogen).
  • Verbetering: De AI Engine biedt een 7,7× verbetering in energie-efficiëntie ten opzichte van de CPU-baseline.
• Precisie: Geverifieerd tegen de dubbel-precisie referentie van MG5aMC.
  • Gemiddelde relatieve fout: 1,43 ppm (deelen per miljoen).
  • Maximale relatieve fout: 168 ppm.
  • Dit niveau van precisie wordt voldoende geacht voor Leading-Order (LO)-berekeningen waarbij fysieke onzekerheden (schaalvariatie, PDF's) numerieke fouten domineren.
• Resourcegebruik:
  • Programmageheugen: Stap 1 is de bottleneck met 94,7% gebruik (15.514 bytes).
  • Programmable Logic: Beperkt gebruik (4,72% LUTs, 2,87% Registers), waardoor ruimte overblijft voor extra logica.

5. Betekenis en Toekomstig Werk

• Betekenis: Dit werk bewijst dat AI Engine-arrays levensvatbaar zijn voor HEP eventgeneratie met hoge doorvoer en energie-efficiëntie, met name voor omgevingen met vermogensbeperkingen (bijv. online triggersystemen of edge computing aan de LHC) waar GPU-vermogensomhulsels prohibitief zijn. Het vestigt een systematische methodologie voor het partitioneren van complexe fysiek-kernen over tiles met beperkt geheugen.
• Beperkingen: De huidige implementatie is beperkt tot Leading-Order (LO)-processen. De latentiecijfers zijn gebaseerd op cycle-approxieme simulatoren in plaats van directe hardware-timing van het volledige array.
• Toekomstige Richtingen:
  • Heliciteitsfiltering: Vooraf berekenen van geldige heliciteitsmaskers om de iteraties in de binnenste lus te verminderen, wat de doorvoer mogelijk kan verdubbelen.
  • Hogere Multiplicititeit: Uitbreiden van de pipeline-diepte voor complexere processen (bijv. $t\bar{t}ggg$).
  • NLO-integratie: Aanpassen van de architectuur voor Next-to-Leading-Order-berekeningen met lusintegralen.
  • Hardware-evolutie: Benutten van Versal-apparaten van de volgende generatie met grotere arrays of hogere kloksnelheden.

Concluderend presenteert het paper een robuust, energie-efficiënt alternatief voor GPU-versnelling voor specifieke HEP-werklasten, waarbij gebruik wordt gemaakt van de unieke cascade-mogelijkheden van de AMD Versal AI Engine om de uitdagingen van geheugenpartitionering inherent aan complexe matrixelementberekeningen op te lossen.