Resumo Técnico: Emulação de Componente-Produto BF16 para GEMM FP32 e FP64 em Intel AMX

1. Definição do Problema

As CPUs modernas integram cada vez mais motores de matriz de alto rendimento (como o Intel Advanced Matrix Extensions, ou AMX) otimizados para cargas de trabalho de IA de baixa precisão (ex: BF16, FP8). No entanto, muitas aplicações de computação científica em campos como dinâmica de fluidos computacional e química quântica ainda dependem de GEMM (Multiplicação de Matrizes Gerais) FP32 e FP64 para estabilidade numérica e reprodutibilidade. Isso cria um descompasso de hardware: os motores de baixa precisão de alto desempenho não são nativamente utilizáveis para tarefas científicas de alta precisão.

A questão central abordada por este artigo é se o rendimento de matriz de baixa precisão projetado para IA pode ser algoritmicamente reaproveitado para acelerar o GEMM de alta precisão em CPUs. O desafio reside em emular aritmética FP32 e FP64 usando componentes BF16 sem incorrer em custos excessivos de decomposição, movimentação de dados e reconstrução, mantendo simultaneamente uma precisão numérica comparável às bibliotecas padrão como a Intel oneMKL.

2. Metodologia

O artigo propõe um framework de emulação de componente-produto que decompõe operandos de alta precisão em múltiplas matrizes BF16, computa produtos de baixa precisão selecionados usando AMX e reconstrói o resultado no formato alvo.

2.1 Formatos de Ponto Flutuante e Domínios

• BF16: Mantém o intervalo de expoente de 8 bits do FP32, mas reduz a precisão da mantissa para 8 bits.
• Emulação FP32: Opera sobre operandos FP32 padrão.
• Emulação FP64: Opera em um domínio restrito ($FB_{64}$), onde os valores FP64 residem dentro do intervalo de expoente normal do BF16. Esta restrição evita a necessidade de fatores de escala por painel, preservando a margem de desempenho.

2.2 Estratégias de Decomposição

• Caminho AMX-FP32: Utiliza uma decomposição de resíduo de três componentes BF16. Cada operando FP32 é dividido em três fatias BF16 ($A_0, A_1, A_2$) usando arredondamento de resíduo direto. Isso segue a abordagem "estilo Henry", que explora o intervalo de expoente compartilhado para garantir a subtração exata de resíduos em FP32.
• Caminho AMX-FP64: Utiliza uma decomposição Ozaki simplificada e fixa de seis fatias. Cada operando FP64 é dividido em seis fatias BF16. Ao contrário dos esquemas Ozaki completos que usam escala dinâmica, este método armazena os pedaços extraídos diretamente como matrizes BF16 (assumindo que permaneçam normais), eliminando a necessidade de metadados de escala e pesos de reconstrução.

2.3 Escalonamento de Produto e Reconstrução

• Reconstrução FP32: O método avalia seis produtos de componentes selecionados (as três primeiras diagonais da matriz de interação de componentes). Crucialmente, esses produtos são acumulados diretamente em tiles AMX de FP32. A implementação emprega um escalonamento de reuso de operandos (ex: $A_1B_1 \to A_1B_0 \to \dots$) que mantém os tiles acumuladores residentes, minimizando o tráfego de carga/armazenamento de tiles.
• Reconstrução FP64: O método avalia um número variável de produtos (6, 10, 15 ou 21). Diferente do caminho FP32, cada produto BF16 selecionado é computado em FP32, armazenado na memória, alargado para FP64 e acumulado em um acumulador de bloco FP64. Isso evita que erros de acumulação FP32 contaminem o resultado de alta precisão, mas incorre em um overhead significativo devido às repetidas operações de armazenamento e conversões de formato.

2.4 Detalhes de Implementação

A implementação utiliza uma estrutura de GEMM bloqueada com blocos de saída de $256 \times 256$ e micro-tiles de $32 \times 32$. Ela depende de buffers de componentes pré-computados e contíguos em bloco para otimizar os padrões de acesso à memória. Os operandos A são armazenados em layout row-major padrão, enquanto os operandos B são pré-empacotados no layout VNNI exigido pelas instruções de produto escalar BF16 do AMX.

3. Principais Contribuições

1. Ponte Algorítmica: Demonstra um método prático para emular GEMM FP32 e FP64 usando instruções AMX BF16, unindo a lacuna entre o hardware otimizado para IA e os requisitos de computação científica.
2. Emulação FP32 com Precisão de Nível FP32: Mostra que um escalonamento de seis produtos BF16 com acumulação FP32 pode atingir uma precisão comparável ao SGEMM da oneMKL (medida por erro relativo normativo e erro por componente escalonado pelo GEMM), sem reivindicar identidade bit a bit.
3. Emulação FP64 Controlada: Introduz um tradeoff de precisão-desempenho ajustável para a emulação FP64, variando o número de produtos de componentes retidos (6 a 21) dentro de uma decomposição de seis fatias fixa.
4. Movimentação de Dados Otimizada: Desenvolve estratégias de escalonamento específicas (reuso de operandos para FP32) e layouts de dados (empacotamento VNNI, buffers pré-computados) para mitigar o overhead de decomposição e reconstrução em motores de matriz de CPU.

4. Resultados Experimentais

Experimentos foram conduzidos em um servidor Intel Xeon Platinum 8462Y+ de socket duplo, comparando os kernels propostos contra o SGEMM e DGEMM da Intel oneMKL.

• Desempenho FP32: A implementação AMX-FP32 supera o oneMKL SGEMM em todos os tamanhos de matriz quadrada testados ($N=256$ a $N=32768$). Os ganhos de velocidade variam de 1,14$\times$ a 2,56$\times$, com os maiores ganhos observados em matrizes maiores, onde a vantagem de rendimento do BF16 amortiza o custo de decomposição.
• Desempenho FP64: Os resultados são mais limitados.
  • AMX-FP64-6 (6 produtos) alcança ganhos de velocidade de até 1,72$\times$ sobre o DGEMM da oneMKL para matrizes grandes ($N \ge 16384$).
  • AMX-FP64-10 fornece ganhos marginais apenas para as maiores matrizes.
  • AMX-FP64-15 e AMX-FP64-21 não superam o DGEMM nas configurações testadas, pois os custos de armazenar, converter e acumular produtos adicionais consomem a margem de desempenho disponível.
• Precisão:
  • O AMX-FP32 atinge erros na ordem de $10^{-7}$, consistentes com a precisão de nível FP32.
  • As variantes AMX-FP64 mostram diminuição do erro conforme o número de produtos aumenta, com a variante de 21 produtos aproximando-se da precisão do DGEMM, mas com um custo de desempenho que anula o ganho de velocidade.

5. Significância e Alegações

O artigo posiciona este trabalho como um estudo algorítmico inicial, em vez de um substituto completo para bibliotecas de alta precisão nativas. Sua significância reside em demonstrar que:

• Motores de matriz de baixa precisão em CPU podem ser efetivamente utilizados para computação científica de alta precisão através de cuidadosa decomposição algorítmica e otimização de movimentação de dados.
• Existe um teto de desempenho para a emulação: enquanto a emulação FP32 é altamente eficaz devido à acumulação eficiente em FP32, a emulação FP64 é limitada pelo overhead de alargamento e acumulação de resultados fora dos registros de tile AMX.
• O método oferece um tradeoff ajustável para FP64, permitindo que aplicações selecionem um número de produtos que equilibre seus requisitos específicos de precisão contra as necessidades de desempenho.

Os autores explicitamente notam limitações: o protótipo atual suporta apenas matrizes quadradas, requer operandos dentro do intervalo de expoente BF16 (para FP64) e não trata overflow/underflow ou valores excepcionais. O trabalho futuro é identificado como a extensão do framework para matrizes retangulares, domínios de entrada mais amplos via decomposição escalonada e adaptação para outros motores de matriz de CPU (ex: Arm SME).