Resumo Técnico: Geração de Testes Baseada em Dados para Fuzzing de Compiladores de IA

Declaração do Problema

Compiladores de Inteligência Artificial (IA) (ex: TVM, TensorRT, ONNXRuntime) são essenciais para implantar modelos de IA em diversos hardwares, mas permanecem propensos a bugs que comprometem a confiabilidade do compilador e a correção do modelo. O processo de compilação envolve três estágios distintos: (1) Carregamento do Modelo, onde modelos de bibliotecas como PyTorch ou Keras são convertidos para Representações Intermediárias (IR) de alto nível; (2) Otimização de Alto Nível, envolvendo transformações de grafo independentes de hardware; e (3) Otimização de Baixo Nível, que aplica transformações específicas de hardware para gerar código executável.

Técnicas de teste existentes (ex: NNSmith, HirGen, Tzer) focam primariamente em gerar entradas válidas e diversas, mas falham em fornecer uma cobertura abrangente e específica para cada estágio. Especificamente, elas enfrentam dificuldades com:

1. Cobertura de Frontend Diversificada: Garantir que toda a lógica de conversão de operadores de várias bibliotecas de IA seja tratada corretamente durante o carregamento do modelo.
2. Consciência do Estágio de Otimização: Gerar grafos computacionais que acionem efetivamente otimizações de alto nível complexas e sensíveis ao contexto.
3. Validação Específica de Hardware: Criar testes que incorporem recursos necessários para acionar otimizações de baixo nível específicas de hardware.

Metodologia

Os autores propõem o OPERA, um framework de teste unificado baseado em dados que integra três componentes complementares, cada um visando um estágio específico do pipeline de compilação de IA:

1. OPERA: Teste de Carregamento de Modelo Baseado em Migração

O OPERA aborda o estágio de carregamento de modelo migrando testes existentes de bibliotecas de IA.

• Abordagem: Ele instrumenta as APIs no código-fonte da biblioteca de IA para executar os testes da biblioteca e extrair instâncias de operadores. Essas instâncias são encapsuladas em modelos de operador único via geração baseada em templates.
• Otimização: Para gerenciar o alto custo de execução dos testes migrados, o OPERA emprega um algoritmo de priorização consciente da diversidade. Ele agrupa operadores por semântica para maximizar a cobertura enquanto elimina redundâncias.
• Objetivo: Maximizar a cobertura de diversas lógicas de conversão de operadores, aproveitando a semântica central já verificada pelos testes da biblioteca.

2. OATest: Teste de Otimização de Alto Nível Baseado em Síntese

O OATest visa otimizações de alto nível através da síntese de grafos computacionais conscientes da otimização.

• Abbordagem: Ele instrumenta o código-fonte do compilador para extrair padrões de otimização de testes criados por desenvolvedores de compiladores. O processo envolve identificar o grafo computacional correspondente a uma otimização e derivar padrões baseados na granularidade da otimização.
• Síntese: O OATest combina esses padrões extraídos com contextos diversos (grafos semente). Ele seleciona aleatoriamente um padrão e uma semente, inserindo o padrão em um ponto de síntese aleatório para aumentar a diversidade estrutural.
  ോടെ
• Garantias de Validade: Para evitar testes inválidos causados por entradas indefinidas, o OATest emprega duas estratégias de correção: (1) reutilizar entradas/saídas compatíveis de nós existentes e (2) criar novos nós para estabelecer conexões corretas entre padrões e contextos.

3. HARMONY: Teste de Otimização de Baixo Nível Baseado em Mutação

O HARMONY foca em otimizações específicas de hardware usando uma abordagem baseada em mutação.

• Geração de Sementes: Em vez de gerar testes do zero, o HARMONY utiliza Grandes Modelos de Linguagem (LLMs) guiados por restrições de operadores coletadas de múltiplas fontes de conhecimento para gerar sementes de IR de baixo nível válidas e diversas. Essa verificação cruzada mitiga riscos associados a documentações desatualizadas.
• Mutação: Um LLM de análise extrai padrões de otimização de documentações e exemplos de código. Um LLM de código então realiza mutações leves nas sementes de IR existentes de acordo com esses padrões.
• Objetivo: Produzir testes semanticamente válidos que sejam especificamente propensos a acionar otimizações de baixo nível direcionadas (ex: ocultação de latência de memória, mapeamento de intrínsecos), minimizando o risco de saídas inválidas causadas por alucinações de LLM.

Principais Contribuições

O artigo apresenta três técnicas específicas que formam um framework unificado:

1. OPERA: Uma abordagem baseada em migração que aproveita testes de bibliotecas de IA existentes para descobrir bugs na lógica de conversão de operadores.
2. OATest: Uma técnica de síntese consciente da otimização que extrai padrões semânticos de testes criados por desenvolvedores para gerar grafos computacionais sensíveis ao contexto.
3. HARMONY: Uma técnica baseada em mutação que utiliza documentação e LLMs para gerar e mutar sementes de IR de baixo nível para testes de otimização de hardware específicos.

Resultados

O framework foi avaliado em quatro compiladores de IA amplamente utilizados (TVM, TensorRT, OpenVINO e ONNXRuntime), detectando 266 bugs anteriormente desconhecidos. O detalhamento por componente é o seguinte:

• OPERA: Descobriu 170 bugs (90 confirmados por desenvolvedores) em oito frontends. A estratégia de priorização melhorou a eficiência de teste em 11,9%–47,4% em comparação com estratégias gerais. As causas raiz dos bugs abrangeram categorias como Problemas de Forma de Tensor, Problemas de Tipo e Lógica de Código Incorreta.
• OATest: Descobriu 56 bugs (42 confirmados, 24 corrigidos) no TVM e ONNXRuntime. Alcançou 60,2% mais cobertura de branch e 66,98% mais cobertura de linha em comparação com NNSmith e WhiteFox.
• HARMONY: Detectou 40 bugs (26 confirmados e corrigidos) no TVM. Demonstrou cobertura substancialmente maior no estágio de otimização de baixo nível em comparação com fuzzers de estado da arte como NNSmith, WhiteFox e Tzer.

Significância e Alegações

O artigo afirma que este trabalho fornece um framework abrangente e consciente dos estágios que aborda sistematicamente os desafios específicos de cada etapa de compilação, em vez de tratar o compilador como uma caixa preta monolítica. Ao integrar técnicas de migração, síntese e mutação impulsionadas por dados (testes de biblioteca, padrões de desenvolvedores e documentação), o framework aumenta significativamente a cobertura e a eficácia dos testes.

Os autores posicionam este trabalho como um passo para garantir a confiabilidade dos compiladores de IA em todo o pipeline. Eles observam que, embora os resultados atuais sejam significativos, o trabalho futuro estenderá o framework para compiladores emergentes (ex: Triton, MLC-LLM) e investigará a localização automática de bugs e reparo para reduzir ainda mais os esforços de manutenção dos desenvolvedores.