MLGO RegAlloc Priority: What's the plan for higher-dim features?

[FENGJIAOLEE]

I've been experimenting with MLGO for register allocation in our compiler,
and wanted to share and ask a few questions:

✅ Eviction policy: Trained and integrated, but saw limited end-to-end gains.
Hypothesis: eviction only triggers on a subset of functions in our workloads,
so the model has fewer opportunities to influence the final code.

🎯 Priority policy: This seems more promising since it applies to nearly every
function. I've patched a few things locally to unblock testing, but I noticed
the current LLVM integration (main branch) only supports 3D features
for the priority policy. Two quick questions if anyone has context:

• Is the priority policy integration still evolving on the LLVM side, are there
  plans to support higher-dimensional features?
• Any rough timeline for when such support might land in main?
  (Totally understand if it's uncertain – just trying to plan our integration efforts)

Thanks for the great work on MLGO! 🙏

[mtrofin]

Re. eviction policy, what did you use as reward function? The current RegAllocScore is a very poor proxy for out of order architectures. If you can run representative benchmarks (emphasis on representative), that would be the most practical current alternative. @boomanaiden154 has a trace-based offline alternative in the works, but it's quite far from being upstreamed

Re. priority policy: we aren't actively working on that at the moment, so if you wanted to evolve it, it'd be a great contribution to the llvm repo. In fact, the whole hope was that folks would apply the same methodology to other passes and that way we can improve both compiler and the methodology.

[FENGJIAOLEE]

Re. eviction policy: Thanks for the hints! I have moved the score calc after postra-scheduler and fine-tuned for our backend. In our tests, it seems to perform better than the default score. – nothing fancy, just better aligned with our backend.

Re. priority policy: I'd love to help evolve this, but due to company policy, I'm not able to share code or detailed internal progress.

One more question: is there a paper or design doc explaining the eviction feature extraction design? The main MLGO paper focuses on inline, so I'd like to understand the rationale behind the eviction feature choices.

[mtrofin]

The paper was about how to stick ML in an industrial compiler, inlining and feature extraction were really less important, especially since all that changes (esp. features). Same here - we took inspiration from the features that went into computing the LR weights, plus some stats about the profile. Nothing more insightful than that. The whole point is to experiment with more features and evolve it, but as we discovered we needed to focus more on the reward signal, didn't spend much time on evolving the features.

Re. priority policy: I'd love to help evolve this, but due to company policy, I'm not able to share code or detailed internal progress.

Ack - it is what it is :) Since there was some refactoring done to set up an advisor design for the priority policy, it should be possible to have a local (to your repo) priority advisor, and if there's a renewed upstream effort here, well - you'll have 2 implementations to pick from :) Low probability of (at least) our group doing that in the short term (i.e. I don't think you risk spending time doing it in your repo, just to find out we're doing the same thing, too).

[FENGJIAOLEE]

Thanks for the details! RL reward design is still new to me .I'm in learning mode for now. Looking forward to seeing how this evolves in MLGO.

We'll give the priority policy a shot locally – good to know we won't collide with upstream anytime soon.