Source-linked AI summary
Learning Performance-Improving Code Edits
Alexander Shypula, Aman Madaan, Yimeng Zeng, Uri Alon, Jacob Gardner, Milad Hashemi, Graham Neubig, Parthasarathy Ranganathan, Osbert Bastani, Amir Yazdanbakhsh
TL;DR
High-level program optimization remains difficult because algorithms and API choices require semantic understanding, while reliable performance measurement and open datasets are limited. This paper introduces PIE and simulator-based evaluation to adapt LLMs for code optimization, finding that suitable adaptation strategies improve performance and can make open models competitive with private ones.
Problem
High-level performance decisions remain largely the responsibility of programmers, while narrow techniques, limited datasets, and unreliable measurement hinder automated optimization research.
Method
The paper builds PIE, a dataset of over 77 K C++ program pairs, and evaluates LLM optimization strategies with deterministic gem5-based performance measurement.
Results
Fine-tuned CODELLAMA 13B reached 66.56% %OPT and 5.65× SPEEDUP for BEST@8, while GPT-3.5 with synthetic data reached 87.63% %OPT and 6.86× SPEEDUP for BEST@8.
Takeaways & Limitations
With suitable adaptation strategies, open models can be competitive with private models, and LLMs can optimize code beyond optimizing compilers given a correctness oracle.
Takeaways & Limitations
The experiments assume gem5 terminates; a two-minute timeout may introduce slight variability, and changing it could change the results.
Abstract
from arXiv · showhide
With the decline of Moore's law, optimizing program performance has become a major focus of software research. However, high-level optimizations such as API and algorithm changes remain elusive due to the difficulty of understanding the semantics of code. Simultaneously, pretrained large language models (LLMs) have demonstrated strong capabilities at solving a wide range of programming tasks. To that end, we introduce a framework for adapting LLMs to high-level program optimization. First, we curate a dataset of performance-improving edits made by human programmers of over 77,000 competitive C++ programming submission pairs, accompanied by extensive unit tests. A major challenge is the significant variability of measuring performance on commodity hardware, which can lead to spurious "improvements." To isolate and reliably evaluate the impact of program optimizations, we design an environment based on the gem5 full system simulator, the de facto simulator used in academia and industry. Next, we propose a broad range of adaptation strategies for code optimization; for prompting, these include retrieval-based few-shot prompting and chain-of-thought, and for finetuning, these include performance-conditioned generation and synthetic data augmentation based on self-play. A combination of these techniques achieves a mean speedup of 6.86 with eight generations, higher than average optimizations from individual programmers (3.66). Using our model's fastest generations, we set a new upper limit on the fastest speedup possible for our dataset at 9.64 compared to using the fastest human submissions available (9.56).
1 INTRODUCTION
The paper targets high-level program optimization, where existing tools and learning methods are constrained by limited datasets and unreliable performance measurement. It introduces PIE and evaluates data-driven LLM adaptation strategies for this task.
- High-level algorithm and API choices remain largely the programmer’s responsibility, while existing learning approaches are narrow or difficult to extend.
- PIE collects performance-improving edits from evolving C++ competitive-programming submissions.
- gem5 provides deterministic and reproducible performance evaluation, avoiding spurious improvements caused by commodity-hardware variability.
- PIE contains more than 77 K C++ program pairs and supports prompting and fine-tuning strategies for optimization.
2 PERFORMANCE IMPROVING EDITS (PIE) DATASET
The PIE dataset turns human performance-improving edits into reproducibly measured C++ program pairs. It filters for meaningful improvements, validates correctness with extensive tests, and separates problems across dataset splits.
- PIE represents performance-improving edits from C++ programs written for competitive-programming tasks.
- Pairs are retained only when the measured relative time improvement exceeds 10%.
- More than 42.8 million gem5 simulations provide the dataset’s annotated execution times.
- The dataset contains 77,967 training pairs, 2,544 validation pairs, and 978 test pairs, with problems kept exclusive to each split.
- Correctness is evaluated with unit tests, expanded to a median of 82.5, 75, and 104 test cases per problem across training, validation, and test data.
- Real-hardware benchmarking can report false optimizations because execution variance makes identical programs appear different.
3 ADAPTING CODE LLMS TO PROGRAM OPTIMIZATION
The paper adapts code LLMs through prompting and fine-tuning strategies that exploit PIE’s slow–fast program pairs. These include retrieval of analogous examples, performance conditioning, and synthetic-data augmentation, with gem5 termination as a stated caveat.
- Few-shot prompting: Chain-of-thought prompting asks the model to reason about optimization before producing the revised program.
- Few-shot prompting: Dynamic retrieval-based few-shot prompting selects analogous PIE programs using C++ embeddings and FAISS to construct prompts during inference.
- Fine-tuning: Fine-tuning uses PIE directly and also addresses dataset imbalance through a 4,085-pair high-quality subset.
- Caveat: The fine-tuning setup assumes gem5 terminates, and its two-minute timeout may introduce slight variability or change results if altered.
- Fine-tuning: Performance-conditioned generation tags fast programs with information about their optimal achievable performance.
- Fine-tuning: Synthetic augmentation generates novel competitive-programming examples with GPT-3.5 and filters duplicates before training.
4 EXPERIMENTS
The experiments compare prompting and fine-tuning strategies for adapting code LLMs to performance optimization, using correctness-aware speedup evaluation and BEST@k sampling. Data-driven adaptation, especially fine-tuning and performance conditioning, substantially outperforms baseline prompting, while the fastest model generation slightly exceeds the fastest human submission speedup.
- Few-shot prompting: Dynamic retrieval-based few-shot prompting outperforms baseline prompting, though increased speedup may cost correctness.The preferred setting retrieves K = 4 prompts.
- Fine-tuning: 5.65× SPEEDUP is achieved by fine-tuned CODELLAMA 13B with performance-conditioned generation at BEST@8.The corresponding result is 66.56% %OPT.
- Fine-tuning: Performance-conditioned generation improves both %OPT and SPEEDUP over traditional fine-tuning for fine-tuned CODELLAMA models.The performance-improvement information helps distinguish optimal from sub-optimal solutions.
- Fine-tuning: Synthetic self-play data improves %OPT and often SPEEDUP, particularly with BEST@1, suggesting improved generalization.The reported effect is marginal and is described for both CODELLAMA and GPT-3.5.
- Discussion and key takeaways: 9.64× is the model’s fastest PIE-dataset speedup, slightly above the 9.56× fastest human speedup and establishing a new upper limit.The comparison covers 39,129 model generations and 118,841 human accepted solutions.
- Discussion and key takeaways: Open CODELLAMA models can be competitive with private GPT-3.5 when paired with appropriate adaptation strategies.CODELLAMA 13B with performance-conditioned generation approached GPT-3.5 with synthetic data.
- Discussion and key takeaways: For models of the same size, fine-tuning significantly outperforms prompting, although retrieval makes prompting effective.The comparison is explicitly between prompting with retrieval and fine-tuning.
5 RELATED WORK
Machine learning has been applied to multiple forms of performance engineering, spanning code refactoring, compiler transformations, parameter search, vectorization, GPU optimization, and algorithm selection.
- Related work: Prior machine-learning approaches target refactoring, compiler transformations, parameter search, vectorization, GPU code, and algorithm selection.These applications cover both source-code changes and performance-engineering decisions.
6 CONCLUSION
The work is an initial step toward using LLMs for high-level computing-stack opportunities, including algorithmic efficiency and code optimization beyond optimizing compilers when correctness can be checked.
- LLMs can improve algorithmic efficiency and enable automatic code optimization beyond optimizing compilers given a correctness oracle.
A.1 ADDITIONAL ANALYSIS OF GENERATED CODE EDITS
The generated edits span algorithmic, I/O, data-structure, and miscellaneous transformations, with algorithmic changes the largest category. Most observed speedups appear substantive, but some derive from reducing array-allocation constants and may reflect incomplete test coverage.
- Algorithmic Transformations: 34.15% of analyzed edits were algorithmic transformations, including recursive-to-dynamic-programming changes and formula-based replacements.
- Input/Output Operations: 26.02% of edits changed input/output operations, often replacing C++ standard I/O with faster C-standard methods.
- Data Structure Modifications: 21.14% of edits modified data structures, including transitions from vectors to arrays, pointer removal, and hashmap use when appropriate.
- Miscellaneous Optimizations: 18.70% of edits were miscellaneous optimizations such as omitting unnecessary initialization and replacing costly functions with constants.
- Caveat: 10 of 120 analyzed speedups came from reducing array-allocation constants, which may indicate insufficient test coverage rather than genuine algorithmic improvement.
- Examples: Figure 3 illustrates eliminating sorting by computing the maximum start and minimum end in one O(n) pass instead of O(n log n) sorting.
A.2 EXAMPLES OF OPTIMIZATIONS
The examples show PIE-generated optimizations that remove unnecessary work through direct computation, avoided precomputation, modular arithmetic, and loop restructuring.
- Figure 4: PIE removes a wasteful nested loop by directly computing the range of 2k −1 consecutive numbers centered around x.
- Figure 5: The optimized factorial implementation computes only the requested input instead of precomputing values through 10^5.
- Figure 6: The modular-arithmetic optimization reduces complexity from O(10^6) to O(B) by checking (A × i) mod B = C for i in [0, B −1].
- Figure 7: The optimized loop improves performance by restructuring control flow to avoid an additional conditional check inside the loop.
A.3 CONVERGENCE OF GPT3.5 FINE-TUNED MODELS WITH ADDITIONAL GENERATIONS
With more generations, GPT-3.5 models trained on HQ data and HQ plus self-play converge toward similar performance, while self-play slightly changes the speedup–correctness trade-off.
- HQ-only and HQ + Self-Play models gradually converge to similar performance as the number of generations increases.
- Self-Play changes speedup from 6.74 to 6.86 and correctness from 86.71 to 87.63, alongside a slight drop in correctness after training.
A.4 ERROR ANALYSIS
The analysis identifies correctness failures as the main source of unsuccessful optimization attempts, while longer descriptions and programs slightly reduce accuracy and speedup. It also examines whether strict performance conditioning limits optimization diversity.
- Failure causes: Approximately 60% of GPT-3.5 Self-Play failures break a unit test.About 30% produce correct programs that are slower or miss the speedup threshold, while about 10% have another failure cause.
- Program properties: Longer problem descriptions and source programs show mild negative correlations with optimization accuracy.The reported correlations are -0.15 for description length and -0.26 for source-program length.
- Program properties: Average speedup also has mild negative correlations with problem description length (-0.16) and source program length (-0.11).These results suggest that longer inputs are associated with slightly lower average speedup in the analysis.
- Performance conditioning: The authors investigate whether conditioning generation only on a 10/10 optimization rate constrains the available optimizations for each input.They compare six generations from 10/10 conditioning with combinations of generations conditioned on 10/10, 9/10, and 8/10.
A.5 PIE DATASET DETAILS
PIE is built from performance-improving C++ submission pairs and supports reproducible optimization experiments with simulator-based runtime annotations. The section also documents prompting, self-play, retrieval, fine-tuning, and measurement-quality details.
- Synthetic data: Self-play prompting generated 10,000 programs, of which 6,553 were outside the PIE splits and formed 3,314 equivalence sets.Evaluating these generations required more than 1.4 million binary input pairs and took under 72 hours on a 24-core processor.
- Fine-tuning: Fine-tuning experiments include CODELLAMA 7B and 13B variants trained with FSDP across eight 48GB GPUs.The supplied implementation details describe convergence training for approximately 4,000 high-quality examples and one-epoch full-data or performance-conditioned training.
- Measurement and adaptation caveats: Duplicate CodeNet submissions can report different CPU runtimes despite identical code and environments, motivating caution about raw runtime annotations.One example reports a 2.44× speedup between identical programs, while the LoRA section reports a hypothesized gap because optimization examples may be rare in training data.
- Prompting: Retrieval-based prompting selects analogous programs or challenges, using CodeBertScore to retrieve close source programs and their slow-to-fast pairs.The experiments vary the number of retrieved prompts across K = {1, 2, 4}.