Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

Resumo Técnico: Diagnóstico Hierárquico Baseado em Compilador para Otimização de Kernels Triton via LLM

1. Definição do Problema

Embora os Grandes Modelos de Linguagem (LLMs) tenham avançado na geração e otimização automatizada de kernels, as abordagens existentes dependem predominantemente de "sinais de superfície", como erros de compilação, métricas de tempo de execução e dados de profiling. Esses sinais indicam que um kernel está lento, mas raramente explicam o porquê de o compilador de backend falhar em realizar uma otimização proveitosa. Isso é particularmente agudo em aceleradores emergentes como NPUs (ex: Ascend), onde a pilha de compilação é em constante evolução, específica para o alvo e opaca para os desenvolvedores.

O desafio central identificado é um gap semântico entre o código-fonte de nível de operador e o comportamento de desempenho do backend. Falhas de desempenho frequentemente emergem apenas após o código passar por estágios de representação intermediária (IR), decisões de escalonamento e passes de compilação específicos do backend. Consequentemente, agentes de LLM, tipicamente pré-treinados em padrões centrados em GPU, comportam-se como otimizadores de tentativa e erro: eles propõem edições de código plausíveis, mas carecem de evidências para mapear sintomas de tempo de execução a causas específicas no nível do compilador.

2. Metodologia

O artigo propõe um framework de otimização hierárquico e fundamentado em compilador que trata a otimização de kernels como um problema de diagnóstico progressivo entre camadas. Em vez de tratar o embasamento do compilador como um recurso sempre ativo, o sistema escala do evidenciamento leve para a análise profunda do compilador apenas quando necessário.

O Fluxo de Trabalho Hierárquico:
O sistema opera através de um pipeline de ciclo fechado organizado em quatro níveis ascendentes:

• L1: Triagem de Padrões (Pattern Triage): O sistema primeiro verifica hipóteses fortes baseadas em padrões usando um índice curado de padrões de otimização, heurísticas e regras de especialistas (incluindo heurísticas internas destiladas de execuções anteriores). Se uma reescrita de baixo risco for justificada por evidências existentes, o processo para aqui.
• L2: Diagnóstico de Profiling: Se os padrões falharem, o sistema coleta evidências de tempo de execução (latência, vazão, contadores de hardware) para localizar gargalos (ex: movimentação de memória, atividade escalar, sincronização). Isso identifica o sintoma, mas não necessariamente o mecanismo.
• L3: Atribuição de IR: Quando o profiling é insuficiente, o sistema inspeciona as Representações Intermediárias (IR) para traduzir sintomas de tempo de execução em diagnósticos estruturais (ex: cargas de passo de memória não unitário, falha de fusão, falhas de vetorização induzidas pelo layout). Ele reconstrói a trajetória de redução (lowering) do operador para identificar transformações estágio a estágio.
• L4: Escalação Compilador-Fonte: Como último recurso, o sistema invoca a análise fundamentada no compilador para explicar por que um padrão diagnosticado bloqueia a otimização do backend. Ele consulta o conhecimento do compilador (comportamentos de passes, regras de redução, restrições) para recuperar as condições no nível de código-fonte que uma reescrita bem-sucedida deve satisfazer.

Arquitetura do Sistema e Gestão de Estado:

• Design Centrado em Agente: Implementado como um ambiente centrado em agente (integrado ao repositório de habilidades do CannBot) em vez de um serviço de compilador bespoke. Utiliza um invólucro de execução (execution harness) fino para gerenciar validação, correção e coleta de evidências, enquanto "habilidades" específicas de domínio lidam com a lógica de otimização.
• Rodadas com Estado (Stateful Rounds): O processo de otimização é organizado em rodadas explícitas e auditáveis. Cada rodada registra sua hipótese, fontes de evidência e resultado de validação. Este design com estado permite que o sistema seja retomável e garante que análises mais profundas (IR/Compilador) sejam anexadas como evidências reutilizáveis, em vez de observações efêmeras.
• Síntese Retrospectiva: Além do loop online, o sistema realiza uma passagem de síntese cross-operator mais lenta. Ele mapeia históricos de otimização validados para famílias de padrões existentes ou cria novos padrões reutilizáveis, permitindo que o nível "mais barato" (Triagem de Padrões) melhore ao longo do tempo.

3. Principais Contribuições

1. Formulação do Problema: Os autores formulam a otimização de kernels baseada em LLM como um problema de diagnóstico progressivo entre camadas, identificando a falta de evidências fundamentadas no compilador como a razão primária pela qual agentes baseados apenas em profiling falham em otimizações específicas de backend.
2. Fluxo de Trabalho Hierárquico: Eles descrevem um fluxo que escala da triagem de padrões para profiling, atribuição de IR e escalação compilador-fonte, propondo reescritas de nível de fonte baseadas em evidências apenas quando evidências mais profundas são necessárias.
3. Evidência de Avaliação: Apresentam evidências de avaliação padronizadas e conscientes da distribuição em um benchmark Ascend 950 derivado do NPUKernelBench, reportando ganhos agregados, melhor rodada contabilizada e traços representativos entre camadas.

4. Resultados Experimentais

O sistema foi implementado para Triton para NPUs Ascend e avaliado em 37 operadores convertidos com sucesso de um benchmark padronizado Ascend 950 (derivado do NPUKernelBench). Cada operador foi testado em 5 casos (185 casos no total).

• Ganhos de Desempenho:
  • Média Geométrica de Aceleração (Speedup): 4.35× (do kernel Triton inicial para o otimizado).
  • Mediana de Aceleração: 2.73×.
  • Distribuição: 22/37 entradas excederam 2× de aceleração; 13/37 excederam 5×. Os resultados mostram uma distribuição de cauda longa, com ganhos substanciais coexistindo com entradas próximas ao baseline.
• Contabilidade do Processo:
  • O índice de melhor rodada mediana foi 10, indicando que resultados ótimos frequentemente exigiram múltiplas iterações.
  • Profundidade de Evidência: Embora 33 das selecionadas "melhores" rodadas tenham sido atribuídas à Triagem de Padrões, o sistema demonstrou a utilidade de camadas mais profundas: 32 trajetórias atingiram o Diagnóstico de Profiling, 3 atingiram a Atribuição de IR e 4 atingiram a Escalação Compilador-Fonte.
  • Estudo de Caso: Um traço representativo (padding de tamanho de Hyena FFT) mostrou como o profiling identificou um gargalo escalar, a atribuição de IR identificou custos de sincronização específicos na redução (lowering) e a análise compilador-fonte explicou a restrição, levando a um robusto aumento de velocidade de 2.08×.

5. Significância e Alegações

O artigo afirma que o embasamento do compilador deve ser seletivo e não sempre ativo. Ao tratar a otimização como um diagnóstico progressivo, o sistema torna o raciocínio profundo do compilador prático, invocando-o apenas quando evidências mais leves não conseguem produzir uma reescrita suficientemente fundamentada.

Os autores enfatem que seus resultados demonstram um espaço considerável para melhoria (headroom) em um subconjunto de operadores, mas não alegam melhoria uniforme em todos os casos. Eles declaram explicitamente que a avaliação atual não fornece isolamento causal (ex: comparando a escalação progressiva contra estratégias de apenas profiling ou de sempre-profundo sob um toolchain congelado). O artigo posiciona-se como um "argumento de sistemas" apoiado por benchmarks transparentes e evidências de processo, reservando a decomposição causal da hierarquia para experimentos controlados futuros.

O trabalho destaca que a relação entre o código-fonte e o desempenho em tempo de execução é mediada por transformações invisíveis do compilador, e que preencher esse gap com uma abordagem hierárquica e baseada em evidências é essencial para otimizar kernels em aceleradores emergentes e opacos, como as NPUs.