MLGO: a Machine Learning Guided Compiler Optimizations Framework
Paper · 2021 · Mircea Trofin, Yundi Qian, Eugene Brevdo et al.
Summary
Academic research has repeatedly shown that machine learning can outperform hand-written compiler heuristics, yet production-grade compilers had not adopted it. The authors present MLGO, a set of design principles and an implementation for bringing learned policies into LLVM without disrupting how developers build and ship software. Training runs offline on a large corpus of IR modules. The instrumented compiler logs features and decisions, and these logs feed the learning algorithm. The finished model is compiled into the compiler as native code, so builds stay deterministic and need no new dependencies. As a pilot, they replaced the inlining-for-size decision heuristic with a neural policy trained by Policy Gradient and by Evolution Strategies. Compared with LLVM -Oz, the learned policy shrinks binaries by up to 7%. A model trained on one corpus also transfers to unrelated software and still helps after four months of changes to both code and compiler. This matters because it suggests ML-driven optimization can be deployed in industry without frequent retraining.
Key points
- MLGO is described as the first complete integration of ML into a complex pass of a real-world compiler. It ships in the main LLVM repository as a build-time opt-in.
- Two training approaches were compared: Policy Gradient (reinforcement learning) and Evolution Strategies (black-box optimization).
- The learned inlining-for-size policy achieves up to 7% size reduction over LLVM -Oz, with results reported on several real-world targets and on SPEC 2006.
- Generalization was tested in two ways. Policies were applied to software outside the training corpus, and they were re-evaluated four months later on newer code and a newer LLVM (Clang selfhost 4ca60915bcc). In the later test the gains shrank somewhat but stayed clearly positive.
- The authors plan to extend MLGO to inlining for speed and to register allocation. For these tasks they intend to use profile-based reward approximations instead of noisy benchmark runs. The code is at github.com/google/ml-compiler-opt.
Institutions
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksInformation Systems · 1,888 worksHardware and Architecture · 183 works
Cite
Cite as Mircea Trofin, Yundi Qian, Eugene Brevdo, Zinan Lin, Krzysztof Choromanski and David Li. 2021. MLGO: a Machine Learning Guided Compiler Optimizations Framework. arXiv:2101.04808.
@misc{trofin2021mlgo,
title = {{MLGO: a Machine Learning Guided Compiler Optimizations Framework}},
author = {Mircea Trofin, Yundi Qian, Eugene Brevdo, Zinan Lin, Krzysztof Choromanski and David Li},
year = {2021},
doi = {10.48550/arxiv.2101.04808},
eprint = {2101.04808},
archivePrefix = {arXiv},
url = {https://doi.org/10.48550/arxiv.2101.04808}
}References
- Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. 2016. Tensorflow: A system for large-scale machine learning. In 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16).…
- Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Tzu-mao Li, Michael Charbi, Benoit Steiner, Steven Johnson, Kayvon Fa-tahlian, Fredo Durand, and Jonathan Regan-kelley. 2019. Learning to optimize Halide with tree search and random programs (ACM Transactions on Graphics). 2–12. find
- Uri Alon, Meital Zilberstein, Omer Levy, and Eran Yahav. 2019. code2vec: Learning distributed representations of code. Proceedings of the ACM on Programming Languages 3, POPL (2019), 1–29. find
- Michael Bain and Claude Sammut. 1995. A Framework for Behavioural Cloning.. In Machine Intelligence 15. 103–129. find
- Rudy Bunel, Alban Desmaison, M Pawan Kumar, Philip HS Torr, and Pushmeet Kohli. 2016. Learning to superoptimize programs. arXiv preprint arXiv:1611.01787 (2016). find
- John Cavazos and Michael FP O’Boyle. 2005. Automatic tuning of inlining heuristics. In SC’05: Proceedings of the 2005 ACM/IEEE Conference on Supercomputing. IEEE, 14–14. find
- Tianqi Chen, Lianmin Zheng, Eddie Q. Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. Learning to Optimize Tensor Programs. CoRR abs/1805.08166 (2018). arXiv:1805.08166 http://arxiv.org/abs/1805.08166 ML Techniques ↗
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/trofin-2021-mlgo-machine-learning-guided-compiler/