LLM Program Optimization via Retrieval Augmented Search

This codebase contains the code for the paper "LLM Program Optimization via Retrieval Augmented Search" by Sagnik Anupam, Alexander Shypula, and Osbert Bastani. It implements the Retrieval-Augmented Search (RAS), Atomic Edit Guided Search (AEGIS), Dynamic Retrieval and Instruct-Only methods described in the paper for two datasets, PIE (Shypula et al., 2024) and Mercury (Du et al., 2024).

Note: for Mercury, only the code for methods reported in the paper (RAS, No Contextual, Instruct-Only) is supported.

Installation

1. The code requires OpenAI organization and API keys (saved in ORGANIZATION and OPENAI_API_KEY environment variables respectively), and an OpenRouter API key (saved in OPENROUTER_API_KEY environment variable). The code requires setting up three environments: one for OpenRouter sampling, which uses Python 3.11.13, one for evaluating Mercury programs, using Python 3.12.3, and one for evaluating PIE programs, which uses Python 3.8.17.
2. Create the Mercury evaluation environment using conda env create -f mercury/mercury_eval_environment.yml (default environment name: mercuryeval). Then, after activating environment with conda activate mercuryeval, cd mercury/human-eval and run pip install -e .. Then, copy the conda environment Python path (the output of which python after activating the environment, by default looks something like <PATH_TO_CONDA>/envs/mercuryeval/bin/python on Linux) to MercuryPaths's python_path attribute in config.py.
3. Create the PIE evaluation environment using conda env create -f pie_eval_environment.yml (default environment name: py38_gem5) and copy its Python path (usually <PATH_TO_CONDA>/envs/py38_gem5/bin/python on Linux) to PiePaths's python_path attribute in config.py.
4. Finally, create the OpenRouter sampling environment using conda env create -f openrouter_environment.yml (default environment name: langchain_for_pie), and activate it with conda activate langchain_for_pie.
5. (Optional) To use any of our retrieval datasets or experiment "starter" files, unzip the PIE files provided here into experiments/common_dataset and/or the Mercury files provided here into mercury/experiments/common.

Experiment Setup

In our experiments, we assume access to training and test sets comprising of (slow, fast) program pairs. We denote the slow program using variables with the abbreviation src (for source), and the faster program using variables with the abbreviation tgt (for target). To see how to construct an example retrieval dataset, see decomposition.py. The retrieval datasets provided in the installation steps above were the GPT-4o-generated ones used in the paper experiments. The retrieval_dataset and target_programs/target_progs attributes in config.py point to the PIE (decomposed)/Mercury training set and human-written original training set target programs respectively. Variables titled model_name usually refer to the OpenRouter model ID (e.g. qwen/qwen3-coder for Qwen3-Coder), while model_folder typically refers to a sanitized folder name for saving the model experiments.

For each experiment, we recommend starting from a common "starter" file which contains a shared step at iteration 1. For example, for running RAS, this shared step file contains the prompts for generating test set program descriptions. In order to create a starter file for an experiment, execute the corresponding starter file creation function in starter_file_generation.py e.g. generate_mercury_none_experiment_starter_file(mercury=True) for the Mercury Instruct-Only Baseline, and then execute the experiment in ras_aegis_experiments.py using the corresponding experiment function (in this case, model_mercury_none_experiment(model_name=model_name, model_folder=model_folder, env=env)). The starter files used in the paper experiments are provided in the Installation steps. To create starter files with your own programs, add Python programs to the filepath in MercuryPaths's common_task_dataset attribute in config.py (for C++ programs, add the programs to the filepath in PiePaths's common_task_dataset attribute).

Mercury-Specific Setup Steps

1. The Mercury evaluator downloads the mercury dataset from huggingface, so set hf_cache_dir in config.py to the absolute path to the mercury subfolder before running, to ensure that the subprocess call in ras_aegis_experiments.py can download to the directory.
2. Run setup_mercury_experiments in starter_file_generation.py prior to running any Mercury experiments - it executes the source programs in the test set to create a reference file to use to measure the beyond metric.

PIE-Specific Setup Steps

1. For PIE, since gem5 evaluations are slow, it is recommended to use a reference runtime file to avoid running the source and human-optimized target test set programs after each optimization iteration to calculate relative speedups. To setup a reference runtime file, run at least one experiment with include_src_tgt = True. Note that if a reference runtime file does not exist at the filepath specified in config.py's utilsConfig class's reference_runtime_path attribute, the PIE experiments will not work correctly when include_src_tgt is set to False, since relative speedups won't be calculated.
2. After running a PIE evaluation, it is possible that the Docker container stalls - it is safe to terminate the process (i.e. all runtime files have been saved) once the following line is printed after a PIE evaluation: "Mean performance for the different types of generations is: ". Note that between successive PIE experiment runs, it is necessary to stop the docker container that was used to run gem5 (since the code will try to create a new docker container listening at the same port) - otherwise, you may see an Internal Server Error where the bind failed with a port is already allocated message.

Acknowledgements

Thank you to the authors of the PIE and Mercury for making their code publicly available.

Funding

This work is funded in part by the UPenn DARPA MOCHA CR PO-0074571/10101353 subcontract of Peraton Labs. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of DARPA or Peraton Labs.

Citation

@article{anupam2026llm,
  author    = {Anupam, Sagnik and Shypula, Alexander and Bastani, Osbert},
  title     = {LLM Program Optimization via Retrieval Augmented Search},
  journal   = {Findings of the Association for Computational Linguistics: ACL 2026},
  year      = {2026},
  url       = {https://arxiv.org/abs/2501.18916},
}