In plain words: A language model reads unoptimized assembly and picks compiler flags that shrink the code; while training it also predicts instruction counts and writes optimized code, which sharpens its understanding. It cut instruction counts 3.0% more than the compiler, beating baselines needing thousands of compilations.
Abstract
We explore the novel application of Large Language Models to code optimization. We present a 7B-parameter transformer model trained from scratch to optimize LLVM assembly for code size. The model takes as input unoptimized assembly and outputs a list of compiler options to best optimize the program. Crucially, during training, we ask the model to predict the instruction counts before and after optimization, and the optimized code itself. These auxiliary learning tasks significantly improve the optimization performance of the model and improve the model's depth of understanding. We evaluate on a large suite of test programs. Our approach achieves a 3.0% improvement in reducing instruction counts over the compiler, outperforming two state-of-the-art baselines that require thousands of compilations. Furthermore, the model shows surprisingly strong code reasoning abilities, generating compilable code 91% of the time and perfectly emulating the output of the compiler 70% of the time.
Chris Cummins, Volker Seeker, Dejan Grubisic, Mostafa Elhoushi, Youwei Liang, Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Kim Hazelwood, Gabriel Synnaeve, Hugh Leather
arXiv:2309.07062 · cs.PL, cs.AI, cs.CL, cs.LG · submitted Sep 11, 2023
abstract · pdf · html
Writing code to unroll a loop is trivial. The limitations of compilers are that almost all currently existing languages are too low level for optimization. ML has the potential to extract back this lost information. Basically the opposite of lowering an IR.