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

Questo lavoro presenta un'architettura a pipeline a cascata in cinque stadi implementata su array AI Engine AMD Versal per valutare in modo efficiente gli elementi di matrice di ordine principale per il processo $\gamma\gamma \to t\bar{t}g$, raggiungendo una velocità di elaborazione prevista di $1.0\times10^6$ valutazioni al secondo con un aumento di velocità di $34\times$ e un miglioramento dell'efficienza energetica di $7.7\times$ rispetto a un singolo core CPU, mantenendo al contempo una precisione numerica dell'ordine di parti per milione.

L'essenziale

Immagina di dover prevedere l'esito di una collisione massiccia e caotica tra due particelle minuscole (come i protoni) all'interno di un gigantesco acceleratore di particelle. Per fare ciò, i fisici utilizzano una ricetta matematica complessa chiamata "elemento di matrice". Calcolare questa ricetta è come risolvere un gigantesco puzzle a più fasi. Il problema è che per ottenere una risposta affidabile, devono risolvere questo stesso puzzle milioni di volte, ogni volta con condizioni iniziali leggermente diverse.

Attualmente, fare questo su processori informatici standard (CPU) è come cercare di risolvere questi puzzle uno per uno con una sola persona. È accurato, ma incredibilmente lento e consuma molta energia, specialmente man mano che l'acceleratore di particelle diventa più potente.

Questo articolo presenta un nuovo modo per risolvere questi puzzle utilizzando un tipo speciale di chip informatico chiamato AMD Versal AI Engine. Invece di avere una sola persona che risolve l'intero puzzle, gli autori hanno costruito una catena di montaggio direttamente all'interno del chip.

Ecco come funziona la loro soluzione, scomposta in concetti semplici:

1. Il problema della "Catena di Montaggio"

La ricetta matematica per questa specifica collisione di particelle (due gluoni che si trasformano in un quark top, un antiquark top e un altro gluone) è troppo grande per entrare nella memoria di un singolo processore minuscolo sul chip. Pensa a come cercare di inserire un manuale di istruzioni di 38 pagine in una tasca che può contenere solo 16 pagine.

La Soluzione: Gli autori hanno diviso il manuale in cinque capitoli. Hanno creato una catena di montaggio a cinque stadi.

• Stadio 1: Legge gli ingredienti grezzi (i dati della collisione) e prepara i primi passaggi.
• Stadio 2 e 3: Passano il lavoro lungo la linea, aggiungendo ulteriori passaggi al calcolo.
• Stadio 4 e 5: Completano i calcoli finali ed emettono la risposta.

2. Il "Nastro Trasportatore" (Pipeline a Cascata)

Questi cinque stadi sono collegati da un nastro trasportatore dedicato e superveloce chiamato interfaccia a cascata.

• Immagina una fabbrica in cui i lavoratori non si fermano per parlare o attendere il permesso di passare una scatola alla persona successiva. Scivolano semplicemente la scatola lungo una canna fumaria istantaneamente.
• In questo chip, le "scatole" sono blocchi di dati chiamati token.
• Gli autori hanno progettato un regolamento rigoroso (un "contratto deterministico") per garantire che i lavoratori non rimangano mai bloccati in attesa l'uno dell'altro. Ogni lavoratore sa esattamente quando passare una scatola e quando riceverne una, così la linea non si inceppa mai.

3. La "Super-Fabbrica" (80 Linee Contemporaneamente)

Il chip che hanno utilizzato (il VCK190) è come un enorme magazzino contenente 400 piccoli lavoratori (chiamati tile).

• Invece di costruire una sola catena di montaggio, ne hanno costruite 80 identiche una accanto all'altra.
• Ogni linea ha 5 lavoratori. $80 \text{ linee} \times 5 \text{ lavoratori} = 400 \text{ lavoratori}$.
• Stanno tutti lavorando contemporaneamente, risolvendo 80 puzzle diversi simultaneamente.

4. I Risultati: Velocità ed Efficienza

Gli autori hanno testato questa "fabbrica" contro due altri metodi: un processore informatico standard (CPU) e una scheda grafica di fascia alta (GPU).

• Velocità: La loro fabbrica a 80 linee è 34 volte più veloce di un singolo core informatico standard.
  • Nota: Una scheda grafica di punta (GPU) è ancora più veloce in assoluto (circa 22 volte più veloce del loro chip), ma la GPU è una macchina molto più grande e costosa.
• Energia: È qui che il loro metodo brilla. Poiché la catena di montaggio è così efficiente e specializzata, utilizza pochissima elettricità.
  • Per risolvere un puzzle, il loro chip utilizza 7,7 volte meno energia di un processore informatico standard.
  • È meno efficiente dal punto di vista energetico rispetto alla gigantesca GPU, ma la GPU consuma una quantità enorme di energia per farlo. Il metodo del chip è un "punto dolce" per situazioni in cui hai bisogno di velocità ma non puoi collegare una macchina enorme e avida di energia.

5. Controllo di Precisione

Hanno assicurato che la loro "catena di montaggio" non commettesse errori. Hanno confrontato le risposte del loro chip con un calcolo "gold standard" in doppia precisione.

• I risultati corrispondevano quasi perfettamente. La differenza era così piccola (circa 1 parte su un milione) che è considerata trascurabile per i calcoli fisici che stanno eseguendo.

Riassunto

In breve, gli autori hanno preso un calcolo fisico complesso che era troppo grande per un singolo chip informatico, lo hanno tagliato in cinque parti gestibili e hanno costruito 80 catene di montaggio parallele per risolverle tutte contemporaneamente. Questo approccio crea un "punto dolce" di alta velocità e basso consumo energetico, offrendo un'alternativa potente per eseguire le simulazioni necessarie per comprendere l'universo al Large Hadron Collider.

Sommario tecnico

1. Enunciazione del Problema

I moderni generatori di eventi per la Fisica delle Alte Energie (HEP), come MadGraph5_aMC@NLO (MG5aMC), affrontano un significativo collo di bottiglia computazionale nella valutazione dell'elemento di matrice ($|M|^2$) per le collisioni di particelle. Mentre il Large Hadron Collider (LHC) entra nella sua fase ad alta luminosità, la domanda per questi calcoli cresce in modo non lineare, mentre la scalabilità delle CPU rimane limitata.

• Il Collo di Bottiglia: La valutazione dell'elemento di matrice rappresenta il 30–40% del tempo totale di generazione degli eventi, in particolare per i processi multi-jet che coinvolgono emissioni reali aggiuntive.
• La Sfida: Le soluzioni esistenti che utilizzano GPU (ad esempio CUDACPP) offrono un'alta velocità di elaborazione ma consumano molta energia. Le Field-Programmable Gate Arrays (FPGA) offrono efficienza energetica ma faticano con il vincolo di 16 kB di memoria di programma (PM) per tile sulle moderne matrici AI Engine. Un'implementazione monolitica di processi complessi (come $gg \to t\bar{t}g$) supera questo limite di memoria, impedendo la mappatura diretta su un singolo tile.

2. Metodologia

Gli autori propongono un'architettura a pipeline a cascata sulla matrice AMD Versal AI Engine (AIE) (specificamente la piattaforma VCK190) per superare i vincoli di memoria e massimizzare il parallelismo.

A. Piattaforma Target e Architettura

• Hardware: AMD Versal XCVC1902 ACAP dotato di 400 tile AI Engine disposti in una griglia $50 \times 8$, con clock a 1,25 GHz.
• Decomposizione della Pipeline: Il processo $gg \to t\bar{t}g$ (che coinvolge 16 diagrammi di Feynman e 10 funzioni HELAS distinte) è decomposto in una pipeline a cinque stadi.
  • Stadio 1: Generazione della funzione d'onda (spinori/vettori esterni) e inizializzazione dei token.
  • Stadi 2 e 3: Valutazioni dei vertici fermione-vettore (suddivisione di 12 diagrammi per bilanciare la memoria).
  • Stadio 4: Valutazioni del vertice a tre gluoni (incluso un generatore di bosoni off-shell differito).
  • Stadio 5: Termini di contatto a quattro gluoni e riduzione della matrice di colore.
• Gestione della Memoria: Per rientrare nel limite di 16 kB, gli autori hanno adottato la partizionamento della memoria di programma e la valutazione differita. Ad esempio, il generatore di bosoni off-shell ($FFV1P0\_3$) è stato spostato dallo Stadio 1 allo Stadio 4, riducendo l'uso di memoria dello Stadio 1 da 17,8 kB a 15,5 kB.

B. Comunicazione Inter-Tile (Protocollo a Cascata)

• Meccanismo: Gli stadi comunicano tramite l'interfaccia a cascata unidirezionale a 384 bit (larghezza di banda di 60 GB/s).
• Protocollo dei Token: Viene utilizzato un protocollo deterministico, privo di deadlock, in cui gli stadi scambiano "token" strutturati contenenti funzioni d'onda e ampiezze parziali.
  • Token Esteso (Stadi 1–4): Trasporta 5 funzioni d'onda esterne, 3 propagatori precalcolati e 6 ampiezze di flusso di colore (18 beat/elicità).
  • Token Ridotto (Stadi 4–5): Trasporta solo 5 funzioni d'onda e ampiezze (12 beat/elicità) dopo la valutazione locale del vertice a tre gluoni.
• Determinismo: Il sistema impone un "contratto a cascata" con strutture di loop identiche, scritture incondizionate e conteggi di token staticamente abbinati per garantire un'operazione senza sovraccarichi e senza controllo di flusso.

C. Adattamenti Software

• Porting della Libreria HELAS: La libreria HELAS standard di MG5aMC (originariamente C++ in doppia precisione scalare) è stata adattata agli intrinseci vettoriali AI Engine (precisione singola float32).
• Ottimizzazioni:
  • Vettorializzazione: Le funzioni d'onda mappate su vettori SIMD a 8 vie.
  • Divisione Complessa: Sostituzione del metodo di Smith (2 divisioni) con un'unica istruzione hardware di reciproco.
  • Caching dell'Elicità: Precalcolo di 10 funzioni d'onda per 32 configurazioni di elicità, selezionandole tramite ricerca indicizzata a bit per ridurre le valutazioni di 16 volte.
  • Riduzione del Colore: Incorporazione delle divisioni di normalizzazione del colore in costanti a tempo di compilazione.

D. Distribuzione del Sistema

• Scalabilità: 80 pipeline indipendenti sono state mappate sui 400 tile disponibili (5 tile per pipeline).
• I/O: Un'architettura a commutazione di pacchetti nella Programmable Logic (PL) distribuisce i punti dello spazio delle fasi alle pipeline e raccoglie i risultati.

3. Contributi Chiave

1. Architettura a Pipeline Guidata dalla Memoria: Introduzione di una nuova pipeline a cascata a 5 stadi che partiziona con successo un calcolo complesso di elementi di matrice multi-diagramma su più tile AI Engine, superando il vincolo di 16 kB della PM.
2. Contratto a Cascata Deterministico: Sviluppo di un protocollo di comunicazione privo di deadlock utilizzando token di funzione d'onda e strutture di loop statiche, eliminando la necessità di hardware complesso di controllo di flusso.
3. Porting Completo di HELAS: Porting riuscito dell'intera libreria di ampiezze HELAS agli intrinseci vettoriali AI Engine, incorporando ottimizzazioni complesse come il caching dell'elicità indicizzato binariamente e la divisione complessa ridotta.
4. Distribuzione Scalabile: Dimostrazione di una distribuzione teorica di 80 pipeline che utilizza il 100% delle risorse di calcolo AI Engine del VCK190.

4. Risultati

• Velocità di Elaborazione: La velocità di elaborazione proiettata per la matrice a 80 pipeline è di $1.0 \times 10^6$ valutazioni di elementi di matrice al secondo (ME/s).
  • Ciò rappresenta un accelerazione di 34 volte rispetto a un singolo core CPU (Intel i5-10600).
  • Sebbene inferiore a una GPU NVIDIA A100 ($2.18 \times 10^7$ ME/s), la soluzione AI Engine è significativamente più efficiente dal punto di vista energetico.
• Efficienza Energetica:
  • AI Engine: 54,8 µJ/ME (a 54,8 W di potenza nel dominio AIE).
  • CPU: 422 µJ/ME.
  • GPU: 7,3 µJ/ME (ma a 159 W di potenza).
  • Miglioramento: L'AI Engine offre un miglioramento di 7,7 volte nell'efficienza energetica rispetto alla linea di base CPU.
• Precisione: Validata contro il riferimento in doppia precisione di MG5aMC.
  • Errore relativo medio: 1,43 ppm (parti per milione).
  • Errore relativo massimo: 168 ppm.
  • Questo livello di precisione è considerato sufficiente per i calcoli di Leading-Order (LO) dove le incertezze fisiche (variazione di scala, PDF) dominano gli errori numerici.
• Utilizzo delle Risorse:
  • Memoria di Programma: Lo Stadio 1 è il collo di bottiglia con un utilizzo del 94,7% (15.514 byte).
  • Programmable Logic: Utilizzo modesto (4,72% LUT, 2,87% Registri), lasciando spazio per logica aggiuntiva.

5. Significato e Lavori Futuri

• Significato: Questo lavoro dimostra che le matrici AI Engine sono valide per la generazione di eventi HEP ad alta velocità e ad alta efficienza energetica, in particolare per ambienti con vincoli di potenza (ad esempio sistemi di trigger online o edge computing all'LHC) dove i profili di potenza delle GPU sono proibitivi. Stabilisce una metodologia sistematica per partizionare kernel fisici complessi su tile con memoria vincolata.
• Limitazioni: L'implementazione attuale è limitata ai processi di Leading-Order (LO). Le cifre di latenza si basano su simulatori approssimati per ciclo piuttosto che sul timing hardware diretto dell'intera matrice.
• Direzioni Future:
  • Filtraggio dell'Elicità: Pre-calcolo di maschere di elicità valide per ridurre le iterazioni del ciclo interno, potenzialmente raddoppiando la velocità di elaborazione.
  • Maggiore Multiplicità: Estensione della profondità della pipeline per processi più complessi (ad esempio $t\bar{t}ggg$).
  • Integrazione NLO: Adattamento dell'architettura per calcoli Next-to-Leading-Order che coinvolgono integrali di loop.
  • Evoluzione Hardware: Sfruttamento di dispositivi Versal di prossima generazione con matrici più grandi o velocità di clock superiori.

In conclusione, il documento presenta un'alternativa robusta ed efficiente dal punto di vista energetico all'accelerazione GPU per carichi di lavoro HEP specifici, sfruttando le capacità uniche di cascata dell'AI Engine AMD Versal per risolvere le sfide di partizionamento della memoria intrinseche nei calcoli complessi degli elementi di matrice.