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

2024-04-27

Abstract: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).

Software Engineering,Artificial Intelligence,Machine Learning,Performance

What problem does this paper attempt to address?

The problem that this paper attempts to solve is how to improve the execution efficiency of software by optimizing program performance in the context of Moore's Law gradually becoming ineffective. Specifically, the paper focuses on high - level program optimization, such as the selection and improvement of APIs and algorithms, which are difficult to achieve due to the difficulty in understanding code semantics. Although pre - trained large - scale language models (LLMs) have shown strong capabilities in solving various programming tasks, their effectiveness in high - level program optimization remains an open research question. Therefore, the paper aims to explore how to adapt these large - scale language models for high - level program optimization.

To achieve this goal, the paper proposes a framework, which includes:
1. Constructed a data set containing more than 77,000 pairs of competitive C++ programming submissions written by human programmers, and each submission pair has undergone extensive unit tests.
2. Designed an environment based on the gem5 full - system simulator to reliably evaluate the impact of program optimization, thus avoiding significant variability when measuring performance on ordinary hardware.
3. Proposed a variety of adaptation strategies, including prompting strategies (such as retrieval - based few - shot prompting and chain - of - thought prompting) and fine - tuning strategies (such as performance - conditional generation and self - play - based synthetic data augmentation).

Through these methods, the paper demonstrates the potential of large - scale language models in program optimization and achieves significant performance improvements in multiple benchmark tests, even exceeding the best human performance.