Learning Performance-Improving Code Edits
LEARNING PERFORMANCE -IMPROVING CODE EDITS
Aman Madaan1, Alexander Shypula2, Uri Alon1
Learning Performance-Improving Code Edits
The waning of Moore's Law has shifted the focus of the tech industry towards
alternative methods for continued performance gains. While optimizing compilers
are a standard tool to help increase program efficiency, programmers continue
to shoulder much responsibility in crafting and refactoring 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. We investigate these
questions by curating 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 evaluate and improve the capacity of large language
models. Specifically, use examples from 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. By
leveraging PIE, we find that both CODEX and CODEGEN can generate
performance-improving edits, with speedups of more than 2.5x for over 25% of
the programs, for C++ and Python, even after the C++ programs were compiled
using the O3 optimization level. Crucially, we show that PIE allows CODEGEN, an
open-sourced and 10x smaller model than CODEX, to match the performance of
CODEX on this challenging task. Overall, this work opens new doors for creating
systems and methods that can help programmers write efficient code.
Aman Madaan Alexander Shypula Uri Alon Milad Hashemi Parthasarathy Ranganathan Yiming Yang Graham Neubig Amir Yazdanbakhsh