LP-GEMM: Integrating Layout Propagation into GEMM Operations
Este artigo apresenta o LP-GEMM, uma técnica que integra a propagação de layout de dados entre operações sequenciais de multiplicação de matrizes (GEMM) para eliminar o reempacotamento redundante, resultando em acelerações significativas em workloads de aprendizado de máquina em arquiteturas x86 e RISC-V.
Autores originais:César Guedes Carneiro, Lucas Alvarenga, Guido Araujo, Sandro Rigo
Imagine que você é um chef de cozinha de elite (o computador) e sua tarefa é preparar milhões de pratos complexos (cálculos matemáticos) para um restaurante lotado (Inteligência Artificial).
O problema que os cientistas deste artigo encontraram é que, na cozinha atual, cada prato passa por um processo desnecessário e repetitivo:
O Problema (A Cozinha Bagunçada): Atualmente, quando você precisa fazer uma sequência de pratos (como na Inteligência Artificial, onde um cálculo leva ao próximo), o chef pega os ingredientes crus, organiza-os em bandejas específicas (isso se chama "empacotamento" ou packing), cozinha o prato, e depois desmonta tudo para colocar no prato final.
O pior é que, para o próximo prato, ele pega esse prato pronto, desmonta tudo de novo, organiza os ingredientes em novas bandejas, cozinha e desmonta novamente. É como se você tivesse que tirar a comida do prato, colocar na mesa, organizar em caixas, e só então começar a cozinhar o próximo. Isso gasta muito tempo e energia apenas movendo coisas, sem cozinhar nada.
A Solução (LP-GEMM): Os autores criaram uma nova técnica chamada LP-GEMM. A ideia é simples: não desmonte o prato entre as etapas.
Eles dividiram o processo de cozinhar em três etapas inteligentes:
O Chef Inicial (Ini-GEMM): Ele pega os ingredientes crus, organiza nas bandejas certas e cozinha o primeiro prato. Mas, em vez de desmontar tudo, ele deixa o prato pronto em um formato especial que já está pronto para ser usado no próximo passo.
O Chef do Meio (Mid-GEMM): Ele pega o prato que já está organizado (sem precisar desmontar nem reorganizar nada) e continua a cozinha. Como ele já sabe como os ingredientes estão organizados, ele pula a etapa chata de "arrumar a mesa" e vai direto para o fogo.
O Chef Final (End-GEMM): Só no último prato é que ele organiza tudo de volta no formato padrão para servir ao cliente.
A Analogia da Fita de Montagem: Pense em uma linha de montagem de carros.
Método Antigo: A cada peça que chega, você tira da caixa, monta na carroceria, tira da carroceria, coloca em outra caixa, e só então monta a próxima peça.
Método LP-GEMM: Você deixa a peça montada na esteira. A próxima máquina pega a peça exatamente como ela está, sem precisar tirar da caixa ou mudar de lugar. Só no final, quando o carro está pronto, você o coloca na caixa de entrega.
Os Resultados (A Velocidade): Ao testar isso em computadores modernos (como os chips da Intel e os novos chips RISC-V), eles descobriram que:
Em tarefas de Inteligência Artificial (como fazer o modelo de linguagem Llama 3.2 funcionar), a nova técnica ficou até 2 vezes mais rápida em computadores comuns e até 5 vezes mais rápida em computadores menores (como os usados em dispositivos móveis ou borda).
Eles conseguiram fazer isso sem precisar de hardware novo, apenas mudando a "receita" de como os dados são movidos.
Por que isso importa? Hoje, a Inteligência Artificial consome muita energia e tempo porque gasta uma parte enorme apenas "arrumando a mesa" (movendo dados) antes de realmente "pensar" (fazer o cálculo). O LP-GEMM mostra que, se a gente parar de fazer essa arrumação repetida, podemos fazer os computadores muito mais rápidos e eficientes, gastando menos bateria e energia.
Resumo em uma frase: O LP-GEMM é como descobrir que, em vez de desmontar e remontar a sua mala a cada parada de viagem, você pode deixá-la pronta e apenas adicionar mais roupas, economizando muito tempo e esforço na jornada da Inteligência Artificial.
Título: LP-GEMM: Integrando Propagação de Layout em Operações GEMM
1. Problema Identificado
Em cargas de trabalho de Computação Científica e Aprendizado de Máquina (ML), como Redes Neurais (MLPs e Transformers), é comum encontrar sequências de multiplicações de matrizes dependentes (GEMMs - General Matrix Multiplications).
Ineficiência Atual: Bibliotecas de álgebra linear de ponta (como OpenBLAS, MKL, cuBLAS) são otimizadas para chamadas individuais de GEMM. No entanto, elas operam sob a restrição da API BLAS, que exige que cada chamada:
Empacote (pack) os dados de entrada em um layout específico e otimizado para o microkernel.
Execute o cálculo.
Desempacote (unpack) o resultado de volta para um layout canônico (padrão) na memória.
O Gargalo: Em uma sequência de GEMMs onde a saída de uma operação é a entrada da próxima, ocorre um ciclo redundante de empacotamento e desempacotamento em cada fronteira. O resultado de uma operação é desempacotado pelo produtor e imediatamente reempacotado pelo consumidor. Isso desperdiça ciclos de CPU, aumenta o tráfego de memória e reduz a eficiência computacional, especialmente em arquiteturas com vetores (AVX-512, RISC-V RVV).
2. Metodologia: LP-GEMM
O artigo propõe o LP-GEMM (Layout Propagation GEMM), uma decomposição do kernel GEMM do OpenBLAS que permite que o layout de dados empacotado seja propagado entre operações sequenciais, eliminando a necessidade de reempacotamento desnecessário.
Arquitetura da Solução
O GEMM padrão é dividido em três kernels especializados:
Ini-GEMM (Inicial):
Realiza o empacotamento inicial dos dados de entrada.
Executa o cálculo e armazena o resultado no layout propagado (otimizado para o microkernel), em vez do layout canônico BLAS.
Mid-GEMM (Intermediário):
Assume que a matriz multiplicadora já está no layout correto (proveniente do kernel anterior).
Pula a etapa de empacotamento da entrada.
Executa o cálculo e mantém o resultado no mesmo layout propagado.
É o kernel mais eficiente em cadeias longas.
End-GEMM (Final):
Realiza a última operação da cadeia.
Executa o cálculo e realiza o desempacotamento final, convertendo o resultado de volta para o layout canônico exigido pela API ou pelo usuário.
Integração com Microkernels
O sistema utiliza dois tipos de microkernels:
Propagate-Layout µKernel: Mantém o layout propagado durante o cálculo.
Default µKernel: Restaura o layout original (usado no End-GEMM).
Compatibilidade com Operações Não-GEMM: O trabalho aborda como lidar com operações intermediárias (como Softmax, RoPE em Transformers) que não são GEMMs. O LP-GEMM adapta essas operações para funcionarem eficientemente com o layout propagado, ajustando os strides (passos de acesso à memória) para evitar overhead de tradução de endereços.
3. Principais Contribuições
Decomposição do Kernel: Introdução dos kernels ini-GEMM, mid-GEMM e end-GEMM que exploram a propagação de layout de dados.
Implementação Prática: Desenvolvimento de uma versão autônoma em C++ do caminho de inferência do modelo Llama-3.2, utilizando exclusivamente chamadas GEMM de nível BLAS com propagação de layout.
Otimização de Operações ML: Implementação de outras operações comuns (Softmax, RoPE, RMSNorm) adaptadas para o layout propagado em arquiteturas AVX-512.
Validação Multi-Arquitetura: Avaliação extensiva em plataformas x86 (AVX-512) e RISC-V (RVV 1.0).
4. Resultados Experimentais
Os testes foram realizados comparando o LP-GEMM com bibliotecas de ponta (OpenBLAS, Intel MKL, BLIS, FlashGEMM).
GEMMs Isolados (Microbenchmarks):
O kernel Mid e End do LP-GEMM alcançaram um speedup mediano de 1.5x sobre o OpenBLAS, com o terceiro quartil ultrapassando 2x.
Isso demonstra que o custo de empacotamento/desempacotamento é frequentemente o fator dominante, superando até mesmo a otimização do microkernel em si.
Carga de Trabalho de Attention (Llama-3.2):
x86 (Intel Xeon): Aceleração de até 2.25x em média para GEMMs sequenciais. O ganho é maior em sequências curtas e diminui conforme o tamanho da matriz aumenta (devido à complexidade O(n3) da multiplicação vs. O(n2) do empacotamento), mas ainda é significativo.
RISC-V (SpacemiT X60): Aceleração ainda mais pronunciada, chegando a 5x em cargas de trabalho tipo Attention. Isso ocorre porque o kernel de referência do OpenBLAS no RISC-V tem ineficiências graves no desempacotamento (acessos de memória fora de ordem), que o LP-GEMM elimina completamente.
Sequências de 3 GEMMs: Em benchmarks simulando blocos de convolução, o LP-GEMM superou consistentemente o OpenBLAS e o FlashGEMM, provando a eficácia da propagação de layout mesmo em cenários complexos.
5. Significado e Conclusão
O LP-GEMM representa uma mudança de paradigma na otimização de álgebra linear:
Do Isolado para o Pipeline: Em vez de otimizar apenas chamadas individuais de GEMM, o foco muda para otimizar o fluxo de dados entre operações dependentes.
Portabilidade: A técnica é independente de microarquitetura específica (funciona bem tanto em x86 quanto em RISC-V) e pode ser combinada com kernels de baixo nível altamente otimizados.
Impacto Prático: A implementação no Llama-3.2 demonstra que é viável integrar essa otimização em modelos de IA reais sem reescrever toda a infraestrutura, exigindo apenas mudanças modestas na chamada de APIs.
Em resumo, o LP-GEMM prova que eliminar o trabalho redundante de reorganização de dados (packing/unpacking) entre operações dependentes é uma das formas mais eficazes de melhorar o desempenho em cargas de trabalho modernas de ML e HPC, oferecendo ganhos de performance competitivos com bibliotecas proprietárias de alto custo.