[RFC] MLGO Regalloc: learned eviction policy for regalloc

[RFC] MLGO for Regalloc: learned eviction policy for regalloc

alphabetical list of authors: Eugene Brevdo, David Li, Yundi Qian, Mircea Trofin

Summary: continuing from our experience with the inliner for size, we were able to learn live range eviction for the ‘greedy’ register allocator, with measurable, positive results on some datacenter-type applications (from 0.3% up to 1.5% improvement in QPS (query per second) on selected internal ThinLTO + FDO apps).

Just like in the inliner case, and salient to our approach, opting in to the learned policy maintains critical compiler invariants, like determinism and timeliness, and does not add new runtime dependencies.


The proposed changes are:

  • the main change (https://reviews.llvm.org/D113525) is a NFC - factoring eviction as an ImmutablePass, RegAllocEvictionAdvisorAnalysis. It can be queried for an implementation of RegAllocEvictionAdvisor. By default, that implementation is today’s behavior.

This design is similar to what we did for the inliner with InlineAdvisor. We expect the factoring to create opportunities for other initiatives, just like it did in the inliner case (see @modimo’s inline replay work, for example).

  • all the ML stuff is an alternative implementation of RegAllocEvictionAdvisor, and needs to be opted in. Size of change - wise, the ML support will be relatively small, as we are leveraging the existing infrastructure we have already introduced with the inliner work (including leveraging the same build bots)

MLGO-specific details:

We plan to first introduce ‘development mode’, i.e. the support for training a policy, then ‘release mode’, i.e. the way in which a pre-trained policy is meant to be leveraged in production. This greatly simplifies integration on our side, and should also allow anyone wishing to train a policy with their training algorithm to do so. Separately, we have been upstreaming the training infrastructure pertinent to the regalloc work to github.com/google/ml-compiler-opt.