Google Releases MLGO Framework


The MLGO framework is trained with the internal Google code base and tested on the code of Fuchsia

Google’s new Machine Learning Guided Optimization (MLGO) is an industrial-grade general framework for systematically integrating machine-learning (ML) techniques in a compiler, particularly in LLVM. This open-source industrial compiler infrastructure is ubiquitous for building mission-critical, high-performance software.

Compiling faster and smaller code can significantly reduce the operational cost of large data centre applications. The compiled code size matters the most to mobile and embedded systems or software deployed on secure boot partitions, where the compiled binary must fit tight code-size budgets.

In a standard compiler, the decisions about the optimisations are made by heuristics, but heuristics become increasingly difficult to improve over time. Heuristics are algorithms that, empirically, produce reasonably optimal results for complex problems, within pragmatic constraints (e.g. “reasonably fast”). In the compiler case, heuristics are widely used in optimisation passes, even those leveraging profile feedback, such as inlining and register allocation. Such passes significantly impact the performance of a wide variety of programs. These problems are often NP-hard and searching for optimal solutions may require exponentially increasing amounts of time or memory.

Recent research has shown that ML can help with these tasks and unlock more code optimisation than complicated heuristics can. In actual code, during the inlining phase, the compiler traverses a huge call graph because thousands of functions call each other. This operation is performed on all caller-callee pairs, and the compiler decides whether to inline a caller-callee pair. This is a sequential decision process as previous inlining decisions will alter the call graph, affecting later decisions and the final result.

Reinforcement learning (RL) is a family of ML techniques that may be applied to find increasingly optimal solutions through an automated iterative exploration and training process. MLGO uses RL to train neural networks to make decisions that can replace heuristics in LLVM. The MLGO framework supports only two kinds of optimisations: inline-for-size and register-allocations-for-performance.

The MLGO framework is trained with the internal Google code base and tested on the code of Fuchsia — a general-purpose, open-source operating system designed to power a diverse ecosystem of hardware and software where binary size is critical. For inline-for-size optimization, MLGO achieves a 3% to 7% size reduction. Similar to the inlining-for-size policy, the register allocation (regalloc-for-performance) policy, MLGO achieves up to 1.5% improvement in queries per second (QPS) on a set of internal large-scale data centre applications.

This framework is still in the research phase. Google says its future goals are to expand the number of optimisations and apply better ML algorithms.