PaperlayerView exact revision on arXiv

Source-linked AI summary

MLGO: a Machine Learning Guided Compiler Optimizations Framework

Mircea Trofin, Yundi Qian, Eugene Brevdo, Zinan Lin, Krzysztof Choromanski, David Li

arXiv:2101.04808v1cs.PLcs.LG

TL;DR

Machine learning had shown promise for compiler optimization, but had not been adopted in production compilers. MLGO applies learned policies to LLVM’s inlining-for-size optimization; policies trained on one software application reduced native size to some extent across other applications and SPEC 2006.

  • Problem

    Machine-learning compiler optimization had shown promise, but its adoption in production compilers remained limited, and real-world deployment posed design challenges.

    Source: §1 · p. 1

  • Method

    MLGO provides a framework for replacing LLVM optimization heuristics with learned policies, and the inlining-for-size pilot trains policies using Policy Gradient and Evolution Strategies.

    Source: §3 · p. 3, §6.2 · p. 9

  • Results

    All three policies trained on the search application reduced native size to some extent across internal applications, Clang, and SPEC 2006.

    Source: §6.2.1 · p. 9

  • Takeaways & Limitations

    The authors are applying MLGO’s principles to inlining for speed and register allocation, with the aim of refining a framework compiler engineers can use broadly.

    Source: §9 · p. 11

    The authors had less experience with policy improvement and ship-blocker workflows, and derived their direction for these from other domains.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Although machine learning has shown promise for compiler optimization, it had not been adopted in production compilers. MLGO pilots a learned replacement for LLVM’s inlining-for-size heuristic, separating policy training from day-to-day compiler use.

1 Introduction

  • Machine learning may scale beyond human-manageable benchmark sets and manually maintained feature combinations, but evolving ML techniques requires different practices.The authors seek policies that generalize well without frequent retraining or online training in production.Source: §1 · p. 1
  • MLGO targets the production-adoption gap with a pilot designed to inform how machine learning can be framed and deployed in an industrial compiler.The pilot was not intended to advance the state of the art beyond the tip-of-tree production compiler.Source: §1 · p. 1
  • The pilot replaces LLVM’s inlining-for-size decision heuristic, a challenging optimization with real-world importance, including for mobile software.The authors chose size because it is relatively easy to measure and non-noisy, aiding rapid prototyping.Source: §1 · p. 1
  • MLGO keeps production compiler use separate from policy training: the compiler embeds a trained model, while training iteratively updates policies from inliner logs.Training uses a large, representative corpus of intermediate-representation modules.Source: §1 · p. 2

2 Background

The background motivates trial-and-error learning for compiler decisions without known optimal labels and describes LLVM’s layered inlining heuristics. The pilot replaces that manual decision process with learned models.

2.1 Machine Learning Techniques for Replacing Compiler Optimization Heuristics

  • Reinforcement learning fits compiler heuristics because optimal decisions are unlabeled, while repeated compilation allows strategies to be explored and improved through experience.In this framing, the compiler is the agent and compilation is its environment.Source: §2.1 · p. 2
  • Evolution Strategies are black-box optimization methods that, like reinforcement learning, improve strategies through trial and error.The authors tried ES as an alternative for compiler optimization problems.Source: §2.1 · p. 2

2.2 High-level Overview of the Current LLVM

  • LLVM’s inliner processes call-graph strongly connected components bottom-up, then considers newly added call sites top-down and runs cleanup optimizations.Those cleanups can affect later inlining decisions.Source: §2.2 · p. 2
  • LLVM estimates post-inline cost by simulating cleanup passes, then compares that cost with a threshold influenced by hotness and inline attributes.Single-block callees or high SIMD instruction shares can receive bonuses; some decisions are deferred to gain more context.Source: §2.2 · p. 3
  • LLVM replaces its years-tuned manual inlining decision process with machine-learned models.The decision is one component among several inlining heuristics, including traversal and cleanup choices.Source: §2.2 · p. 3

3 MLGO

MLGO organizes machine-learning replacement of compiler heuristics around compiler-user and compiler-engineer needs. It separates fixed-policy production use from iterative policy development, while acknowledging less experience with some engineering workflows.

3 MLGO

  • MLGO provides guidelines and requirements for applying machine learning to replace manual compiler-optimization heuristics.The guidelines derive from the needs and scenarios of compiler users and engineers.Source: §3 · p. 3

3.1 Persona: Compiler User

  • Compiler users expect correct, performant code and unchanged determinism, build timeliness, and build-and-release infrastructure.The stated goal is to introduce no changes to the compiler user.Source: §3.1 · p. 3
  • Production policies are trained offline, embedded as fixed native code, and run by inference, avoiding training during compilation and changes to user release pipelines.The framework also requires policies to generalize across code bases and changes without frequent retraining.Source: §3.1 · p. 3

3.2 Persona: Compiler Engineer

  • MLGO forgoes human comprehensibility of learned policies, and the authors have less experience with policy improvement and ship-blocker workflows.Their proposed direction for those workflows draws partly on experience in other domains.Source: §3.2 · p. 4
  • For policy improvement, MLGO proposes adding regression cases to the training corpus, then investigating missing features or alternative algorithms if retraining fails.The feature-diagnosis step retains a manual role for compiler engineers.Source: §3.2 · p. 4
  • For ship blockers, engineers can revert to an earlier policy or manual heuristics, and may locally train a specialized policy for a pathological case.The authors note that local ML training infrastructure is a trade-off compared with changing heuristic flags.Source: §3.2 · p. 4

4 MLGO Policy Training

This section frames inlining decisions as sequential policy learning and describes policy-gradient training, practical adaptations, and a behavioral-cloning warmstart.

4.1.1 RL Problem Formulation

  • Inlining is formulated as sequential decisions at call sites, with the changing call graph as state, inline/no-inline as actions, and native-size reduction as reward.The stated reward is impractical to measure during inlining because native size is not yet known.Source: §4.1.1 · p. 5

4.1.2 Policy Gradient Algorithm

  • Policy Gradient updates the policy toward higher expected reward using Monte Carlo trajectories, alternating compilation under the current policy with parameter updates.The training loop collects fresh trajectories before each update.Source: §4.1.2 · p. 5

4.1.3 Challenges and Implementation Details

  • The 11-feature state reduces call-graph input; total module reward avoids unavailable per-step native sizes, but needs more data and may lower final quality.The reduced state also does not support a value network that predicts future returns.Source: §4.1.3 · p. 6

4.1.4 Warmstart with Behavioral Cloning Policy

  • Behavioral cloning trains a neural network to imitate LLVM’s heuristic decisions, providing a reasonably performing warmstart policy for faster RL training.It uses the RL features as inputs and heuristic inlining decisions as labels.Source: §4.1.4 · p. 6

4.2 Evolution Strategies

Evolution Strategies trains an inlining policy as a black-box optimization problem, using total reward to estimate updates to policy parameters.

4.2 Evolution Strategies

  • Evolution Strategies treats policy parameters as the optimization variables and total native-size reduction under a policy as its black-box reward.The reward is evaluated for a module after inlining.Source: §4.2 · p. 6
  • ES optimizes a smoothed objective and uses a Monte Carlo gradient estimate from sampled Gaussian perturbations to update policy parameters.The perturbation vectors are sampled from a zero-mean, identity-covariance multivariate normal distribution.Source: §4.2 · pp. 6–7
  • Each iteration samples perturbations, compiles the perturbed policies to obtain rewards, and updates parameters, alternating reward collection with policy updates.The algorithm is described as iterating between data collection and policy update, similarly to Policy Gradient.Source: §4.2 · p. 7

4.3 PG v.s. ES

The comparison highlights ES’s simpler data and problem-structure requirements, while Policy Gradient is substantially more sample-efficient in the inlining-for-size task.

Complexity

  • ES needs only total reward, avoiding trajectory logging and easing use when an optimization can be scored but not formulated as an MDP.This reduces engineering complexity and storage/network requirements.Source: §4.3 · p. 7

Sample Efficiency

  • Over 20X computational resources are needed to train an ES policy of similar quality to a Policy Gradient policy.The authors expect the sample-efficiency gap to be even larger when partial rewards are available after each decision.Source: §4.3 · p. 7

4.4 Training Infrastructure

Training alternates between policy-guided compilation and policy updates over an IR corpus. A parallel data collector produces the compilation traces and native-size rewards used for training.

4.4 Training Infrastructure

  • Each iteration sends the current policy to a data collector, which samples IR files, compiles them, and returns training data for policy updates.After several iterations, the trainer exports the trained policy.Source: §4.4 · p. 7
  • Parallel data collection improves training efficiency because collection is the inlining problem’s training bottleneck.The collector takes pre-inlining IR and an optional policy, performs inlining and subsequent optimization, then compiles to native code and measures module size.Source: §4.4 · p. 7
  • Without a supplied policy, the collector uses heuristic inlining and logs its trace; ES training needs only the native-size reward, not that log.The heuristic trace can support behavioral-cloning warmstarts, while heuristic inlining also serves as a baseline.Source: §4.4 · p. 7

5 LLVM Implementation

MLGO adds reusable LLVM support for learned inlining policies, separating model evaluation and feature extraction from the inliner. Release and development modes support deployment and training workflows, respectively.

5 LLVM Implementation

  • LLVM 8 support includes release and development modes plus continuous-integration build bots; development mode requires the TensorFlow C library at run time.Release mode instead requires the SavedModel compiler during the LLVM build.Source: §5 · p. 8, §5.2 · p. 8, §5.3 · p. 9

5.1 RL Driven Inlining (for Size) in LLVM

  • The InlineAdvisor abstraction selects heuristic or learned inlining, while a persistent module analysis tracks module-wide features across the inlining pipeline.FunctionPropertiesAnalysis extracts call-site and module-wide features, including call-graph position and estimated instructions removed.Source: §5.1 · p. 8
  • Both modes exchange fixed-sized input and output tensors with the model, allowing its internal architecture to change without changing that interface.LLVM handles user directives and correctness without model evaluation; a hard instruction-growth threshold guards against unexpectedly expensive downstream processing.Source: §5.1 · p. 8

5.2 Release Mode Implementation in LLVM

  • Release-mode model integration uses a reusable build mechanism that can also support other RL-driven optimization passes.A build rule compiles the SavedModel into native code for LLVM; model weights are checked in as source.Source: §5.2 · p. 8

5.3 Development Mode Implementation in LLVM

  • Development mode loads models through a reusable C++ wrapper and produces training traces of observed features and decisions.Its logs support both heuristic-policy warmstarts and the policy under training; exploratory randomness must be removed and effectiveness revalidated before shipping.Source: §5.3 · pp. 8–9

6 Evaluation

The evaluation measures compiler overhead, policy training and size results, and generalization across software and time. The trained policies reduce native size on diverse software and retain wins over current -Oz after four months of active development.

6.1 Compilation Overhead

  • On a 33MB IR module, release-mode inlining time increased 10%; because inlining is 10-15% of total compile time, its net contribution was 1%.Runtime memory utilization increased 0.65%, and the compiled model added 115KB, or 0.08%, to clang binary size.Source: §6.1 · p. 9
  • The development-mode overhead was not formally measured; training data collection processed 26K IR modules in parallel in around 10 minutes on a 72-thread machine.The run stayed below half of the available 192GB RAM.Source: §6.1 · p. 9

6.2 Inlining for Size Results

  • The policies were trained on an internal search application containing over 28,000 IR modules with varied code patterns.The authors identify this pattern diversity as relevant to generalization across time and software domains.Source: §6.2 · p. 9
  • PG consumes ~5% training resources of ES (100 *12 v.s. 488 *60), while larger neural networks may yield better policies at greater training cost.The paper also reports that a reasonable PG policy can be trained on a 72-thread machine within a day.Source: §6.2 · p. 9, Table 2 · p. 9

6.2.1 Generalizability across Software

  • All three trained policies reduce native size to some extent across tested software, with ES(L)>PG>ES for most software.The evaluation applies policies trained on an internal search application to three other internal applications and Clang; Figure 5 reports SPEC 2006 results.Source: §6.2.1 · p. 9, Figure 4 · p. 10

6.2.2 Generalizability across Time

  • After four months of active development, the policies’ effectiveness may degrade somewhat, but they still achieve decent wins over current -Oz.The same three policies were evaluated on the same software as Figure 4, using an LLVM version from four months later.Source: §6.2.2 · p. 10

7 Related Work

Prior research explored machine learning for compiler optimization, but this work focuses on design constraints for production deployment. Related approaches address distinct tasks or use different forms of integration.

7 Related Work

  • The paper frames its contribution as identifying problem framing and design constraints for applying ML compiler optimizations in production.It situates this work among earlier academic efforts to replace hand-crafted compiler heuristics with machine learning.Source: §7 · p. 10
  • Prior compiler work includes supervised prediction of unrolling, instruction scheduling, and mapping decisions, as well as ML-based scheduling and TensorFlow cost models [7].Other explored areas include instruction scheduling in straight-line code.Source: §7 · p. 10
  • Some related work applies ML to memory prefetching or stochastic-search super-optimization [5], rather than replacing inliner heuristics.The paper also distinguishes pragma-based vectorization guidance, which is a prebuild step and not transparently deployable to compiler users.Source: §7 · p. 10
  • Milepost GCC predicts useful compiler flags but does not use ML-trained policies in its implementation.Inlining-specific precedents include databases of observed decisions and parameterized heuristics.Source: §7 · p. 10

8 Future Directions

The future directions address reward design for speed optimization and several extensions to ML techniques. They favor task-specific reward approximations and identify richer state representations and partial-reward policy gradients as research directions.

8.1 Applying MLGO to Speed Optimizations

  • Benchmark-run rewards are slower and noisier than size measurements, so speed optimization favors avoiding them in favor of problem-specific approximations.For inlining for speed, the proposed approximation combines a per-critical-call-graph working-set estimate with dynamic instruction count; profiling is assumed.Source: §8.1 · pp. 10–11
  • Register allocation could use a block-frequency-weighted sum of introduced moves as its reward, while speed inlining could combine working-set estimates and dynamic instruction count.Both approaches require profiling information, assumed for workloads concerned with speed.Source: §8.1 · p. 11

8.2 ML Techniques

  • Richer state representations could use code embeddings for callers and callees or graph neural networks on the call-site neighborhood.Source: §8.2 · p. 11
  • Partial-reward policy gradients could improve sample efficiency and trainability, but require encoding global call-graph information and predicting function native size from IR.Source: §8.2 · p. 11

9 Conclusion

The conclusion presents MLGO as a framework for applying machine learning to compiler optimization in real-world settings, with inlining-for-size as its pilot. It describes ongoing applications to speed inlining and register allocation, and aims to refine the framework for broader use by compiler engineers.

9 Conclusion

  • The work frames real-world compiler optimization with ML, proposes MLGO, and applies it to inlining-for-size with an LLVM build-time opt-in implementation, training methodology, and two algorithms.Source: §9 · p. 11
  • The authors are applying the same principles to speed inlining and register-allocation policies, aiming to refine MLGO for broader use by compiler engineers.Source: §9 · p. 11
Loading 2101.04808v1…