MLGO Framework incorporates ML into compiler optimizations

The entry of modern computers gave rise to the need for compiling faster and smaller code. A well-optimized code can decrease the operational cost of big datacenter applications to a greater extent. The size of compiled code is especially important in mobile and embedded systems, as well as software installed on secure boot partitions, where the compiled binary must fit within strict code size budgets.

With advancements in the field, the headroom has been severely constrained by increasingly complex heuristics, hampering maintenance and further enhancements. Heuristics are algorithms that produce empirically reasonable optimal results for difficult problems within pragmatic constraints (For instance, “reasonably fast”).

Heuristics are widely used in compiler optimization passes, including those that use profile feedback, like inlining and register allocation. Passes of this type have a significant impact on the performance of a wide range of programs. These are frequently NP-hard problems, and finding optimal solutions may necessitate exponentially increasing amounts of time or memory.

Recent research has revealed that machine learning can help explore wider opportunities in compiler optimization with the substitution of complex heuristics with ML policies. In real code, the compiler traverses a massive call graph during the inlining phase because thousands of functions call each other. This operation is performed on all caller-callee pairs, and the compiler decides whether or not to inline a caller-callee pair. This is a sequential decision process because previous inlining decisions will change the call graph, influencing subsequent decisions and the final outcome. However, the adoption of ML in general-purpose, industry-strength compilers remains a challenge.

To overcome this challenge, MLGO: a Machine Learning Guided Compiler Optimizations Framework is introduced. This is the first industrial-grade general framework for systematically integrating machine-learning (ML) techniques in a compiler, specifically in LLVM, an open-source industrial compiler infrastructure used to build mission-critical, high-performance software.

More about MLGO

MLGO trains neural networks using reinforcement learning (RL) to make decisions that can replace heuristics in LLVM. Only two types of optimizations are supported by the MLGO framework: inline-for-size and register-allocations-for-performance.

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

This framework is still being researched. Google says its future goals are to increase the number of optimizations and use better machine learning algorithms.