[Paper Review] MLGO: a Machine Learning Guided Compiler Optimizations Framework

Mircea Trofin, Yundi Qian+4 more|arXiv (Cornell University)|Jan 13, 2021

Machine Learning and Data Classification25 references35 citations

TL;DR

MLGO integrates machine learning into the LLVM compiler to replace inlining-for-size heuristics, using RL and Evolution Strategies, achieving up to 7% native size reduction and good generalization across targets.

ABSTRACT

Leveraging machine-learning (ML) techniques for compiler optimizations has been widely studied and explored in academia. However, the adoption of ML in general-purpose, industry strength compilers has yet to happen. We propose MLGO, a framework for integrating ML techniques systematically in an industrial compiler -- LLVM. As a case study, we present the details and results of replacing the heuristics-based inlining-for-size optimization in LLVM with machine learned models. To the best of our knowledge, this work is the first full integration of ML in a complex compiler pass in a real-world setting. It is available in the main LLVM repository. We use two different ML algorithms: Policy Gradient and Evolution Strategies, to train the inlining-for-size model, and achieve up to 7\% size reduction, when compared to state of the art LLVM -Oz. The same model, trained on one corpus, generalizes well to a diversity of real-world targets, as well as to the same set of targets after months of active development. This property of the trained models is beneficial to deploy ML techniques in real-world settings.

Motivation & Objective

• Motivate leveraging machine learning to replace hand-crafted compiler optimization heuristics in industry-strength compilers.
• Demonstrate a full integration of ML in LLVM via a pilot inlining-for-size optimization pass.
• Develop and compare ML training approaches (reinforcement learning and evolution strategies) for policy optimization.
• Show that a model trained on one corpus generalizes to diverse real-world targets and remains robust over time.

Proposed method

• Formulate inlining-for-size as a reinforcement learning problem with a sequential decision process over call sites.
• Train policies using policy gradient (PPO) and evolution strategies to maximize native size reduction.
• Introduce a reduced feature space of 11 numerical features to make state representation tractable.
• Warm-start RL with behavioral cloning to imitate LLVM's existing inlining heuristics and accelerate training.
• Embed trained ML policies into LLVM as release-mode models and support development-mode experimentation.

Experimental results

Research questions

• RQ1 Can ML-based policies outperform the hand-tuned LLVM inlining-for-size heuristics in a real-world LLVM pipeline?
• RQ2 Do policies trained on one corpus generalize to other code bases and to future revisions of the same targets without retraining?
• RQ3 What are the trade-offs between reinforcement learning and evolution strategies for training inlining policies in a compiler context?

Key findings

• ML policies achieved up to 7% size reduction compared to state-of-the-art LLVM Oz in the pilot inlining-for-size problem.
• A single model trained on one corpus generalizes well across diverse real-world targets and across months of active development.
• Two training paradigms were explored: reinforcement learning (policy gradient with PPO) and evolution strategies, each with distinct trade-offs.
• Warmstarting RL with behavioral cloning accelerates training by leveraging the existing LLVM inlining heuristics.
• The framework separates policy training from production usage, embedding fixed policies into the compiler for deterministic, release-mode operation.