[llvm] [MLGO] Count LR Evictions Rather than Relying on Cascade (PR #124440)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlgo

Author: Aiden Grossman (boomanaiden154)

<details>
<summary>Changes</summary>

This patch adjusts the mlregalloc-max-cascade flag (renaming it to mlregalloc-max-eviction-count) to actually count evictions rather than just looking at the cascade number. The cascade number is not very representative of how many times a LR has been evicted, which can lead to some problems in certain cases, where we might end up with many eviction problems where we have now masked off all the interferences and are forced to evict the candidate.

This is probably what I should've done in the first place. No test case as this only shows up in quite large functions post ThinLTO and it would be hard to construct something that would serve as a nice regression test without being super brittle. I've tested this on the pathological cases that we have come across so far and it works.

Fixes #<!-- -->122829

---
Full diff: https://github.com/llvm/llvm-project/pull/124440.diff


1 Files Affected:

- (modified) llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp (+23-6) 


``````````diff
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9c6487b40d6061..7a95c36f431e1a 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -63,11 +63,11 @@ static cl::opt<std::string> InteractiveChannelBaseName(
         "outgoing name should be "
         "<regalloc-evict-interactive-channel-base>.out"));
 
-static cl::opt<unsigned>
-    MaxCascade("mlregalloc-max-cascade", cl::Hidden,
-               cl::desc("The maximum number of times a live range can be "
-                        "evicted before preventing it from being evicted"),
-               cl::init(20));
+static cl::opt<unsigned> MaxEvictionCount(
+    "mlregalloc-max-eviction-count", cl::Hidden,
+    cl::desc("The maximum number of times a live range can be "
+             "evicted before preventing it from being evicted"),
+    cl::init(20));
 
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TFLITE
@@ -364,6 +364,8 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 
   using RegID = unsigned;
   mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
+
+  mutable std::unordered_map<unsigned, unsigned> VirtRegEvictionCounts;
 };
 
 #define _DECL_FEATURES(type, name, shape, _)                                   \
@@ -657,7 +659,7 @@ bool MLEvictAdvisor::loadInterferenceFeatures(
       // threshold, prevent the range from being evicted. We still let the
       // range through if it is urgent as we are required to produce an
       // eviction if the candidate is not spillable.
-      if (IntfCascade >= MaxCascade && !Urgent)
+      if (VirtRegEvictionCounts[Intf->reg().id()] > MaxEvictionCount && !Urgent)
         return false;
 
       // Only evict older cascades or live ranges without a cascade.
@@ -803,6 +805,21 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   }
   assert(CandidatePos < ValidPosLimit);
   (void)ValidPosLimit;
+
+  // Update information about how many times the virtual registers being
+  // evicted have been evicted.
+  if (CandidatePos == CandidateVirtRegPos) {
+    VirtRegEvictionCounts[VirtReg.reg()] += 1;
+  } 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)) {
+        VirtRegEvictionCounts[Intf->reg()] += 1;
+      }
+    }
+  }
+
   return Regs[CandidatePos].first;
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/124440