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

Cet article présente une architecture de pipeline en cascade à cinq étapes implémentée sur des réseaux de moteurs IA AMD Versal pour évaluer efficacement les éléments de matrice au premier ordre du processus $\gamma\gamma \to t\bar{t}g$, atteignant un débit projeté de $1.0\times10^6$ évaluations par seconde avec une accélération de $34\times$ et une amélioration de l'efficacité énergétique de $7.7\times$ par rapport à un seul cœur de CPU tout en maintenant une précision numérique de l'ordre de la partie par million.

L'essentiel

Imaginez que vous essayez de prédire le résultat d'une collision massive et chaotique entre deux particules minuscules (comme des protons) à l'intérieur d'un gigantesque accélérateur de particules. Pour ce faire, les physiciens utilisent une recette mathématique complexe appelée « élément de matrice ». Calculer cette recette revient à résoudre un immense puzzle en plusieurs étapes. Le problème est que, pour obtenir une réponse fiable, ils doivent résoudre ce même puzzle des millions de fois, chaque fois avec des conditions de départ légèrement différentes.

Actuellement, effectuer cela sur des processeurs d'ordinateur standard (CPU) revient à essayer de résoudre ces puzzles un par un avec une seule personne. C'est précis, mais incroyablement lent et très énergivore, surtout à mesure que l'accélérateur de particules devient plus puissant.

Cet article présente une nouvelle méthode pour résoudre ces puzzles en utilisant un type spécial de puce informatique appelé le moteur IA AMD Versal. Au lieu d'avoir une seule personne résoudre l'ensemble du puzzle, les auteurs ont construit une chaîne de montage directement à l'intérieur de la puce.

Voici comment leur solution fonctionne, décomposée en concepts simples :

1. Le problème de la « chaîne de montage »

La recette mathématique pour cette collision de particules spécifique (deux gluons se transformant en un quark top, un anti-quark top et un autre gluon) est trop volumineuse pour tenir dans la mémoire d'un seul micro-processeur sur la puce. Imaginez essayer de faire entrer un manuel d'instructions de 38 pages dans une poche ne pouvant en contenir que 16.

La solution : Les auteurs ont divisé le manuel en cinq chapitres. Ils ont créé une chaîne de montage à cinq étapes.

• Étape 1 : Lit les ingrédients bruts (les données de collision) et prépare les premières étapes.
• Étapes 2 et 3 : Transmettent le travail le long de la ligne, ajoutant davantage d'étapes au calcul.
• Étapes 4 et 5 : Terminent les calculs finaux et émettent la réponse.

2. Le « convoyeur » (Pipeline en cascade)

Ces cinq étapes sont reliées par un convoyeur dédié ultra-rapide appelé une interface en cascade.

• Imaginez une usine où les ouvriers ne s'arrêtent pas pour discuter ou attendre la permission de passer un carton à la personne suivante. Ils glissent simplement le carton dans un toboggan instantanément.
• Dans cette puce, les « cartons » sont des blocs de données appelés jetons.
• Les auteurs ont conçu un code de règles strict (un « contrat déterministe ») pour garantir que les ouvriers ne restent jamais bloqués en attendant les autres. Chaque ouvrier sait exactement quand passer un carton et quand en recevoir un, de sorte que la ligne ne se bloque jamais.

3. La « super-usine » (80 lignes simultanées)

La puce utilisée (la VCK190) est comme un immense entrepôt contenant 400 micro-ouvriers (appelés tuiles).

• Au lieu de construire une seule chaîne de montage, ils en ont construit 80 identiques côte à côte.
• Chaque ligne compte 5 ouvriers. $80 \text{ lignes} \times 5 \text{ ouvriers} = 400 \text{ ouvriers}$.
• Ils travaillent tous simultanément, résolvant 80 puzzles différents en même temps.

4. Les résultats : Vitesse et efficacité

Les auteurs ont testé cette « usine » contre deux autres méthodes : un processeur d'ordinateur standard (CPU) et une carte graphique haut de gamme (GPU).

• Vitesse : Leur usine à 80 lignes est 34 fois plus rapide qu'un seul cœur d'ordinateur standard.
  • Note : Une carte graphique haut de gamme (GPU) reste globalement plus rapide (environ 22 fois plus rapide que leur puce), mais le GPU est une machine beaucoup plus grande et plus coûteuse.
• Énergie : C'est ici que leur méthode brille. Grâce à l'efficacité et à la spécialisation de la chaîne de montage, elle consomme très peu d'électricité.
  • Pour résoudre un puzzle, leur puce utilise 7,7 fois moins d'énergie qu'un processeur d'ordinateur standard.
  • Elle est moins économe en énergie que le gigantesque GPU, mais le GPU consomme une quantité massive d'énergie pour y parvenir. La méthode de la puce représente un « juste milieu » pour les situations où vous avez besoin de vitesse mais ne pouvez pas brancher une machine massive gourmande en énergie.

5. Vérification de la précision

Ils ont veillé à ce que leur « chaîne de montage » ne fasse pas d'erreurs. Ils ont comparé les réponses de leur puce à un calcul de référence « gold standard » en double précision.

• Les résultats correspondaient presque parfaitement. La différence était si infime (environ 1 partie sur un million) qu'elle est considérée comme négligeable pour les calculs de physique qu'ils effectuent.

Résumé

En bref, les auteurs ont pris un calcul de physique complexe trop volumineux pour une seule puce informatique, l'ont découpé en cinq pièces gérables, et ont construit 80 chaînes de montage parallèles pour les résoudre toutes simultanément. Cette approche crée un « juste milieu » de haute vitesse et de faible consommation d'énergie, offrant une alternative puissante pour exécuter les simulations nécessaires à la compréhension de l'univers au Grand collisionneur de hadrons.

Résumé technique

1. Énoncé du problème

Les générateurs d'événements modernes de physique des hautes énergies (HEP), tels que MadGraph5_aMC@NLO (MG5aMC), font face à un goulot d'étranglement computationnel significatif dans l'évaluation de l'élément de matrice ($|M|^2$) pour les collisions de particules. Alors que le Grand collisionneur de hadrons (LHC) entre dans sa phase de haute luminosité, la demande pour ces calculs croît de manière non linéaire, tandis que la mise à l'échelle des CPU reste limitée.

• Le goulot d'étranglement : L'évaluation de l'élément de matrice représente 30 à 40 % du temps total de génération d'événements, en particulier pour les processus multi-jets impliquant des émissions réelles supplémentaires.
• Le défi : Les solutions existantes utilisant des GPU (par exemple, CUDACPP) offrent un débit élevé mais consomment une puissance significative. Les réseaux de portes programmables sur site (FPGA) offrent une efficacité énergétique mais peinent avec la contrainte de mémoire de programme (PM) de 16 ko par tuile sur les réseaux de moteurs d'IA modernes. Une implémentation monolithique de processus complexes (comme $gg \to t\bar{t}g$) dépasse cette limite de mémoire, empêchant le mappage direct sur une seule tuile.

2. Méthodologie

Les auteurs proposent une architecture de pipeline en cascade sur le réseau de moteurs d'IA (AIE) AMD Versal (spécifiquement la plateforme VCK190) pour surmonter les contraintes de mémoire et maximiser le parallélisme.

A. Plateforme cible et architecture

• Matériel : AMD Versal XCVC1902 ACAP comportant 400 tuiles de moteur d'IA disposées en une grille de $50 \times 8$, cadencées à 1,25 GHz.
• Décomposition du pipeline : Le processus $gg \to t\bar{t}g$ (impliquant 16 diagrammes de Feynman et 10 fonctions HELAS distinctes) est décomposé en un pipeline à cinq étapes.
  • Étape 1 : Génération de fonctions d'onde (spinors/vecteurs externes) et initialisation des jetons.
  • Étapes 2 et 3 : Évaluations des vertex fermion-vecteur (répartition de 12 diagrammes pour équilibrer la mémoire).
  • Étape 4 : Évaluations des vertex triple-gluon (incluant un générateur de boson hors couche différé).
  • Étape 5 : Termes de contact à quatre gluons et réduction de la matrice de couleur.
• Gestion de la mémoire : Pour tenir dans la limite de 16 ko, les auteurs ont employé le partitionnement de la mémoire de programme et l'évaluation différée. Par exemple, le générateur de boson hors couche ($FFV1P0\_3$) a été déplacé de l'étape 1 vers l'étape 4, réduisant l'utilisation de la mémoire de l'étape 1 de 17,8 ko à 15,5 ko.

B. Communication inter-tuile (Protocole de cascade)

• Mécanisme : Les étapes communiquent via l'interface de cascade unidirectionnelle de 384 bits (bande passante de 60 Go/s).
• Protocole de jeton : Un protocole déterministe, sans blocage, est utilisé où les étapes échangent des « jetons » structurés contenant des fonctions d'onde et des amplitudes partielles.
  • Jeton étendu (Étapes 1–4) : Transporte 5 fonctions d'onde externes, 3 propagateurs précalculés et 6 amplitudes de flux de couleur (18 battements/hélicité).
  • Jeton réduit (Étapes 4–5) : Transporte uniquement 5 fonctions d'onde et amplitudes (12 battements/hélicité) après l'évaluation locale du vertex triple-gluon.
• Déterminisme : Le système impose un « contrat de cascade » avec des structures de boucle identiques, des écritures inconditionnelles et des nombres de jetons correspondants statiquement pour assurer une opération sans surcharge et sans contrôle de flux.

C. Adaptations logicielles

• Portage de la bibliothèque HELAS : La bibliothèque HELAS standard de MG5aMC (à l'origine C++ scalaire double précision) a été portée sur les intrinsèques vectorielles des moteurs d'IA (virgule flottante simple précision float32).
• Optimisations :
  • Vectorisation : Les fonctions d'onde sont mappées sur des vecteurs SIMD de largeur 8.
  • Division complexe : La méthode de Smith (2 divisions) a été remplacée par une seule instruction de réciproque matérielle.
  • Mise en cache de l'hélicité : Précalcul de 10 fonctions d'onde pour 32 configurations d'hélicité, sélectionnées par recherche indexée par bits pour réduire les évaluations d'un facteur 16.
  • Réduction de couleur : Les divisions de normalisation de couleur ont été repliées dans des constantes de temps de compilation.

D. Déploiement système

• Échelle : 80 pipelines indépendants ont été mappés sur les 400 tuiles disponibles (5 tuiles par pipeline).
• Entrée/Sortie : Une architecture à commutation de paquets dans la logique programmable (PL) distribue les points de l'espace des phases vers les pipelines et collecte les résultats.

3. Contributions clés

1. Architecture de pipeline pilotée par la mémoire : Introduction d'un nouveau pipeline en cascade à 5 étapes qui partitionne avec succès un calcul d'élément de matrice multi-diagramme complexe sur plusieurs tuiles de moteur d'IA, surmontant la contrainte de PM de 16 ko.
2. Contrat de cascade déterministe : Développement d'un protocole de communication sans blocage utilisant des jetons de fonction d'onde et des structures de boucle statiques, éliminant le besoin de matériel de contrôle de flux complexe.
3. Portage complet de HELAS : Portage réussi de la bibliothèque complète d'amplitudes HELAS vers les intrinsèques vectorielles des moteurs d'IA, intégrant des optimisations complexes telles que la mise en cache binaire indexée de l'hélicité et la division complexe réduite.
4. Déploiement évolutif : Démonstration d'un déploiement théorique de 80 pipelines utilisant 100 % des ressources de calcul des moteurs d'IA du VCK190.

4. Résultats

• Débit : Le débit projeté pour le réseau de 80 pipelines est de $1,0 \times 10^6$ évaluations d'élément de matrice par seconde (ME/s).
  • Cela représente une accélération de 34× par rapport à un cœur de CPU monothread (Intel i5-10600).
  • Bien que inférieur à un GPU NVIDIA A100 ($2,18 \times 10^7$ ME/s), la solution des moteurs d'IA est nettement plus économe en énergie.
• Efficacité énergétique :
  • Moteur d'IA : 54,8 µJ/ME (à une puissance de domaine AIE de 54,8 W).
  • CPU : 422 µJ/ME.
  • GPU : 7,3 µJ/ME (mais à une puissance de 159 W).
  • Amélioration : Le moteur d'IA offre une amélioration de 7,7× de l'efficacité énergétique par rapport à la référence CPU.
• Précision : Validée par rapport à la référence double précision de MG5aMC.
  • Erreur relative moyenne : 1,43 ppm (parties par million).
  • Erreur relative maximale : 168 ppm.
  • Ce niveau de précision est jugé suffisant pour les calculs d'ordre dominant (LO) où les incertitudes physiques (variation d'échelle, PDF) dominent les erreurs numériques.
• Utilisation des ressources :
  • Mémoire de programme : L'étape 1 est le goulot d'étranglement avec une utilisation de 94,7 % (15 514 octets).
  • Logique programmable : Utilisation modeste (4,72 % de LUT, 2,87 % de registres), laissant de la place pour une logique supplémentaire.

5. Importance et travaux futurs

• Importance : Ce travail prouve que les réseaux de moteurs d'IA sont viables pour la génération d'événements HEP à haut débit et économe en énergie, en particulier pour les environnements contraints en puissance (par exemple, systèmes de déclenchement en ligne ou informatique de périphérie au LHC) où les enveloppes de puissance des GPU sont prohibitives. Il établit une méthodologie systématique pour partitionner des noyaux physiques complexes sur des tuiles à mémoire contrainte.
• Limites : L'implémentation actuelle est limitée aux processus d'ordre dominant (LO). Les chiffres de latence sont basés sur des simulateurs approximatifs en cycles plutôt que sur un timing matériel direct de l'ensemble du réseau.
• Directions futures :
  • Filtrage de l'hélicité : Précalcul des masques d'hélicité valides pour réduire les itérations de la boucle interne, potentiellement doublant le débit.
  • Multiplicité plus élevée : Extension de la profondeur du pipeline pour des processus plus complexes (par exemple, $t\bar{t}ggg$).
  • Intégration NLO : Adaptation de l'architecture pour les calculs de Next-to-Leading-Order impliquant des intégrales de boucle.
  • Évolution matérielle : Exploitation des dispositifs Versal de nouvelle génération avec des réseaux plus grands ou des fréquences d'horloge plus élevées.

En conclusion, l'article présente une alternative robuste et économe en énergie à l'accélération par GPU pour des charges de travail HEP spécifiques, exploitant les capacités de cascade uniques du moteur d'IA AMD Versal pour résoudre les défis de partitionnement de mémoire inhérents aux calculs complexes d'éléments de matrice.