LLM Program Optimization via Retrieval Augmented Search

Recent work has demonstrated the potential of large language models (LLMs) for program optimization, a key challenge in programming languages.We propose a blackbox adaptation method called Retrieval Augmented Search (RAS) that performs beam search over candidate optimizations; at each step, it retrieves incontext examples from a given training dataset of slow-fast program pairs to guide the LLM.Critically, we find that performing contextual retrieval based on an LLM-generated natural language description significantly outperforms retrieval based on the source code.We also propose AEGIS, a method for improving interpretability by decomposing training examples into "atomic edits" that are significantly more incremental in nature.We show that RAS performs up to 2.06× better than prior state-of-theart blackbox adaptation strategies on optimizing C++ programs, and that AEGIS performs up to 1.37× better while making significantly smaller edits.We also show that using RAS improves the mean runtime percentile of Python programs by 10.27 compared to baselines.