[llvm] 00f692b - Reland "[MLGO] Count LR Evictions Rather than Relying on Cascade (#124440)"

[Aiden Grossman via llvm-commits]

Author: Aiden Grossman
Date: 2025-01-28T07:09:25Z
New Revision: 00f692b94f9aa08ede4aaba6f2aafe17857599c4

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

LOG: Reland "[MLGO] Count LR Evictions Rather than Relying on Cascade (#124440)"

This reverts commit aa65f93b71dee8cacb22be1957673c8be6a3ec24.

This relands commit 8cc83b66e20e72cdb3bb5fbd549c941797b0e0c9.

It looks like this was a transitive include issue.

Added: 
    

Modified: 
    llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9c6487b40d6061..9656774c6eaaef 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -42,6 +42,7 @@
 #include <array>
 #include <bitset>
 #include <memory>
+#include <unordered_map>
 
 using namespace llvm;
 
@@ -63,11 +64,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 +365,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 +674,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 +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;
 }