PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

Each language version is independently generated for its own context, not a direct translation.

Imagine que você é um chef de cozinha de classe mundial (o Programador de IA) que criou uma receita incrível para um prato complexo (um Modelo de IA, como um chatbot ou um gerador de imagens).

O problema é que sua receita está escrita em uma linguagem poética e abstrata. O restaurante (o Hardware, como uma placa de vídeo NVIDIA) não entende essa linguagem. Ele precisa de instruções precisas: "pegue 3 ovos, bata por 10 segundos, coloque na panela X na temperatura Y".

Até agora, para fazer essa tradução, usávamos dois métodos principais:

O Método "Eager" (Ansioso): O chef tenta cozinhar passo a passo na hora. É rápido de começar, mas ineficiente. Ele vai até a geladeira (memória lenta) pegar um ingrediente, volta, pega outro, vai até a geladeira de novo... Muito tempo perdido andando.
O Método "Biblioteca de Pratos Prontos": O chef diz: "Quero um macarrão". O restaurante corre para o armário e pega um macarrão industrializado (uma biblioteca de código feita por especialistas, como a CuDNN da NVIDIA). É rápido, mas se você quiser um macarrão com um formato estranho ou uma combinação nova, o armário não tem. Você teria que reinventar a roda manualmente.

O que é o PolyBlocks?

O PolyBlocks é um super-gerente de cozinha (um compilador) que entra em cena. Ele não usa os pratos prontos do armário. Em vez disso, ele pega a sua receita poética, analisa tudo, e escreve do zero uma instrução de cozinha perfeita e otimizada para aquele restaurante específico.

Aqui está como ele funciona, usando analogias simples:

1. A Fábrica de Blocos (A Infraestrutura)

Imagine que o PolyBlocks é uma fábrica de blocos de montar (Lego).

Outros sistemas (como o Torch Inductor ou XLA) muitas vezes dependem de "peças prontas" (bibliotecas) que já vêm de fora. Se o restaurante mudar (um novo chip de IA), eles podem não ter a peça certa e precisam esperar o fabricante da peça.
O PolyBlocks cria todas as peças na hora. Ele é modular. Se você comprar um novo fogão (um novo chip de IA), o PolyBlocks apenas ajusta o molde para fazer as peças novas. Não precisa esperar ninguém.

2. O Jogo de "Fusão" (Juntar as Coisas)

Na cozinha, o maior desperdício de tempo é ir e voltar da geladeira (memória lenta) para a bancada (memória rápida).

Sem PolyBlocks: Você pega o tomate, volta para a bancada, pega a cebola, volta para a geladeira, pega o alho...
Com PolyBlocks (Fusão): O PolyBlocks olha para a receita e diz: "Ei, você vai usar tomate, cebola e alho juntos. Vamos pegar os três de uma vez e deixar tudo na bancada antes de começar a cozinhar".
- Ele funde várias etapas em uma só. Isso significa que o chip de IA não precisa ir até a memória lenta milhares de vezes. Ele faz tudo na "bancada" (memória rápida do chip), o que é muito mais rápido.

3. O "Tiling" (Cortar em Pedaços Perfeitos)

Imagine que você tem que pintar uma parede gigante.

Se você tentar pintar tudo de uma vez, você cansa e a tinta seca antes de terminar.
O PolyBlocks usa uma técnica chamada "Tiling" (Ladrilhamento). Ele divide a parede em quadrados pequenos (blocos). Ele pinta um quadrado, termina, e só então vai para o próximo.
- Isso é perfeito para a memória do chip. Ele carrega apenas o pedaço que precisa, processa, e joga fora, mantendo tudo organizado e rápido.

4. A Mágica da "Atenção" (O Cérebro da IA)

Os modelos modernos (como o que gera este texto) usam algo chamado "Camada de Atenção". É como se o chef precisasse ler um livro inteiro para responder a uma pergunta, mas só pode segurar algumas páginas na mão.

Fazer isso manualmente é um pesadelo de otimização.
O PolyBlocks tem um "truque de mágica" (chamado reduce-reduce fusion e wmma-fusion). Ele reorganiza a leitura do livro de forma que o chef nunca precise largar as páginas na mesa. Ele calcula tudo enquanto segura as páginas, sem precisar voltar à estante. Isso torna a IA muito mais rápida.

O Resultado na Prática

Os autores testaram esse "super-gerente" em placas de vídeo NVIDIA (A100 e A10).

Contra o método "Ansioso" (Eager): O PolyBlocks foi 2x a 3x mais rápido.
Contra os "Pratos Prontos" (Inductor/XLA): Em muitos casos, o PolyBlocks foi tão rápido ou até mais rápido do que as soluções que usam bibliotecas prontas da NVIDIA, mesmo sem usar as bibliotecas prontas!
O Grande Trunfo: Como o PolyBlocks cria o código do zero, ele é flexível. Se amanhã saírem novos chips de IA que ninguém conhece hoje, o PolyBlocks pode ser adaptado para eles rapidamente, enquanto os outros sistemas ficariam presos esperando atualizações das bibliotecas.

Resumo em uma frase

O PolyBlocks é um tradutor inteligente que pega ideias complexas de Inteligência Artificial e as transforma em instruções de cozinha super-eficientes, eliminando desperdícios e criando receitas personalizadas para qualquer tipo de fogão, sem depender de receitas prontas que podem não servir para todos.

Each language version is independently generated for its own context, not a direct translation.

Resumo Técnico: PolyBlocks

1. O Problema

O ecossistema de Inteligência Artificial (IA) enfrenta um desafio crescente: a lacuna entre frameworks de alto nível (como PyTorch, JAX e TensorFlow) e hardware especializado (chips de IA, GPUs).

Dependência de Bibliotecas Manuais: Compiladores atuais (como Torch Inductor e XLA) dependem pesadamente de bibliotecas de baixo nível otimizadas manualmente (ex: cuDNN, cuBLAS, FlashAttention) para obter desempenho. Isso limita a portabilidade para novos chips e impede otimizações cruzadas entre operadores (fusão) que não sejam suportadas por essas bibliotecas.
Falta de Reutilização: Criar compiladores para novos chips exige reescrever transformações complexas, pois as infraestruturas existentes não são facilmente reutilizáveis entre diferentes arquiteturas.
Compromisso Abstração vs. Desempenho: Abordagens de alto nível perdem desempenho ao não explorar otimizações específicas de hardware, enquanto abordagens de baixo nível (CUDA, Triton) exigem esforço manual massivo e conhecimento profundo de hardware.

2. Metodologia e Arquitetura

O PolyBlocks é uma infraestrutura de compilador modular e reutilizável baseada em MLIR (Multi-Level Intermediate Representation). Diferente de compiladores que dependem de bibliotecas externas, o PolyBlocks adota uma abordagem de geração de código totalmente automática (fully code-generating), gerando intrínsecos específicos do alvo a partir de alto nível.

Principais Decisões de Design:

Pipeline de Passes Modular: A infraestrutura é organizada em um pipeline de cinco estágios (S1 a S5) que transforma o IR de entrada até o código de máquina:
- S1-S2 (Frontend): Neutros ao alvo, baixam operações de tensores para semântica de buffers (memref) e depois para laços aninhados afins (affine).
- S3 (Otimizador de Médio Nível): O núcleo do PolyBlocks. Realiza a maioria das otimizações (fusão, tiling, mapeamento para unidades de matriz) usando análise afim leve e modelos de custo.
- S4-S5 (Backend): Convertem para dialectos específicos de alvo (ex: gpu, nvvm para NVIDIA) e geram código final (LLVM/NVPTX).
Análise Afim Leve: Em vez de usar técnicas poliedrais completas (que são computacionalmente caras), o PolyBlocks utiliza análise afim simples e verificações lineares para a maioria dos casos, recorrendo a bibliotecas de conjuntos inteiros apenas quando estritamente necessário. Isso garante tempos de compilação rápidos (segundos).
Abordagem "Full Code-Generating": Não há dependência de bibliotecas de terceiros. O compilador gera todo o código de baixo nível, permitindo otimizações cruzadas (fusão de operadores) que bibliotecas manuais não conseguem realizar.

3. Contribuições e Otimizações Chave

O papel destaca várias transformações inovadoras implementadas no estágio S3:

Fusão Baseada em "Slicing" (Slicing-based Fusion): Vai além da fusão de laços tradicional. Calcula fatias (slices) do produtor necessárias para o consumidor, permitindo fusão mesmo em laços imperfeitos ou com dependências complexas. Isso elimina buffers intermediários e reduz acessos à memória global.
Fusão e Tiling em Duas Fases:
1. Realiza tiling (divisão em blocos) nos laços de destino (ex: matrizes de saída) primeiro.
2. Em seguida, realiza a fusão baseada em fatias, puxando os produtores para dentro dos laços já tilados. Isso preserva a localidade de dados e permite fusão eficiente entre operadores como convoluções e matrizes.
Fusão da Camada de Atenção (Attention Layer): Implementa uma fusão automática de toda a operação de atenção (incluindo $Q \times K^T$ $Q \times K^{T}$ , softmax e multiplicação por $V$ $V$ ).
- Usa passes de reduce-reduce-fusion e wmma-fusion para eliminar idas e vindas à memória DRAM e à memória compartilhada on-chip, mantendo os dados nos registradores.
- Suporta variações de atenção (como FlexAttention) sem esforço adicional significativo.
Empacotamento "On-the-Fly" para Convoluções: Transforma convoluções em multiplicações de matrizes (GEMM) sem materializar a matriz expandida (que consumiria muita memória). O PolyBlocks gera código que empacota os dados de entrada e pesos diretamente na memória rápida (on-chip) durante a execução, suportando todas as configurações de convolução (stride, dilatação, padding).
Mapeamento para Unidades de Matriz (Tensor Cores): Transforma automaticamente laços afins para utilizar instruções de multiplicação de matriz (ex: WMMA em GPUs NVIDIA), sem exigir que o usuário especifique manualmente o mapeamento.

4. Resultados Experimentais

Os testes foram realizados em GPUs NVIDIA A10 e A100, comparando o PolyBlocks com Torch Inductor, XLA, TensorRT e execução eager (PyTorch padrão).

Desempenho Geral (PyTorch):
- Em batch size 1 (cenário de inferência em tempo real), o PolyBlocks foi 2.15x mais rápido que a execução eager e 1.4x mais rápido que o Torch Inductor.
- Em batch size 8, foi 1.8x mais rápido que o eager e competitivo (0.97x) com o Inductor.
- Superou o TensorRT em vários casos.
Desempenho em Operadores Individuais:
- Para convoluções, o código gerado pelo PolyBlocks foi competitivo com a biblioteca cuDNN em centenas de casos, superando-a em mais de 2x em cerca de 50 casos.
- Para multiplicação de matrizes (matmul), o desempenho foi comparável ao cuBLAS e ao Triton (escrito manualmente).
JAX: O PolyBlocks também foi aplicado ao JAX, mostrando um ganho de 2.12x sobre a execução eager e 1.15x sobre o XLA.
Estudo de Ablação: A análise mostrou que a fusão cruzada de operadores (cross-operator fusion) sozinha trouxe um ganho de 2.87x sobre código já otimizado, destacando a importância de reduzir acessos à memória global. O uso de Tensor Cores trouxe um ganho médio de 17x.

5. Significado e Impacto

O PolyBlocks representa um avanço significativo na compilação de IA:

Portabilidade e Reutilização: Permite criar rapidamente compiladores para novos chips de IA reutilizando a mesma infraestrutura de transformações, sem depender de bibliotecas de fornecedores que podem não existir para novos hardwares.
Otimização Automática Superior: Demonstra que um compilador totalmente gerador de código pode igualar ou superar bibliotecas manuais altamente otimizadas, especialmente em cenários de fusão de operadores complexos (como camadas de atenção) onde bibliotecas tradicionais falham.
Escalabilidade: A abordagem baseada em MLIR e análise afim leve permite compilar modelos grandes (com milhares de laços) em segundos, tornando a compilação JIT (Just-In-Time) viável para fluxos de trabalho de produção.
Futuro da Infraestrutura de IA: O trabalho sugere que a dependência de bibliotecas manuais pode ser superada por compiladores inteligentes, facilitando a adoção de aceleradores de IA personalizados e heterogêneos.

Em resumo, o PolyBlocks preenche a lacuna entre a produtividade de frameworks de alto nível e o desempenho máximo de hardware especializado, oferecendo uma infraestrutura modular que automatiza otimizações complexas de forma eficiente e portátil.