Meta Large Language Model Compiler: Foundation Models of Compiler Optimization (2407.02524v1)

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

Meta LLM Compiler: Foundation Models of Compiler Optimization

The paper introduces the Meta LLM Compiler (LLM Compiler), a set of sophisticated, pre-trained models specifically optimized for code and compiler optimization. The authors, including Chris Cummins and his colleagues from Meta AI, have addressed the potential of LLMs in the niche but crucial domain of compiler optimizations.

Key Contributions

The key contributions of the paper include:

1. Model Creation: LLM Compiler builds on the Code Llama architecture, extending its capabilities to understand compiler intermediate representations (IRs), assembly language, and optimization techniques. The model was trained on an extensive corpus of 546 billion tokens, predominantly comprising LLVM-IR and assembly code.
2. Model Availability: The authors emphasize the openness of their models, releasing them under a bespoke commercial license to encourage widespread reuse by both academic researchers and industry practitioners. Two model sizes are provided: 7 billion and 13 billion parameters.
3. Evaluation and Results: Fine-tuned versions of the model were evaluated in two specific tasks: optimizing code size and disassembling x86_64 and ARM assembly back to LLVM-IR. The models demonstrated impressive capabilities, achieving 77% of the optimizing potential of an autotuning search and a 45% rate of successful disassembly round-trip, with 14% exact match. These results highlight the practical utility and efficiency of the LLM Compiler.
4. Training Pipeline: The training pipeline is meticulously detailed, presenting a multi-stage process where the models are first pre-trained on unlabelled IRs and assembly code, followed by instruction fine-tuning on compiler optimization tasks and further adaptation for specific downstream tasks like flag tuning and disassembly.

Methodology

The authors provide in-depth methodological details. The training involves multiple stages:

• Pretraining: The model is initially pre-trained on a large-scale dataset of 401 billion tokens comprising LLVM-IR and assembly code. This ensures a comprehensive understanding of compiler-specific languages.
• Instruction Fine-Tuning: Two critical tasks are used for fine-tuning. First, compiler emulation, where the model learns the effects of various optimization passes. Second, specific tasks like optimization flag tuning and disassembly, ensuring the model can generate effective optimization pass lists and reverse-engineer assembly code into IR.

Evaluation

The models were rigorously evaluated on the MiBench benchmark suite, demonstrating superior performance compared to both GPT-4 Turbo and Code Llama - Instruct. For the flag tuning task, the models improved code size reduction by approximately 5% beyond what was achieved with -Oz optimization. The ability to round-trip disassembled code back to assembly with high fidelity (up to 14% exact match) further underscores the practical capabilities of the LLM Compiler.

Implications and Future Directions

Practical Implications: The LLM Compiler has significant implications for optimizing compiler workflows, potentially reducing the need for extensive and time-consuming autotuning. This can result in more efficient compiler pipelines, benefiting both academia and industry.

Theoretical Implications: The training methodologies and fine-tuning strategies employed can inform future research on LLMs for specialized tasks beyond natural language processing. Particularly, the integration of detailed domain-specific knowledge (like compiler optimizations) into LLM frameworks sets a precedence for future models targeting specialized technical applications.

Future Developments: The authors suggest that future work could expand the scope of the LLM Compiler to other aspects of code optimization, such as run-time performance improvements. Additionally, extending context windows and improving the fidelity of disassembly could be areas of focus, addressing some limitations noted in the paper.

Conclusion

The Meta LLM Compiler is a remarkable advancement in the application of LLMs to the domain of compiler optimization, offering both theoretical insights and practical tools for advancing the efficiency and capability of compilers. The detailed training strategies, robust evaluation, and open model access position LLM Compiler as a significant contribution to the fields of AI and compiler technology, providing a strong foundation for ongoing research and development.