[MLGO] Count LR Evictions Rather than Relying on Cascade

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,23 @@ 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, _)                                   \
@@ -656,8 +673,8 @@ bool MLEvictAdvisor::loadInterferenceFeatures(
       // large amount of compile time being spent in regalloc. If we hit the
       // 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)
+      // eviction if the candidate is not spillable. 
+      if (getEvictionCount(Intf->reg()) > MaxEvictionCount && !Urgent)
         return false;
 
       // Only evict older cascades or live ranges without a cascade.
@@ -803,6 +820,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;
 }