AstraAI: LLMs, Retrieval, and AST-Guided Assistance for HPC Codebases

1. O Problema

O desenvolvimento e manutenção de software para Computação de Alto Desempenho (HPC) são tarefas complexas devido a bases de código extensas, dependências intrincadas, paralelismo de baixo nível e restrições rigorosas de desempenho. Ferramentas de IA generativa existentes (como assistentes de código gerais) frequentemente falham em fornecer sugestões seguras para HPC, pois carecem de:

• Compreensão Semântica Profunda: Não entendem as nuances específicas de frameworks científicos.
• Garantias de Correção: Tendem a gerar código que viola padrões de projeto ou introduz erros sutis em estruturas de dados paralelas.
• Contexto Estrutural: Falham em manter a consistência com a organização do código existente, o que é crítico para a estabilidade de compilação e execução em escala.

2. Metodologia: A Arquitetura do AstraAI

O AstraAI é uma ferramenta de interface de linha de comando (CLI) projetada para operar diretamente em terminais Linux, integrando três pilares principais para gerar código contextualmente consciente:

• Geração Aumentada por Recuperação (RAG):

  • O sistema indexa o repositório de código-fonte (ex: AMReX) em "chunks" (fatias) de código com metadados ricos (tipo de tarefa, intenções do usuário, palavras-chave de API).
  • Durante a inferência, a solicitação do usuário é convertida em um vetor de embedding (usando o modelo all-MiniLM).
  • O sistema recupera os trechos de código mais semanticamente relevantes do banco de dados para fornecer exemplos de padrões específicos do framework e uso de APIs ao modelo de linguagem (LLM).

• Análise Baseada em Árvore de Sintaxe Abstrata (AST):

  • Para garantir a correção estrutural, o AstraAI realiza análise estática do código.
  • Utiliza ferramentas como intercept-build e clang-query para gerar um banco de dados de compilação (compile_commands.json) e extrair a AST.
  • Isso permite a extração precisa de assinaturas de funções, variáveis de membro, especificadores de acesso e intervalos de linhas de origem.
  • Essas informações estruturais são injetadas no prompt, dando ao LLM uma visão explícita da interface do programa e reduzindo ambiguidades.

• Construção de Prompt de Alta Fidelidade:

  • O prompt final enviado ao LLM é uma composição estratificada contendo:
    1. Instruções gerais do sistema (estilo de codificação).
    2. Contexto recuperado via RAG (exemplos de código).
    3. Informações estruturais extraídas via AST.
    4. A solicitação original do usuário.
  • O modelo pode ser hospedado localmente (via Hugging Face e Ollama, otimizado para GPUs como A100) ou acessado via API (American Science Cloud).

• Edição Segura:

  • O código gerado é formatado com marcadores de conflito estilo GitHub, permitindo que os desenvolvedores revisem e integrem as alterações mantendo a estabilidade da construção do projeto.

3. Principais Contribuições

• Integração Híbrida: A principal inovação é a combinação sinérgica de RAG (para conhecimento de domínio) e AST (para precisão estrutural) para guiar a geração de código em ambientes HPC.
• Framework Específico para HPC: Diferente de ferramentas genéricas, o AstraAI é desenhado para manter a consistência com padrões de projeto específicos (como o AMReX, infraestrutura de software para aplicações exascale).
• Flexibilidade de Implantação: Suporta tanto modelos locais (privacidade, controle, conformidade com restrições de HPC) quanto modelos de ponta via API.
• Pipeline Automatizado: Oferece um fluxo de trabalho totalmente automatizado em Python, desde a extração de contexto até a inferência, facilitando a reprodutibilidade.

4. Resultados e Avaliação

O sistema foi avaliado no framework AMReX utilizando quatro modelos de linguagem: CodeLlama-13B, CodeLlama-34B, GPT-OSS-120B e Claude-Sonnet-4.6-High. A métrica principal foi a similaridade de cosseno entre o código gerado e uma implementação de referência verificada (com normalização de variáveis locais via AST para focar na estrutura semântica).

• Desempenho em Modelos Menores: Modelos menores (CodeLlama-13B/34B) apresentaram baixa similaridade na linha de base (apenas prompt do usuário), indicando falta de conhecimento prévio sobre AMReX. Com o AstraAI (RAG + AST), a similaridade aumentou drasticamente (ex: de 0.32 para 0.78 na conversão de loops 3D; de 0.47 para 0.37 em um caso específico, mas com melhoria geral significativa em tarefas complexas como solvers multigrid, onde o CodeLlama-34B atingiu 1.00).
• Desempenho em Modelos de Ponta: Modelos maiores já possuíam bom desempenho na linha de base devido ao treinamento prévio em repositórios open-source. O AstraAI ainda ofereceu melhorias ou manteve o alto desempenho, garantindo consistência estrutural.
• Conclusão Quantitativa: A adição de contexto recuperado e restrições estruturais transformou modelos menores em ferramentas viáveis para tarefas específicas de HPC, reduzindo a necessidade de treinamento fino (fine-tuning) massivo para cada novo projeto.

5. Significado e Impacto

O AstraAI representa um avanço significativo na adoção de IA para o desenvolvimento de software científico:

• Segurança e Confiabilidade: Ao ancorar a geração de código em estruturas reais do projeto (AST) e exemplos validados (RAG), reduz o risco de introduzir erros que quebrariam simulações científicas complexas.
• Produtividade: Permite que desenvolvedores realizem modificações complexas (como portar código sequencial para paralelo GPU/MPI) com menos esforço manual de verificação.
• Acesso a Código Proprietário: Diferente de modelos puramente baseados em dados públicos, o AstraAI permite a geração de código assistida por IA para bases de código proprietárias ou internas, desde que o contexto seja fornecido localmente.
• Futuro: O trabalho abre caminho para refatoração automatizada de múltiplos arquivos e suporte a modelos ainda maiores via paralelismo de modelos (DeepSpeed), consolidando a IA como uma ferramenta essencial para a próxima geração de aplicações exascale.