Resumo Técnico: SOLAR – Análise de Desempenho de Velocidade da Luz Baseada em IA

1. Declaração do Problema

Os aceleradores de aprendizagem profunda modernos alcançam throughputs de pico na ordem de petaFLOPS, mas as cargas de trabalho do mundo real raramente se aproximam desses limites teóricos. Existe uma lacuna crítica em entender quão rápido um modelo poderia rodar no hardware alvo versus quão rápido ele realmente roda. As ferramentas atuais falham em abordar isso de forma abrangente:

• Contadores de FLOPs (ex: fvcore, ptflops) sofrem de cobertura de operadores estreita (75–84% no KernelBench), ignoram o tráfego de I/O e não conseguem derivar tempos mínimos de execução teóricos (limites de Velocidade da Luz ou SOL).
• Profilers de hardware (ex: NCU, NSys) medem o desempenho alcançado de implementações específicas, em vez dos limites teóricos, podendo reportar "folga" (headroom) que é fisicamente inalcançável devido a restrições algorítmicas.
• Estimativa pura por LLM é propensa a erros, particularmente em cargas de trabalho compostas, e carece das capacidades de busca sistemática necessárias para análise sensível ao cache.
• Frameworks de modelagem de desempenho existentes (ex: Timeloop, AccelForge) exigem especificação manual de carga de trabalho, o que é propenso a erros e desconectado do desenvolvimento rápido de modelos.

Não existe ferramenta existente que derive automaticamente limites de SOL validados diretamente a partir de código-fonte de alto nível (PyTorch/JAX) com ampla cobertura de operadores.

2. Metodologia

SOLA (Speed-of-Light Analysis for Runtime) é um framework de fonte-para-SOL que preenche essa lacuna através de um pipeline de três estágios que separa a tradução generativa da análise determinística.

Estágio 1: Frontend Agêntico Validado

O primeiro estágio traduz programas de origem arbitrários em um IR de Loop Afim executável.

• LLM como Parser: Um agente de LLM traduz o código-fonte para o IR de Loop Afim, uma representação intermediária independente de linguagem.
• Estrutura do IR: O IR divide os programas em:
  • Kernels de Loop Afim Restritos: Loops perfeitamente aninhados sobre tensores de dimensões nomeadas com limites estáticos e expressões de índice afim, projetados para levantamento mecânico para einsums.
  • Camada de Composição Irrestrita: Lida com controle de fluxo, movimentação de dados (reshape, slice) e operações esparsas.
• Loop de Validação: Para garantir a correção, o sistema executa o programa original para capturar saídas "golden", gera o IR via LLM, executa o IR (usando Numba) e compara as saídas. Se ocorrer uma incompatibilidade, o feedback de diagnóstico guia o LLM para corrigir a tradução. Este loop de "gerar-então-verificar" garante que o pipeline downstream receba apenas traduções numericamente validadas.

Estágio 2: Geração de Einsum Determinística

Uma vez validado o IR, um compilador determinístico o eleva para um Grafo de Einsum.

• Construção do Grafo: O backend rastreia o IR para construir um Grafo Acíclico Dirigido (DAG) onde os nós são einsums estendidos (suportando reduções como max) e as arestas representam dependências de dados de tensores.
• Derivação Mecânica: A partir da estrutura de subscritos dos einsums, o sistema deriva mecanicamente fórmulas exatas de contagem de MAC e tráfego de memória.
• Análise de Fusibilidade: A estrutura do grafo identifica operadores adjacentes que podem compartilhar buffers on-chip, permitindo que o sistema calcule o custo de operações fundidas versus execução independente.

Estágio 3: Análise de SOL

O backend analítico computa tempos mínimos teóricos de execução ($T_{SOL}$) contra especificações de hardware alvo (throughput de pico $\Pi$ e largura de banda de DRAM $\beta$).

• Modelo de Roofline: $T_{SOL} = \max(\frac{\text{FLOPs}}{\Pi}, \frac{\text{Traffic}}{\beta})$.
• Modos de Multi-Fidelidade:
  1. Não Fundido (Unfused): Assume execução independente com viagens de ida e volta completas à DRAM para intermediários.
  2. Fundido (Fused): Elimina o tráfego de DRAM intermediário assumindo capacidade infinita de buffer on-chip para operadores fundidos.
  3. Fundido Sensível ao Cache (OROJENESIS): Modela a capacidade finita de buffer on-chip. Utiliza programação linear inteira para otimizar estratégias de tiling, contabilizando re-leituras da DRAM quando o conjunto de trabalho excede os limites do cache.
• Roofline Inverso: Dado um tempo de latência alvo, o sistema deriva as especificações mínimas de hardware necessárias ($\Pi_{min}$ e $\beta_{min}$).

3. Principais Contribuições

1. Primeiro Framework Fonte-para-SOL: O SOLAR é a primeira ferramenta a derivar limites de SOL validados diretamente do código-fonte PyTorch e JAX, alcançando 100% de cobertura de análise no KernelBench.
2. Pipeline Híbrido Generativo-Determinístico: Ele combina de forma única um LLM para tradução independente de linguagem (com validação numérica) e um compilador determinístico para análise de desempenho reproduzível. Isso evita os erros de estimativas puras por LLM enquanto contorna o esforço manual da modelagem tradicional.
3. Avaliação Abrangente: O framework fornece evidências em casos de uso de kernel, algoritmo e arquitetura, demonstrando capacidades em análise de headroom, identificação de otimização, exploração cross-platform e provisionamento de hardware.

4. Resultados

A avaliação foi conduzida no KernelBench (270 problemas), modelos JAX/Flax e cargas de trabalho de robótica no Jetson Thor.

• Análise de Headroom: O SOLAR revela lacunas de otimização significativas. No KernelBench, a lacuna entre o tempo de execução medido e o SOL fundido varia de 4.3× (L1) a 54.6× (L3) para a execução eager do PyTorch. Mesmo o torch.compile deixa um headroom substancial (ex: 47.7× no L3), indicando que a fusão em nível de grafo é crítica para cargas de trabalho compostas.
• Insights de Multi-Fidelidade:
  • Fusão: A análise fundida revela até 7.8× de headroom adicional comparado à análise por operador no KernelBench.
  • Caching: Os limites sensíveis ao cache tornam o SOL mais rigoroso em até 2.25× em relação aos limites fundidos ingênuos, identificando o tiling como um gargalo em cenários específicos.
  • Otimização Algorítmica: O reordenamento de cadeias de einsum (ex: no DeepSeek MLA) alcançou uma redução de 2.04× em FLOPs em relação ao ordenamento ingênuo.
• Projeção Cross-Platform: O SOLAR projeta o desempenho em quatro alvos de hardware (B200, H100, A6000, Jetson Thor) sem acesso físico. Os benefícios de fusão variam significamente por plataforma (1.8× no B200 vs. 14.5× no Jetson Thor), destacando a interação entre largura de banda e fusão.
• Provisionamento de Hardware (Roofline Inverso): Para modelos de robótica visando loops de controle de 500 Hz, o SOLAR identificou que a largura de banda é a principal restrição. Por exemplo, o modelo DreamZero WAM requer uma melhoria de 19.7× na largura de banda atual para atingir as metas de latência, sugerindo que upgrades de silício sozinhos são insuficientes sem compressão de modelo ou mudanças algorítmicas.

5. Significância e Alegações

O artigo afirma que o SOLAR muda fundamentalmente como a otimização de desempenho é abordada, tornando os limites teóricos acessíveis e acionáveis sem modelagem manual.

• Preenchendo a Lacuna: Ele fecha o descompasso entre o desenvolvimento rápido de modelos e a análise de desempenho de hardware, fornecendo um teto de otimização "validado".
• Insights Acionáveis: Ao oferecer análise de multi-fidelidade, ajuda engenheiros a distinguir entre gargalos causados por falta de fusão, má estratégia de tiling ou ineficiências algorítmicas fundamentais.
• Design Independente de Hardware: Permite que arquitetos e designers de algoritmos explorem espaços de design e provisionem requisitos de hardware antes mesmo do silício existir, evitando o desperdício de computação e tokens em trabalhos com pouco headroom restante.

Os autores enfatizam que, embora a análise seja baseada em formas de tensores (não em valores) e assuma comportamento ideal de hardware (ignorando throttling térmico), ela fornece a primeira visão automatizada, validada e abrangente da velocidade da luz teórica para cargas de trabalho de deep learning.