[llvm] f1dbc45 - [MLGO] Properly Handle Counting Evictions of Candidates

[Aiden Grossman via llvm-commits]

Author: Aiden Grossman
Date: 2025-03-05T07:19:50Z
New Revision: f1dbc45210cec766bed7dd320ed9420484ac3ec6

URL: https://github.com/llvm/llvm-project/commit/f1dbc45210cec766bed7dd320ed9420484ac3ec6
DIFF: https://github.com/llvm/llvm-project/commit/f1dbc45210cec766bed7dd320ed9420484ac3ec6.diff

LOG: [MLGO] Properly Handle Counting Evictions of Candidates

This patch makes it so that onEviction actually gets called when the
model ends up selecting the candidate to evict. Where we were handling
this previously ended up being dead code as we would return earlier with
MCRegister::NoRegister.

Fixes #129841.

Added: 
    

Modified: 
    llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 8215d07b74dba..1cde094d78e23 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -860,6 +860,7 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   // Regs[CandidatePos].second)
   assert(Regs[CandidatePos].second);
   if (CandidatePos == CandidateVirtRegPos) {
+    onEviction(VirtReg.reg());
     assert(!MustFindEviction);
     return MCRegister::NoRegister;
   }
@@ -869,15 +870,11 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   // Update information about how many times the virtual registers being
   // evicted have been evicted so that we can prevent the model from evicting
   // the same ranges continually and eating compile time.
-  if (CandidatePos == CandidateVirtRegPos) {
-    onEviction(VirtReg.reg());
-  } else {
-    for (MCRegUnit Unit : TRI->regunits(Regs[CandidatePos].first)) {
-      LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, Unit);
-      const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff);
-      for (const LiveInterval *Intf : reverse(IFIntervals)) {
-        onEviction(Intf->reg());
-      }
+  for (MCRegUnit Unit : TRI->regunits(Regs[CandidatePos].first)) {
+    LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, Unit);
+    const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff);
+    for (const LiveInterval *Intf : reverse(IFIntervals)) {
+      onEviction(Intf->reg());
     }
   }