CelloAI: Leveraging Large Language Models for HPC Software Development in High Energy Physics

TL;DR

CelloAI is a locally hosted AI assistant leveraging large language models and retrieval-augmented generation to improve documentation and code development in high energy physics HPC environments, addressing legacy software challenges.

Contribution

This work introduces CelloAI, a novel local LLM-based tool with retrieval capabilities tailored for HEP HPC code documentation and generation, ensuring privacy and large context handling.

Findings

Effective code documentation generation for HEP applications.

Improved code retrieval accuracy with syntax-aware chunking.

Enhanced code understanding and refactoring suggestions.

Abstract

Next-generation High Energy Physics (HEP) experiments will generate unprecedented data volumes, necessitating High Performance Computing (HPC) integration alongside traditional high-throughput computing. However, HPC adoption in HEP is hindered by the challenge of porting legacy software to heterogeneous architectures and the sparse documentation of these complex scientific codebases. We present CelloAI, a locally hosted coding assistant that leverages Large Language Models (LLMs) with retrieval-augmented generation (RAG) to support HEP code documentation and generation. This local deployment ensures data privacy, eliminates recurring costs and provides access to large context windows without external dependencies. CelloAI addresses two primary use cases, code documentation and code generation, through specialized components. For code documentation, the assistant provides: (a) Doxygen style comment generation for all functions and classes by retrieving relevant information from RAG sources (papers, posters, presentations), (b) file-level summary generation, and (c) an interactive chatbot for code comprehension queries. For code generation, CelloAI employs syntax-aware chunking strategies that preserve syntactic boundaries during embedding, improving retrieval accuracy in large codebases. The system integrates callgraph knowledge to maintain dependency awareness during code modifications and provides AI-generated suggestions for performance optimization and accurate refactoring. We evaluate CelloAI using real-world HEP applications from ATLAS, CMS, and DUNE experiments, comparing different embedding models for code retrieval effectiveness. Our results demonstrate the AI assistant's capability to enhance code understanding and support reliable code generation while maintaining the transparency and safety requirements essential for scientific computing environments.