Resumo Técnico: Axon: Um Superotimizador de Síntese para Programas de Tensores

Definição do Problema

Escrever kernels de alto desempenho para aceleradores de IA (como o Trainium da Amazon) exige um conhecimento profundo de detalhes específicos do hardware, incluindo hierarquias de memória, estratégias de tiling, seleção de instruções, layout de dados e fusão de operadores. Embora linguagens de kernel baseadas em tiles como a NKI (Neuron Kernel Interface) abstraiam alguns detalhes de baixo nível, alcançar a utilização máxima do hardware continua sendo um desafio para a maioria dos engenheiros de ML. As abordagens de otimização existentes sofrem de várias limitações:

• Regras manuais: Muitos sistemas dependem de regras de reescrita escritas manualmente para transformações algébricas, que são propensas a erros e difíceis de escalar para novas combinações de operadores.
• Geração baseada em templates: Algumas abordagens utilizam templates escritos à mão para geração de código, o que limita a flexibilidade.
• Escopo limitado: Muitos otimizadores visam apenas CPUs ou GPUs, falhando em abordar as arquiteturas únicas de múltiplos motores e orientadas a DMA dos aceleradores de IA modernos.
• Falta de correção comprovável: Testes probabilísticos ou reordenação baseada em heurísticas podem perder configurações ideais ou introduzir erros semânticos.

Metodologia

O Axon é um superotimizador de síntese projetado para programas de aceleradores de IA baseados em tiles. Ele aceita programas de tensores de alto nível (escritos em uma interface semelhante ao NumPy) e gera automaticamente instruções de destino otimizadas (NKI) através de um pipeline de múltiplos estágios. A inovação central é o $\nu$Graph ("nu-graph"), uma representação interna persistente que rastreia todas as variantes de programa semanticamente equivalentes simultaneamente durante o processo de síntese, adiando as escolhas de otimização para a seleção empírica.

O pipeline consiste nos seguintes estágios:

1. Transformações Algébricas via Propagação de Operadores:

   • O Axon descobre grafos de computação semanticamente equivalentes propagando operadores através do grafo de computação.
   • Ele tenta reordenar operadores adjacentes (por exemplo, mover uma multiplicação elemento a elemento através de uma multiplicação de matriz) para expor paralelismo e oportunidades de fusão.
   • Garantia de Correção: Em vez de usar regras de reescrita manuais, o Axon decompõe operações de tensores em computações granulares e usa um solver SMT (Satisfiability Modulo Theories) sobre tensores ilimitados para verificar de forma provável que a reordenação preserva a semântica. Isso lida com formas de tensores arbitrárias e não requer definições de propriedades manuais.

2. Tiling Simbólico:

   • O sistema realiza o tiling de cada operação de tensor usando parâmetros simbólicos para dimensões de strip, block e tile.
   • Esses parâmetros são restringidos por descrições de hardware (ex: capacidades de buffer on-chip e limites específicos de instrução).
   • O Axon gera múltiplas versões com diferentes configurações de tiling, permitindo a otimização independente de ramos paralelos (ex: separando ramos de RMSNorm e MatMul para rodar em diferentes motores de computação).

3. Síntese de ISA (Seleção de Instrução):

   • Usando uma abordagem de síntese de programa baseada em sketch, o Axon reduz operações de tensores de alto nível para instruções ISA de destino (ex: nc_matmul e dma_copy do Trainium).
   • Ele explora sequências de instruções válidas e transformações de layout de dados (como transposições) necessárias para instruções específicas de hardware.
   • A equivalência entre a operação de alto nível e a sequência de instruções sintetizada é novamente verificada via SMT.

4. Fusão de ISA e Operadores:

   • Uma segunda rodada de transformações algébricas é realizada no grafo de ISA sintetizado para expor novas oportunidades de fusão criadas pelo processo de síntese (ex: fundir uma transposição com uma multiplicação de matriz).
   • Fusão de Nós (Node Fusion): Mescla operações consecutivas com dimensões de tiling compatíveis para manter os resultados intermediários em buffers on-chip (SBUF/PSUM), reduzendo o tráfego de HBM.
   • Fusão de Instruções (Instruction Fusion): Substitui sequências de instruções ISA por variantes fundidas (ex: dma_copy + nc_transpose $\rightarrow$ dma_transpose).

5. Emissão de Código e Seleção Empírica:

   • O Axon instancia parâmetros simbólicos com valores concretos para emitir programas NKI concretos.
   • Cada candidato é compilado pelo compilador Neuron e executado no hardware alvo.
   • O kernel de melhor desempenho é selecionado com base em medições empíricas de latência e validação de correção numérica contra uma implementação de referência.

Principais Contribuições

O artigo apresenta três contribuições primárias:

1. Transformações Algébricas Prováveis: Uma metodologia para descobrir automaticamente grafos de computação semanticamente equivalentes ao propagar operadores e verificar a equivalência usando SMT sobre tensores ilimitados. Isso elimina a necessidade de regras de reescrita manuais.
2. Pipeline Baseado em Síntese: Um pipeline completo para aceleradores de IA baseados em tiles que explora conjuntamente transformações algébricas, seleção de instruções, configurações de tiling e fusão. Ele mantém todas as variantes válidas em um $\nu$Graph e seleciona o kernel ideal empiricamente.
3. Avaliação Abrangente: Uma avaliação no Trainium da Amazon em 20 benchmarks, incluindo operadores individuais e kernels complexos de múltiplos operadores de LLM (ex: Group Query Attention, Gated MLP).

Resultados Experimentais

O Axon foi avaliado contra três baselines: o compilador Neuron da Amazon, kernels NKI otimizados manualmente e o Mirage (um compilador de busca de última geração).

• Vs. Compilador Neuron:

  • Operadores Individuais: O Axon alcança acelerações de até 3.7x (ex: SiLU em grandes dimensões) com uma aceleração de média geométrica de 1.23x–1.90x em vários operadores.
  • Kernels de Múltiplos Operadores: O Axon alcança acelerações de até 19x em kernels como Transpose+MatMul e até 1.47x em RMSNorm+MatMul. As acelerações de média geométrica variam de 1.07x a 1.32x.
  • Raciocínio: Os ganhos são atribuídos a melhores estratégias de tiling, utilização eficiente de DMA e a capacidade de fundir operações entre diferentes motores de computação (ex: executando RMSNorm em motores Scalar/Vector enquanto o MatMul roda no Tensor Engine em paralelo).

• Vs. NKI Otimizado Manualmente:

  • O Axon iguala ou supera kernels escritos por especialistas em até 1.35x em configurações específicas (ex: Cumsum e RoPE), demonstrando que a exploração automatizada pode rivalizar a sintonia manual sem exigir profundo conhecimento de hardware.

• Vs. Mirage:

  • Nos benchmarks suportados por ambos, o Axon alcança acelerações de média geométrica de 2–10%.
  • Crucialmente, o Axon sintetiza com sucesso kernels para 16 benchmarks adicionais que o Mirage não consegue lidar devido à sua dependência de templates manuais.

• Overhead de Compilação:

  • O espaço de busca pode gerar cerca de ~30.000 candidatos para kernels complexos.
  • A Fase 1 (síntese simbólica) é rápida (<82 segundos). A Fase 2 (compilação e execução) domina o tempo (até ~1.7 horas para camadas complexas de LLM). O artigo observa que este é um custo único, comparável a outros compiladores baseados em schedule, e os kernels otimizados podem ser armazenados em cache para reutilização.

Significância e Alegações

O artigo afirma que o Axon representa uma mudança da otimização baseada em heurísticas ou templates para uma abordagem de síntese baseada em semântica. Ao utilizar solvers SMT para garantir a equivalência semântica, o Axon pode descobrir automaticamente transformações algébricas complexas (como a reordenação de operações entre diferentes motores de computação) que são difíceis de codificar manualmente.

A significância reside em:

• Automação: Reduzindo o fardo dos programadores para ajustar manualmente kernels para hardwares específicos.
• Correção: Fornecendo garantias prováveis para transformações, ao contrário de métodos de teste probabilísticos.
• Desempenho: Alcançando acelerações significativas em aceleradores de IA modernos ao otimizar conjuntamente tiling, seleção de instruções e fusão em um espaço de busca unificado.
• Generalidade: A capacidade de lidar com kernels complexos de múltiplos operadores de LLM que compiladores de busca existentes (como o Mirage) não conseguem sintetizar.

Os autores enfatizam que, embora o overhead de compilação não seja trivial, este é um custo único padrão para compiladores de otimização, e os ganhos de desempenho resultantes justificam o investimento, particularmente para cargas de trabalho de LLM em larga escala.