Google’s new Machine Learning Guided Optimization (MLGO) is an industrial-grade general framework for integrating machine-learning (ML) techniques systematically in a compiler and in particular in LLVM, an open-source industrial compiler infrastructure that is ubiquitous for building mission-critical, high-performance software.
The Google AI Blog explored the release.
Compiling faster and smaller code can significantly reduce the operational cost of large datacenter applications. The size of compiled code matters the most to mobile and embedded systems or software deployed on secure boot partitions, where the compiled binary must fit in tight code-size budgets.
In a standard compiler, the decisions about the optimizations are made by heuristics, but heuristics become increasingly difficult to improve over time. Heuristics are algorithms that, empirically, produce reasonably optimal results for hard problems, within pragmatic constraints (e.g. "reasonably fast"). In the compiler case, heuristics are widely used in optimization passes, even those leveraging profile feedback, such as inlining and register allocation. Such passes have a significant impact on the performance of a broad 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 optimization than the complicated heuristics can. In real code, during the inlining phase, the compiler traverses a huge call graph, because there are thousands of functions calling each other. This operation is performed on all caller-callee pairs and the compiler makes decisions on whether to inline a caller-callee pair or not. 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 optimizations: 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 datacenter applications.
This framework is still in a research phase. Google says its future goals are to expand the number of optimizations and apply better ML algorithms.