[MLGO] Add Threshold to Prevent Pathological Compile Time Cases

[boomanaiden154]

This patch adds a threshold flag, -mlregalloc-max-cascade, to prevent live ranges from being evicted more than is necessary.

After deploying a new regalloc model, we ran into some pathological cases where the model decided it wanted to ping-pong evictions, taking up a large amount of compile time. This threshold is mostly a stop gap while we continue to investigate other solutions and work on minimizing/constructing test cases.

[llvmbot]

@llvm/pr-subscribers-llvm-regalloc

@llvm/pr-subscribers-mlgo

Author: Aiden Grossman (boomanaiden154)

Changes

This patch adds a threshold flag, -mlregalloc-max-cascade, to prevent live ranges from being evicted more than is necessary.

After deploying a new regalloc model, we ran into some pathological cases where the model decided it wanted to ping-pong evictions, taking up a large amount of compile time. This threshold is mostly a stop gap while we continue to investigate other solutions and work on minimizing/constructing test cases.

---

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

1 Files Affected:

• (modified) llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp (+15-1)

diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 03f015f8c9e32d..e202dd45cfed50 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -63,6 +63,12 @@ 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));
+
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TFLITE
 #include "RegAllocScore.h"
@@ -643,8 +649,16 @@ bool MLEvictAdvisor::loadInterferenceFeatures(
            RegClassInfo.getNumAllocatableRegs(MRI->getRegClass(VirtReg.reg())) <
                RegClassInfo.getNumAllocatableRegs(
                    MRI->getRegClass(Intf->reg())));
-      // Only evict older cascades or live ranges without a cascade.
+
       unsigned IntfCascade = RA.getExtraInfo().getCascade(Intf->reg());
+      // There is a potential that the model could be adversarial and
+      // continually evict live ranges over and over again, leading to a
+      // large amount of compile time being spent in regalloc. If we hit the
+      // threshold, prevent the range from being evicted.
+      if (IntfCascade >= MaxCascade)
+        return false;
+
+      // Only evict older cascades or live ranges without a cascade.
       if (Cascade <= IntfCascade) {
         if (!Urgent)
           return false;