Goal-persistent coding agents optimize particle tracking code performance
Goal-Persistent Coding Agents as Scientific Performance Engineers: A Fixed-Radius Nearest-Neighbor Case Study
Artificial IntelligencePerformanceSoftware Engineering
Summary
Improving software performance can be tricky and requires careful testing and changes. This paper shows that certain coding agents, powered by AI models like Codex and Claude Code, can work through these tasks on their own. They started with a GPU-based program for tracking particles and made it run faster while keeping the results exactly the same. Their work demonstrates that AI agents can be used as performance engineers, provided clear tests and goals guide them.
What this means in practice
- •For software performance engineers: Use goal-directed AI coding agents to autonomously optimize and speed up GPU-based scientific code with guaranteed correctness and improved runtime.
- •For machine learning practitioners: Apply AI coding agents to refine and accelerate CUDA implementations in ML pipelines that depend on nearest-neighbor search and particle tracking applications.
Authors
Xiangyang Ju
Abstract
Coding agents can pursue persistent objectives across many tool-use turns, but evidence that general-purpose agents can conduct rigorous scientific performance engineering remains limited. We present a repository-scale case study in which off-the-shelf Codex and Claude Code agents optimize fixed-radius nearest-neighbor (FRNN) search for particle tracking. Starting from a PyTorch-dependent CUDA implementation, the agents follow an executable goal that specifies exact-correctness tests, profiling requirements, and acceptance criteria without prescribing code transformations. In the primary sequential trajectory, they autonomously remove the PyTorch dependency and conduct hypothesis-driven optimization experiments. The resulting standalone C++/CUDA library exactly reproduces the targeted reference result. Its synchronous NumPy interface achieved 1.6-fold speedup over the original GPU-resident PyTorch interface, despite including host transfers. Similar speedups were observed across different GPU architectures and software stacks. An independent optimization rerun followed a different sequence of hypotheses and reached even better performance on the target workload. These results show that goal-persistent coding agents can act as experimental performance engineers, and that executable scientific contracts are needed both to guide and to validate their optimization.