Learning Performance Improving Code Edits with

Aman Madaan¹, Alexander Shypula², Uri Alon¹, Milad Hashemi³, Parthasarathy Ranganathan³, Yiming Yang¹, Graham Neubig^1,4, Amir Yazdanbakhsh⁵

¹Carnegie Mellon University ²University of Pennsylvania ³Google
⁴Inspired Cognition ⁵Google Research, Brain Team

Paper Code Data

TLDR: A dataset and models for learning code edits that can improve runtime. Browse through some optimizations performed by codex/codegen using PIE below. The examples are selectively picked for brevity.

Caption 1

About

Q. About PIE

PIE stands for Performance Improving Code Edits. We also train strong baselines, showing that CODEGEN-16B can be finetuned using PIE to match closely with CODEX on speedup. PIE is curated from CodeNet. Thanks to project codenet for making the dataset available! Each problem statement comes with a set of unit tests to check correctness.

Q. What are we releasing?

We are releasing the dataset of trajectories (slow-fast pairs) in Python and C++, our code for preparing the dataset, and our prompts. We are also releasing our evaluation harness, allowing you to evaluate new methods on the dataset. (more details on the repo). We are working on releasing our finetuned CodeGen models but if you are interested in trying them out, please reach out to us.

Q. But ChatGPT/CODEX can already do everything?

While anecdotal examples are fun, they can be misleading. Our benchmark evaluation confirms that codex is pretty good at the task but is far from perfect! Crucially, we leverage pie to finetune CodeGen-16B, and show that it can closely match CODEX in terms of speedup (2.62x for codex vs. 2.45x for codegen-16b finetuned on PIE). Further, CodeGen-16B + PIE is *better* than codex in terms of the lines of code changed (56% changed by CodeGen-16B vs. 70% changed by codex) (please also see the note above on usability of optimization methods and the importance of making minimal-changes).

Q. What makes PIE unique?

We collect trajectories of code edits from CodeNet, where a programmer iteratively modifies an initial program to improve its runtime efficiency. The fact that in each trajectory, a single programmer makes the edits, for a single program and a single goal (improving runtime efficiency) makes this dataset unique.

Q. Why are trajectories important?

Trajectories are important because they capture the evolution of a program as it is being written by a developer. The differences between subsequent edits are localized to a small number of lines of code. The fact that the changes are minimal is critical in training models that can perform surgical code edits. We do not want complete rewrites; we want to build models that can propose minimal changes to a program. We consider this to be an important aspect of our work.

Q. Why is it essential to have minimal changes?

We believe that the ability to make minimal changes is critical for the success of any program repair system. A system that completely rewrites a program is unlikely to be helpful in practice --- the programmer will have to verify that the new program is correct, and that it still achieves the same goal as the original program. In contrast, a system that can propose a small number of edits to a program is more likely to be useful in practice. The programmer can quickly verify that the edits are correct, without re-verifying the entire program.

Q. Who cares if my python script runs in 0.5 seconds instead of 0.4 seconds?

A lot of real-world use cases! Trading algorithms, daemons on cloud servers, query processing, os kernels, etc., are just a few examples where programs frequently run at scale. For programs executed millions of times a day, even a tiny improvement in runtime can significantly impact the overall system performance. Finally, it's not all about tiny changes --- we have examples of ~100x speedups generated using PIE! (see the summation example above)

Q. What if the optimizations add bugs?

Each trajectory in PIE comes with a unit test. We use the unit test to evaluate the correctness of the optimized program. In some cases, the optimization may introduce a bug that the unit tests miss. Therefore, we will also release an extended set of unit tests sourced from code contests.

Q. How large is the dataset?

The filtered dataset used for finetuning has ~40k pairs for Python and ~100k for C++. We also release the entire dataset, which has ~150k pairs for Python and ~200k for C++. The dataset is available on our [github](https://github.com/madaan/pie-perf).

Q. How can I help?

If you are interested in helping, please reach out to the authors!

Abstract

The waning of Moore's Law has shifted the focus of the tech industry towards alternative methods to achieve continued performance gains. To optimize program efficiency, programmers must craft and refactor code with better performance characteristics. In this paper, we investigate the ability of large language models (LLMs) to suggest functionally correct, performance-improving code edits. We hypothesize that language models can suggest such edits in ways that would be impractical for static analysis alone. To evaluate and improve the capacity of large language models, we have curated a large-scale dataset of Performance-Improving Edits (PIE). PIE contains trajectories of programs, where a programmer begins with an initial, slower version and iteratively makes changes to improve the program's performance. We use PIE to fine-tune multiple variants of CODEGEN, a billion-scale Transformer-decoder model. Additionally, we use examples from PIE to prompt OpenAI's CODEX using a few-shot prompting. Our results show that both CODEX and CODEGEN can generate performance-improving edits, with speedups of more than 2.5x for over 25% of the programs. Moreover, PIE allows CODEGEN, an open-sourced and 10x smaller model than CODEX, to match the performance of CODEX on this challenging task.

Dataset

Slow-Fast Pairs

Dataset of slow-fast pairs (tsv) is located here.

Columns description:

user_id: user id
problem_id: problem id. Details about the problems can be found in data/problem_list.csv
language: programming language
submission_id_v0: submission id of the first version of the code
submission_id_v1: submission id of the improved version of the code
cpu_time_v0: cpu time of the first version of the code
cpu_time_v1: cpu time of the second version of the code. cpu_time_v0 > cpu_time_v1 by at least 1% for all the pairs in the dataset. For pairs where the first version was TLE, cpu_time_v0 is set to some high value (e.g. 1000000).
memory_v{0,1}: memory used by the code in the two versions. We can also use memory_v0 > memory_v1 to filter out pairs.
status_v{0,1}: status of the code in the two versions. status_v0 can be "Accepted" or "Time Limit Exceeded", but status_v1 is always "Accepted".
improvement_frac: percentage of improvement of the second version of the code with respect to the first version. improvement_frac is always > 0.

Train/Test/Val Splits for Python and C++

We also provide the splits used in the paper. The splits are in jsonl format.

The train/test/val files contain additional information like the diff. The slower program is the input and the faster program is the target.

Main Results

BibTeX


@article{madaan2023learning,
    title={Learning Performance-Improving Code Edits},
    author={Madaan, Aman and Shypula, Alexander and Alon, Uri and Hashemi, Milad and Ranganathan, Parthasarathy and Yang, Yiming and Neubig, Graham and Yazdanbakhsh, Amir},
    journal={arXiv preprint arXiv:2302.07867},
    year={2023}
}