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

Este artigo apresenta uma arquitetura de pipeline em cascata de cinco estágios implementada em matrizes AI Engine da AMD Versal para avaliar eficientemente elementos de matriz de ordem dominante para o processo $\gamma\gamma \to t\bar{t}g$, alcançando uma taxa de processamento projetada de $1.0\times10^6$ avaliações por segundo com um aumento de velocidade de $34\times$ e uma melhoria de eficiência energética de $7.7\times$ em relação a um único núcleo de CPU, mantendo precisão numérica na ordem de partes por milhão.

O essencial

Imagine que você está tentando prever o resultado de uma colisão massiva e caótica entre duas partículas minúsculas (como prótons) dentro de um acelerador de partículas gigante. Para fazer isso, os físicos usam uma receita matemática complexa chamada "elemento de matriz". Calcular essa receita é como resolver um quebra-cabeça gigante e multifásico. O problema é que, para obter uma resposta confiável, eles precisam resolver esse mesmo quebra-cabeça milhões de vezes, cada vez com condições iniciais ligeiramente diferentes.

Atualmente, fazer isso em processadores de computador padrão (CPUs) é como tentar resolver esses quebra-cabeças um por um com uma única pessoa. É preciso, mas incrivelmente lento e consome muita energia, especialmente à medida que o acelerador de partículas se torna mais poderoso.

Este artigo apresenta uma nova maneira de resolver esses quebra-cabeças usando um tipo especial de chip de computador chamado AMD Versal AI Engine. Em vez de ter uma pessoa resolvendo o quebra-cabeça inteiro, os autores construíram uma linha de montagem de fábrica diretamente dentro do chip.

Veja como a solução deles funciona, dividida em conceitos simples:

1. O Problema da "Linha de Montagem"

A receita matemática para essa colisão de partículas específica (dois glúons se transformando em um quark top, um antiquark top e outro glúon) é grande demais para caber na memória de um único processador minúsculo no chip. Pense nisso como tentar encaixar um manual de instruções de 38 páginas em um bolso que só pode conter 16 páginas.

A Solução: Os autores dividiram o manual em cinco capítulos. Eles criaram uma linha de montagem de cinco estágios.

• Estágio 1: Lê os ingredientes brutos (os dados da colisão) e prepara os primeiros passos.
• Estágios 2 e 3: Passam o trabalho ao longo da linha, adicionando mais etapas ao cálculo.
• Estágios 4 e 5: Finalizam os cálculos finais e lançam a resposta.

2. A "Esteira Rolante" (Pipeline em Cascata)

Esses cinco estágios são conectados por uma esteira rolante super-rápida e dedicada chamada interface em cascata.

• Imagine uma fábrica onde os trabalhadores não param para conversar ou esperar permissão para passar uma caixa para a próxima pessoa. Eles apenas deslizam a caixa por um tobogã instantaneamente.
• Neste chip, as "caixas" são pedaços de dados chamados tokens.
• Os autores projetaram um livro de regras estrito (um "contrato determinístico") para garantir que os trabalhadores nunca fiquem presos esperando uns pelos outros. Cada trabalhador sabe exatamente quando passar uma caixa e quando receber uma, para que a linha nunca trave.

3. A "Superfábrica" (80 Linhas ao Mesmo Tempo)

O chip que eles usaram (o VCK190) é como um armazém massivo contendo 400 trabalhadores minúsculos (chamados de tiles).

• Em vez de construir apenas uma linha de montagem, eles construíram 80 linhas de montagem idênticas lado a lado.
• Cada linha tem 5 trabalhadores. $80 \text{ linhas} \times 5 \text{ trabalhadores} = 400 \text{ trabalhadores}$.
• Todos estão trabalhando ao mesmo tempo, resolvendo 80 quebra-cabeças diferentes simultaneamente.

4. Os Resultados: Velocidade e Eficiência

Os autores testaram essa "fábrica" contra dois outros métodos: um processador de computador padrão (CPU) e uma placa de vídeo de alto desempenho (GPU).

• Velocidade: Sua fábrica de 80 linhas é 34 vezes mais rápida que um único núcleo de computador padrão.
  • Nota: Uma placa de vídeo de ponta (GPU) ainda é mais rápida no geral (cerca de 22 vezes mais rápida que seu chip), mas a GPU é uma máquina muito maior e mais cara.
• Energia: É aqui que o método deles brilha. Como a linha de montagem é tão eficiente e especializada, ela usa muito pouca eletricidade.
  • Para resolver um quebra-cabeça, seu chip usa 7,7 vezes menos energia que um processador de computador padrão.
  • É menos eficiente em termos energéticos que a GPU gigante, mas a GPU consome uma quantidade massiva de energia para fazer isso. O método do chip é um "ponto ideal" para situações onde você precisa de velocidade, mas não pode conectar uma máquina enorme e faminta de energia.

5. Verificação de Precisão

Eles garantiram que sua "linha de montagem" não cometesse erros. Eles compararam as respostas de seu chip com um cálculo de precisão dupla que serve como "padrão ouro".

• Os resultados corresponderam quase perfeitamente. A diferença foi tão pequena (cerca de 1 parte em um milhão) que é considerada negligenciável para os cálculos de física que eles estão fazendo.

Resumo

Em resumo, os autores pegaram um cálculo de física complexo que era grande demais para um único chip de computador, cortaram-no em cinco partes gerenciáveis e construíram 80 linhas de montagem paralelas para resolvê-las todas ao mesmo tempo. Essa abordagem cria um "ponto ideal" de alta velocidade e baixo consumo de energia, oferecendo uma alternativa poderosa para executar as simulações necessárias para entender o universo no Grande Colisor de Hádrons.

Resumo técnico

1. Declaração do Problema

Os geradores de eventos modernos de Física de Altas Energias (HEP), como o MadGraph5_aMC@NLO (MG5aMC), enfrentam um gargalo computacional significativo na avaliação do elemento de matriz ($|M|^2$) para colisões de partículas. À medida que o Grande Colisor de Hádrons (LHC) entra em sua fase de Alta Luminosidade, a demanda por esses cálculos cresce de forma não linear, enquanto a escalabilidade da CPU permanece limitada.

• O Gargalo: A avaliação do elemento de matriz representa 30–40% do tempo total de geração de eventos, particularmente para processos de multi-jatos envolvendo emissões reais adicionais.
• O Desafio: Soluções existentes que utilizam GPUs (por exemplo, CUDACPP) oferecem alta taxa de transferência, mas consomem energia significativa. As Redes de Portas Programáveis em Campo (FPGAs) oferecem eficiência energética, mas lutam com a restrição de 16 kB de memória de programa (PM) por tile em arrays modernos de AI Engine. Uma implementação monolítica de processos complexos (como $gg \to t\bar{t}g$) excede esse limite de memória, impedindo o mapeamento direto para um único tile.

2. Metodologia

Os autores propõem uma arquitetura de pipeline em cascata no array AMD Versal AI Engine (AIE) (especificamente a plataforma VCK190) para superar as restrições de memória e maximizar o paralelismo.

A. Plataforma-Alvo e Arquitetura

• Hardware: AMD Versal XCVC1902 ACAP com 400 tiles de AI Engine dispostos em uma grade de $50 \times 8$, operando a 1,25 GHz.
• Decomposição do Pipeline: O processo $gg \to t\bar{t}g$ (envolvendo 16 diagramas de Feynman e 10 funções HELAS distintas) é decomposto em um pipeline de cinco estágios.
  • Estágio 1: Geração da função de onda (espinores/vetores externos) e inicialização de tokens.
  • Estágios 2 e 3: Avaliações de vértice férmion-vetor (dividindo 12 diagramas para equilibrar a memória).
  • Estágio 4: Avaliações de vértice de três glúons (incluindo um gerador de bóson fora da camada de massa adiado).
  • Estágio 5: Termos de contato de quatro glúons e redução da matriz de cor.
• Gerenciamento de Memória: Para caber no limite de 16 kB, os autores empregaram particionamento da memória de programa e avaliação adiada. Por exemplo, o gerador de bóson fora da camada de massa ($FFV1P0\_3$) foi movido do Estágio 1 para o Estágio 4, reduzindo o uso de memória do Estágio 1 de 17,8 kB para 15,5 kB.

B. Comunicação Inter-Tile (Protocolo de Cascata)

• Mecanismo: Os estágios comunicam-se através da interface unidirecional de cascata de 384 bits (largura de banda de 60 GB/s).
• Protocolo de Token: Um protocolo determinístico, livre de deadlocks, é usado onde os estágios trocam "tokens" estruturados contendo funções de onda e amplitudes parciais.
  • Token Estendido (Estágios 1–4): Carrega 5 funções de onda externas, 3 propagadores pré-calculados e 6 amplitudes de fluxo de cor (18 batidas/helicidade).
  • Token Reduzido (Estágios 4–5): Carrega apenas 5 funções de onda e amplitudes (12 batidas/helicidade) após a avaliação local do vértice de três glúons.
• Determinismo: O sistema impõe um "contrato de cascata" com estruturas de loop idênticas, escritas incondicionais e contagens de tokens correspondidas estaticamente para garantir operação sem sobrecarga e sem controle de fluxo.

C. Adaptações de Software

• Portabilidade da Biblioteca HELAS: A biblioteca HELAS padrão do MG5aMC (originalmente C++ de precisão dupla escalar) foi portada para intrínsecos vetoriais do AI Engine (precisão simples float32).
• Otimizações:
  • Vetorização: Funções de onda mapeadas para vetores SIMD de 8 vias.
  • Divisão Complexa: Substituição do método de Smith (2 divisões) por uma única instrução de recíproco de hardware.
  • Cache de Helicidade: Pré-cálculo de 10 funções de onda para 32 configurações de helicidade, selecionando-as via busca indexada por bits para reduzir as avaliações em 16×.
  • Redução de Cor: Dobramento das divisões de normalização de cor em constantes de tempo de compilação.

D. Implantação do Sistema

• Escala: 80 pipelines independentes foram mapeados nos 400 tiles disponíveis (5 tiles por pipeline).
• E/S: Uma arquitetura comutada por pacotes na Lógica Programável (PL) distribui pontos do espaço de fase para os pipelines e coleta os resultados.

3. Principais Contribuições

1. Arquitetura de Pipeline Orientada a Memória: Introdução de um novo pipeline de cascata de 5 estágios que partitiona com sucesso um cálculo complexo de elemento de matriz multi-diagrama através de múltiplos tiles de AI Engine, superando a restrição de 16 kB de PM.
2. Contrato de Cascata Determinístico: Desenvolvimento de um protocolo de comunicação livre de deadlocks usando tokens de função de onda e estruturas de loop estáticas, eliminando a necessidade de hardware complexo de controle de fluxo.
3. Porta Completa do HELAS: Portagem bem-sucedida da biblioteca completa de amplitudes HELAS para intrínsecos vetoriais do AI Engine, incorporando otimizações complexas como cache de helicidade indexado binariamente e redução de divisão complexa.
4. Implantação Escalável: Demonstração de uma implantação teórica de 80 pipelines utilizando 100% dos recursos de computação do AI Engine do VCK190.

4. Resultados

• Taxa de Transferência: A taxa de transferência projetada para o array de 80 pipelines é de $1,0 \times 10^6$ avaliações de elemento de matriz por segundo (ME/s).
  • Isso representa um acréscimo de velocidade de 34× sobre um único núcleo de CPU de thread única (Intel i5-10600).
  • Embora inferior a uma GPU NVIDIA A100 ($2,18 \times 10^7$ ME/s), a solução AI Engine é significativamente mais eficiente energeticamente.
• Eficiência Energética:
  • AI Engine: 54,8 µJ/ME (com 54,8 W de potência no domínio AIE).
  • CPU: 422 µJ/ME.
  • GPU: 7,3 µJ/ME (mas com 159 W de potência).
  • Melhoria: O AI Engine oferece uma melhoria de 7,7× na eficiência energética em comparação com a linha de base da CPU.
• Precisão: Validado contra a referência de precisão dupla do MG5aMC.
  • Erro relativo médio: 1,43 ppm (partes por milhão).
  • Erro relativo máximo: 168 ppm.
  • Este nível de precisão é considerado suficiente para cálculos de Ordem Principal (LO), onde as incertezas físicas (variação de escala, PDFs) dominam os erros numéricos.
• Utilização de Recursos:
  • Memória de Programa: O Estágio 1 é o gargalo com 94,7% de utilização (15.514 bytes).
  • Lógica Programável: Uso modesto (4,72% LUTs, 2,87% Registradores), deixando espaço para lógica adicional.

5. Significado e Trabalho Futuro

• Significado: Este trabalho prova que arrays de AI Engine são viáveis para geração de eventos HEP de alta taxa de transferência e eficiência energética, particularmente para ambientes com restrições de energia (por exemplo, sistemas de gatilho online ou computação de borda no LHC) onde os envelopes de potência da GPU são proibitivos. Estabelece uma metodologia sistemática para particionar kernels de física complexos em tiles de memória restritos.
• Limitações: A implementação atual é limitada a processos de Ordem Principal (LO). As cifras de latência baseiam-se em simuladores aproximados de ciclos, em vez de temporização direta de hardware do array completo.
• Direções Futuras:
  • Filtragem de Helicidade: Pré-cálculo de máscaras de helicidade válidas para reduzir as iterações do loop interno, potencialmente dobrando a taxa de transferência.
  • Maior Multiplicidade: Estender a profundidade do pipeline para processos mais complexos (por exemplo, $t\bar{t}ggg$).
  • Integração NLO: Adaptar a arquitetura para cálculos de Próxima Ordem Principal (NLO) envolvendo integrais de loop.
  • Evolução de Hardware: Aproveitar dispositivos Versal de próxima geração com arrays maiores ou velocidades de clock mais altas.

Em conclusão, o artigo apresenta uma alternativa robusta e eficiente energeticamente à aceleração por GPU para cargas de trabalho HEP específicas, aproveitando as capacidades únicas de cascata do AI Engine da AMD Versal para resolver os desafios de particionamento de memória inerentes a cálculos complexos de elementos de matriz.