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

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

Author: boomanaiden154

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(100));
 
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TFLITE
@@ -364,6 +364,22 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 
   using RegID = unsigned;
   mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
+
+  mutable std::unordered_map<unsigned, unsigned> VirtRegEvictionCounts;
+
+  void onEviction(Register RegBeingEvicted) const {
+    // If we cannot find the virtual register in the map, we just assume it has
+    // not been evicted before and thus has a value of zero (which is what the
+    // subscript operator returns by default).
+    ++VirtRegEvictionCounts[RegBeingEvicted.id()];
+  }
+
+  unsigned getEvictionCount(Register Reg) const {
+    auto EvictionCountIt = VirtRegEvictionCounts.find(Reg.id());
+    if (EvictionCountIt != VirtRegEvictionCounts.end())
+      return EvictionCountIt->second;
+    return 0;
+  }
 };
 
 #define _DECL_FEATURES(type, name, shape, _)                                   \
@@ -657,7 +673,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 (getEvictionCount(Intf->reg()) > MaxEvictionCount && !Urgent)
         return false;
 
       // Only evict older cascades or live ranges without a cascade.
@@ -803,6 +819,22 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   }
   assert(CandidatePos < ValidPosLimit);
   (void)ValidPosLimit;
+
+  // 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());
+      }
+    }
+  }
+
   return Regs[CandidatePos].first;
 }