Meta Large Language Model Compiler: Foundation Models of Compiler Optimization

Chris Cummins,Volker Seeker,Dejan Grubisic,Baptiste Roziere,Jonas Gehring,Gabriel Synnaeve,Hugh Leather
2024-06-28
Abstract:Large Language Models (LLMs) have demonstrated remarkable capabilities across a variety of software engineering and coding tasks. However, their application in the domain of code and compiler optimization remains underexplored. Training LLMs is resource-intensive, requiring substantial GPU hours and extensive data collection, which can be prohibitive. To address this gap, we introduce Meta Large Language Model Compiler (LLM Compiler), a suite of robust, openly available, pre-trained models specifically designed for code optimization tasks. Built on the foundation of Code Llama, LLM Compiler enhances the understanding of compiler intermediate representations (IRs), assembly language, and optimization techniques. The model has been trained on a vast corpus of 546 billion tokens of LLVM-IR and assembly code and has undergone instruction fine-tuning to interpret compiler behavior. LLM Compiler is released under a bespoke commercial license to allow wide reuse and is available in two sizes: 7 billion and 13 billion parameters. We also present fine-tuned versions of the model, demonstrating its enhanced capabilities in optimizing code size and disassembling from x86_64 and ARM assembly back into LLVM-IR. These achieve 77% of the optimising potential of an autotuning search, and 45% disassembly round trip (14% exact match). This release aims to provide a scalable, cost-effective foundation for further research and development in compiler optimization by both academic researchers and industry practitioners.
Programming Languages,Artificial Intelligence
What problem does this paper attempt to address?
### Problems the Paper Aims to Solve The paper aims to address the under-explored application of large language models (LLMs) in the field of code and compiler optimization. Although LLMs have demonstrated remarkable capabilities in various software engineering and coding tasks, their application in code and compiler optimization remains less explored. Training LLMs requires substantial computational resources and data collection, which is often unaffordable for researchers. To tackle this issue, the authors introduce the **Meta Large Language Model Compiler (LLM Compiler)**, a suite of pre-trained models specifically designed for code optimization tasks. These models are built on Code Llama and trained on a large dataset of compiler intermediate representations (IR) and assembly code to enhance the understanding of compiler behavior. The LLM Compiler offers two models with different parameter sizes (700 million and 1.3 billion parameters) and is fine-tuned to optimize code size and disassembly tasks. These models exhibit excellent performance in optimization potential and disassembly accuracy, achieving 77% and 45% of the performance of auto-tuning search, respectively. ### Main Contributions 1. **Pre-trained Models**: The LLM Compiler is trained on a corpus of 546 billion compiler IR and assembly code, providing a deep understanding of compiler optimization techniques. 2. **Fine-tuned Versions**: Fine-tuned versions are provided for specific downstream tasks such as optimization flag tuning and disassembly, significantly enhancing the performance of these tasks. 3. **Scalability and Cost-effectiveness**: By offering pre-trained models, the computational and data collection costs for researchers and industry practitioners in training large models are reduced. 4. **Commercial License**: The models are released under a custom commercial license, allowing for broad reuse. ### Method Overview 1. **Pre-training Phase**: The models are pre-trained on a large dataset containing compiler IR and assembly code to build a comprehensive understanding of these languages. 2. **Instruction Fine-tuning**: By generating random optimization sequences and training the model to predict the optimized code and code size, the model's understanding of compiler optimization mechanisms is further enhanced. 3. **Downstream Task Fine-tuning**: Fine-tuning is performed for specific tasks such as optimization flag tuning and disassembly to improve the model's performance on these tasks. ### Experimental Results 1. **Optimization Flag Tuning**: The LLM Compiler FTD performs excellently in the task of optimizing code size, with the 1.3 billion parameter model generating binary files smaller than those optimized with -Oz in 61% of cases. 2. **Disassembly Task**: In the disassembly task, the LLM Compiler FTD correctly disassembles code 14% of the time. 3. **Performance Comparison**: Compared to GPT-4 Turbo and Code Llama - Instruct, the LLM Compiler FTD significantly outperforms other models in the optimization flag tuning task. ### Conclusion The LLM Compiler provides a scalable and cost-effective foundation for the field of compiler optimization, paving the way for academic researchers and industry practitioners to further develop code optimization models. Through pre-training and fine-tuning, these models demonstrate excellent performance in code optimization and disassembly tasks, promising to drive further advancements in this field.