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

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 created https://github.com/llvm/llvm-project/pull/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

>From 8d265219897002b481a6b0ec328c5232d8966351 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Sun, 26 Jan 2025 03:09:21 +0000
Subject: [PATCH] [MLGO] Count LR Evictions Rather than Relying on Cascade

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
---
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 29 ++++++++++++++++-----
 1 file changed, 23 insertions(+), 6 deletions(-)

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;
 }