The Need for Speed

These days, more attention is usually given to producing clean, well-structured source code than is given to optimizing the performance or size of that code. The logic behind decisions to this effect is that abundant processing and storage resources are available, so it makes more sense to focus on writing maintainable code than on saving processing cycles. In many cases, this course of action can be wise — poorly written code can be costly to update and is prone to exhibiting unexpected behavior. But there are still many problems that are just at the edge, or even a bit beyond, the capabilities of current computing technologies. Furthermore, a particular use case may necessitate employing a heavily resource-constrained computing device. For situations such as these, compiling faster, smaller code is essential.

Fortunately, there are still research groups giving this area the consideration that it deserves. Engineers at Google AI have recently reported the results of their efforts to use machine learning to help compilers to optimize code for both performance and size. Much of the optimization that is done by modern compilers involves the use of heuristics, which are now impeding maintenance and further improvements due to their rapidly increasing complexity. The team at Google AI asked if these heuristics could be replaced, and improved upon, through the use of learning algorithms.

Their solution, called MLGO, integrated machine learning algorithms into LLVM, an open-source industrial compiler infrastructure. Reinforcement learning was leveraged by MLGO to train artificial neural networks how to make optimization decisions that can replace LLVM’s heuristics. The technique was demonstrated under two conditions: by using inlining to reduce code size, and by using register allocation to improve execution speed.

Inlining uses heuristics to reduce code size by removing redundant code. A traditional compiler traces through the code’s call graph — a map of the calls that functions make to other functions — to collapse and simplify the code, without changing the function. As you can imagine, the heuristics required to achieve this goal have grown very complex over time — so much so that it is now very difficult to build in any further improvements. To solve this problem, the reinforcement learning algorithm was allowed to iterate through possible solutions to optimal inlining. In the process, it was able to learn a method that can reduce code size by 7% when compared with LLVM’s heuristics, and it was also able to do away with the complexity and lack of maintainability that those heuristics introduce.

In their second demonstration, the team took on the problem of assigning variables to physical registers. Since data must be moved to a register in order to perform many operations on it, it is essential that the contents of variables be moved to the appropriate register the moment that it is freed up. Idle time resulting from poor register allocation algorithms will cause the running code to suffer a performance hit. By again taking a reinforcement learning approach, an execution speed improvement of up to 1.5% was observed as compared with traditional heuristics-based approaches.

These initial results are quite encouraging. Improving the performance of existing hardware has the potential to allow us to solve problems that were previously out of reach. Google AI has open-sourced their code, and you can check it out on GitHub.