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

Mahesh Natarajan; Xiaoye Li; Weiqun Zhang

cs.AIMar 28, 2026

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

Authors: Mahesh Natarajan, Xiaoye Li, Weiqun Zhang

Organizations: Applied Mathematics and Computational Research Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA

Abstract

We present AstraAI, a command-line interface (CLI) coding framework for high-performance computing (HPC) software development. AstraAI operates directly within a Linux terminal and integrates large language models (LLMs) with Retrieval-Augmented Generation (RAG) and Abstract Syntax Tree (AST)-based structural analysis to enable context-aware code generation for complex scientific codebases. The central idea is to construct a high-fidelity prompt that is passed to the LLM for inference. This prompt augments the user request with relevant code snippets retrieved from the underlying framework codebase via RAG and structural context extracted from AST analysis, providing the model with precise information about relevant functions, data structures, and overall code organization. The framework is designed to perform scoped modifications to source code while preserving structural consistency with the surrounding code. AstraAI supports both locally hosted models from Hugging Face and API-based frontier models accessible via the American Science Cloud, enabling flexible deployment across HPC environments. The system generates code that aligns with existing project structures and programming patterns. We demonstrate AstraAI on representative HPC code generation tasks within AMReX, a DOE-supported HPC software infrastructure for exascale applications.

Figures & tables

Figure 1: End-to-end AstraAI workflow: User prompts are interpreted to extract intent, and intent-specific pipelines combine retrieval-augmented context, AST-guided constraints, and controlled code generation or editing. Tools used at each stage are highlighted in red, and the workflow is orchestrated via Python scripts.

Panel of Figure 2 — Figure 2: General structure of the chunks for retrieval (top), an example code snippet in AMReX with the metadata (middle), and the code chunk in .json format (bottom).

Panel of Figure 3 — Figure 3: clang-query command examples to extract structure information of the code.

Panel of Figure 4 — Figure 4: General structure of the high-fidelity prompt.

Figure 5: An example user prompt that attempts to port a nested for loop in C++ into AMReX (top), the full prompt to the LLM (middle), and the result of the inference (bottom).

Task	Baseline (user prompt only)
	CodeLlama 13B	CodeLlama 34B	GPT-OSS-120B	Claude-Sonnet-4.6-High
1. 3D nested for loop to AMReX kernel conversion	0.32	0.99	0.92	0.94
2. Creating a copy of a MultiFab from another	0.60	0.65	0.82	0.87
3. MultiFab linear combination	0.77	0.83	0.90	0.85
4. 7-point finite-difference Laplacian kernel	0.88	0.85	0.87	0.97
5. Neumann boundary ghost cell fill	0.70	0.71	0.72	0.79

Table 1: Baseline performance across models showing the cosine similarity metric.

Task	AstraAI (user prompt + RAG chunks + AST info)
	CodeLlama 13B	CodeLlama 34B	GPT-OSS-120B	Claude-Sonnet-4.6-High
1. 3D nested for loop to AMReX kernel conversion	0.78	0.92	0.91	0.92
2. Creating a copy of a MultiFab from another	0.87	0.89	0.96	0.95
3. MultiFab linear combination	0.88	0.98	0.94	0.93
4. 7-point finite-difference Laplacian kernel	0.80	1.00	0.96	0.96
5. Neumann boundary ghost cell fill	0.99	0.99	0.92	0.90

Table 2: Performance of AstraAI across models showing the cosine similarity metric.

Retrieval-Augmented Code Generation

See all 14 papers

An Agentic Evaluation Framework for AI-Generated Scientific Code in PETSc
Mar 16, 2026Hong Zhang, Barry Smith, Satish Balay +4High-Performance Computing Code Generation
Argonne National Laboratory, Lemont, IL, USA
Retrieval-Augmented Generation for Scientific Code Understanding
Sep 10, 2026Aaron Nobile, Andreas Adelmann, Mohsen SadrRetrieval-Augmented Generation LLM Grounding
Department of Physics, ETH Zürich, 8093 Zürich, Switzerland · Paul Scherrer Institute, 5232 Villigen PSI, Switzerland · Massachusetts Institute of Technology, Cambridge, MA 02139, USA
CodegenBench: Can LLMs Write Efficient Code Across Architectures?
Jun 1, 2026Jie Li, Wenzhao Wu, Junqi Hu +5High-Performance Computing Code Generation
1Sun Yat-sen University · 2National Supercomputing Center in Wuxi · 3National Supercomputing Center in Shenzhen +2

Mar 16, 2026J/K move · Enter open · S save

An Agentic Evaluation Framework for AI-Generated Scientific Code in PETSc

Hong Zhang, Barry Smith, Satish Balay, Le Chen+3

Argonne National Laboratory, Lemont, IL, USA

High-Performance Computing Code Generation AI Agent Evaluation Scientific Code Generation Code Generation Evaluation Code Generation Benchmarks

Open paper AskPDF arXiv

While LLMs have accelerated scientific code generation, comprehensively evaluating generated code remains challenging. Many benchmarks emphasize functional correctness or task completion, which is insufficient for code built on production HPC libraries, where solver selection, API conventions, memory management, parallel awareness, and performance also matter. We introduce PETSCAgent-Bench, a multidimensional benchmark and agent-based framework for assessing whether AI-generated scientific code uses a production HPC library as an expert would. A tool-augmented evaluator compiles, executes, and measures code and combines deterministic checks with LLM-based assessments in a 14-evaluator pipeline spanning five categories: correctness, performance, code quality, algorithmic appropriateness, and library-specific conventions. A2A and MCP enable black-box evaluation of compatible coding agents. Across realistic PETSc problems, frontier models generate readable, well-structured code but struggle with correctness on challenging problems and with library-specific conventions even when code compiles and runs---limitations that conventional pass/fail evaluation does not capture.

Abstract

Retrieval-Augmented Code Generation

Explore similar work

An Agentic Evaluation Framework for AI-Generated Scientific Code in PETSc

Retrieval-Augmented Generation for Scientific Code Understanding

CodegenBench: Can LLMs Write Efficient Code Across Architectures?