Resumo Técnico: Rumo a Modelos de Mundo para Compiladores

Declaração do Problema

A otimização de programas de tensores é essencial para sistemas modernos de aprendizado de máquina, mas a identificação do programa mais eficiente é dificultada por um enorme espaço de busca combinatória. Para evitar o custo proibitivo de medições em dispositivo, os auto-agendadores existentes (ex: Ansor, AutoTVM) dependem de modelos de custo aprendidos para classificar candidatos a programas. No entanto, esses modelos tipicamente tratam os candidatos como instantâneos de código estáticos (ex: texto final de TensorIR).

Essa perspectiva estática introduz duas limitações críticas:

1. Perda de Contexto de Transformação: Ela ignora o processo de transformação passo a passo que gerou o programa. Consequentemente, os avaliadores não conseguem entender como uma decisão de agendamento em um passo posterior depende do contexto estabelecido por decisões anteriores.
2. Sensibilidade ao Ruído Sintático: Codificadores estáticos são facilmente induzidos ao erro por variações superficiais de código. Por exemplo, dois programas estruturalmente equivalentes (ex: multiplicação de matrizes com tamanhos de loop $N=1024$ vs. $N=1023$) podem produzir textos de TensorIR significativamente diferentes devido ao tratamento de borda (peeling de loop), fazendo com que o modelo classifique incorretamente seu potencial de desempenho, apesar de possuírem comportamento de hardware idêntico.

Metodologia

Os autores propõem um framework de Modelo de Mundo para Compiladores que reformula a avaliação de agendamento como dinâmica latente condicionada à ação sobre estados de programas. Em vez de avaliar um artefato de código estático final, o framework simula a evolução do estado do programa através de um espaço latente contínuo.

O framework consiste em três componentes aprendidos instanciados dentro do AutoScheduler do TVM:

1. Representação de Estado de Programa de Tensor (Encoder)

• Entrada: Texto TensorIR (estados pré-agendamento e pós-agendamento).
• Arquitetura: Uma arquitetia baseada em CodeBERT treinável, adaptada para entradas longas. Utiliza uma janela deslizante para tokenizar e codificar blocos sobrepostos, agregando-os via média mascarada (masked mean pooling) em uma representação de vetor denso $z = E(s)$.
• Treinamento: Utiliza aprendizado contrastivo para mapear estados de TensorIR similares (provenientes da mesma fonte de carga de trabalho) próximos no espaço latente, enquanto afasta estados dissimilar. Isso cria um espaço de estados onde a equivalência semântica é preservada, apesar das variações sintáticas.

2. Transição de Estado Multietapa Condicionada à Ação

• Objetivo: Prever a representação do estado terminal induzida por uma sequência de ações de agendamento sem materializar explicitamente as ASTs intermediárias.
• Mecanismo: Partindo de um embedding de estado inicial $z_0$, o modelo executa o rollout do efeito de uma sequência de ações de agendamento $a_{1:T}$ passo a passo.
  • As ações são serializadas e codificadas em embeddings densos.
  • Um modelo de transição $F(\cdot)$, implementado usando uma arquitetura baseada em TransH (originalmente da aprendizagem de representação de grafos de conhecimento), prevê o próximo estado latente: $z_t = F(z_{t-1}, a_t)$.
  • Isso modela as ações de agendamento como transformações geométricas em um hiperplano específico de relação, rastreando efetivamente a proveniência da decisão e resolvendo ambiguidades referenciais (ex: o que uma variável de loop representa após a fusão).
• Treinamento: Utiliza perda de rollout de múltiplas etapas combinando um termo de correspondência de estado (minimizando a distância para estados futuros reais) e um termo de ranking baseado em margem (garantindo que o estado previsto esteja mais próximo do estado real do que de amostras negativas).

3. Avaliação de Candidatos Baseada em Ranking

• Entrada: A representação do estado terminal prevista $\hat{z}_T$, características estruturadas de ação (ex: contagem de ações, comprimento da sequência) e características de hardware opcionais.
• Saída: Um score escalar para priorização relativa de candidatos dentro da mesma carga de trabalho.
• Arquitetura: Um modelo de ranking XGBoost treinado com um objetivo LambdaRank.
• Estratégia de Treinamento: O modelo não prevê o tempo de execução absoluto. Em vez disso, ele aprende a classificar candidatos dentro de um grupo de carga de trabalho com base em tempos de execução medidos, tratando o problema como uma tarefa de aprendizado de ranking (learning-to-rank).

Principais Contribuições

1. Perspectiva de Modelo de Mundo para Compiladores: O artigo adapta o paradigma de modelo de mundo para otimização de tensores, estabelecendo um framework onde os agendamentos candidatos são modelados através da evolução do estado latente condicionado à ação, em vez de instantâneos estáticos.
2. Dataset de Trajetória Ação-Estado: Construção de um dataset baseado em TVM/TenSet que organiza logs de ajuste (tuning) em trajetórias alinhadas de ação-estado (estado pré-agendamento, sequência de ações, estados intermediários, estado pós-agendamento) para previsão de múltiplas etapas.
3. Integração de Busca Eficiente: Implementação do framework no AutoScheduler do TVM, demonstrando que a modelagem de dinâmicas latentes reduz significativamente os custos de autotuning mantendo alta qualidade de busca.

Resultados Experimentais

O método foi avaliado em uma CPU Intel Xeon Gold 6430 e uma GPU NVIDIA RTX 4090 através de sete modelos de redes neurais (incluindo ResNet, BERT, GPT e OPT) e 103 subgrafos extraídos.

• Desempenho vs. Ansor: Sob um orçamento de 64 tentativas, o método melhora a latência do subgrafo representativo sobre o Ansor em 1.37× na GPU e 1.54× na CPU.
• Eficiência de Amostragem: O método alcança uma qualidade de busca comparável ao Ansor-10K (10.000 tentativas) usando apenas 1.000 tentativas (10× menos medições). Especificamente, supera ligeiramente o Ansor-10K na média aritmética e iguala-se a ele dentro de 2.2% na média geométrica em subgrafos representativos.
• Aceleração End-to-End: Comparado ao PyTorch/PyTorch-opt (cuDNN), o método entrega acelerações de média geométrica de 4.61× (CPU) e 3.67× (GPU), com acelerações de pico atingindo até 58.61×.
• Estudos de Ablação: Remover o módulo de transição aprendido (revertendo para codificação de estado inicial ou final estática) resultou em aumentos significativos de latência (mediana de ~56–69%), confirmando que a modelagem da evolução do estado intermediário é o componente mais crítico.

Significância e Limitações

Significância:
O artigo argumenta que tratar a otimização de compiladores como um processo dinâmico de transição de estados, em vez de um problema de avaliação de código estático, permite que os avaliadores capturem as dependências semânticas entre decisões de agendamento. Essa abordagem suprime efetivamente o ruído sintático e resolve ambiguidades referenciais, levando a uma busca mais eficiente com menos medições de hardware.

Limitações (conforme declarado pelos autores):

1. Dependência do Espaço de Busca: O método melhora a avaliação de candidatos, não a geração. Sua eficácia é limitada pela capacidade do agendador subjacente de propor agendamentos fortes; um avaliador melhor não pode recuperar agendamentos que nunca foram propostos.
2. Acúmulo de Erro: Erros de previsão podem se acumular ao longo de trajetórias de agendamento longas ou complexas no espaço latente, potencialmente reduzindo a confiabilidade em casos difíceis.
3. Ranking vs. Predição: O modelo é treinado para ranking dentro da carga de trabalho, não para predição de latência absoluta. Isso o torna ideal para priorização online, mas menos adequado para configurações que exigem estimativas de desempenho calibradas entre diferentes cargas de trabalho.