[llvm] [MSSAUpdater] Replace recursion with worklist and cap it. (PR #150543)

[Alina Sbirlea via llvm-commits]

https://github.com/alinas created https://github.com/llvm/llvm-project/pull/150543

Replace the recursion used for finding and removing trivial MemoryPhis with a worklist approach and add a cap for it for pathological cases.

This came up in a profile that also found #Issue150531, where a lot of time was found on copying the MemoryPhi operands during the recusion.

>From 5c8002cbe47c8cf25585e02a628acbafed8adef3 Mon Sep 17 00:00:00 2001
From: Alina Sbirlea <asbirlea at google.com>
Date: Thu, 24 Jul 2025 22:43:47 +0000
Subject: [PATCH] [MSSAUpdater] Replace recursion with worklist and cap it.

Replace the recursion used for finding and removing trivial MemoryPhis
with a worklist approach and add a cap for it for pathological cases.

This came up in a profile that also found #Issue150531, where a lot of
time was found on copying the MemoryPhi operands during the recusion.
---
 llvm/include/llvm/Analysis/MemorySSAUpdater.h |  1 -
 llvm/lib/Analysis/MemorySSAUpdater.cpp        | 75 +++++++++++++++----
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MemorySSAUpdater.h b/llvm/include/llvm/Analysis/MemorySSAUpdater.h
index 96bf99922d848..8c03f241e8609 100644
--- a/llvm/include/llvm/Analysis/MemorySSAUpdater.h
+++ b/llvm/include/llvm/Analysis/MemorySSAUpdater.h
@@ -257,7 +257,6 @@ class MemorySSAUpdater {
   MemoryAccess *
   getPreviousDefRecursive(BasicBlock *,
                           DenseMap<BasicBlock *, TrackingVH<MemoryAccess>> &);
-  MemoryAccess *recursePhi(MemoryAccess *Phi);
   MemoryAccess *tryRemoveTrivialPhi(MemoryPhi *Phi);
   template <class RangeType>
   MemoryAccess *tryRemoveTrivialPhi(MemoryPhi *Phi, RangeType &Operands);
diff --git a/llvm/lib/Analysis/MemorySSAUpdater.cpp b/llvm/lib/Analysis/MemorySSAUpdater.cpp
index ecfecb03c3757..56743e7e7caca 100644
--- a/llvm/lib/Analysis/MemorySSAUpdater.cpp
+++ b/llvm/lib/Analysis/MemorySSAUpdater.cpp
@@ -18,12 +18,17 @@
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include <algorithm>
 
 #define DEBUG_TYPE "memoryssa"
 using namespace llvm;
 
+static cl::opt<uint32_t> TrivialPhiProcessingLimit(
+    "mssaupdater-processing-limit", cl::Hidden, cl::init(20),
+    cl::desc("The worklist limit for trivial phi removal (default = 20)"));
+
 // This is the marker algorithm from "Simple and Efficient Construction of
 // Static Single Assignment Form"
 // The simple, non-marker algorithm places phi nodes at any join
@@ -181,18 +186,6 @@ MemoryAccess *MemorySSAUpdater::getPreviousDefFromEnd(
 
   return getPreviousDefRecursive(BB, CachedPreviousDef);
 }
-// Recurse over a set of phi uses to eliminate the trivial ones
-MemoryAccess *MemorySSAUpdater::recursePhi(MemoryAccess *Phi) {
-  if (!Phi)
-    return nullptr;
-  TrackingVH<MemoryAccess> Res(Phi);
-  SmallVector<TrackingVH<Value>, 8> Uses;
-  std::copy(Phi->user_begin(), Phi->user_end(), std::back_inserter(Uses));
-  for (auto &U : Uses)
-    if (MemoryPhi *UsePhi = dyn_cast<MemoryPhi>(&*U))
-      tryRemoveTrivialPhi(UsePhi);
-  return Res;
-}
 
 // Eliminate trivial phis
 // Phis are trivial if they are defined either by themselves, or all the same
@@ -230,9 +223,61 @@ MemoryAccess *MemorySSAUpdater::tryRemoveTrivialPhi(MemoryPhi *Phi,
     removeMemoryAccess(Phi);
   }
 
-  // We should only end up recursing in case we replaced something, in which
-  // case, we may have made other Phis trivial.
-  return recursePhi(Same);
+  // Continue traversal in a DFS worklist approach, in case we might find
+  // other trivial Phis.
+  if (!Same)
+    return nullptr;
+
+  TrackingVH<MemoryAccess> Result(Same);
+
+  // Worklist approach to recursively removing trivial Phis.
+  SmallVector<TrackingVH<Value>, 5> Worklist;
+
+  for (auto U : Same->users()) {
+    if (dyn_cast<MemoryPhi>(&*U))
+      Worklist.push_back(TrackingVH<Value>(U));
+  }
+
+  while (!Worklist.empty() || Worklist.size() > TrivialPhiProcessingLimit) {
+    MemoryPhi *RecPhi = dyn_cast<MemoryPhi>(&*(Worklist[Worklist.size() - 1]));
+    Worklist.pop_back();
+
+    if (!RecPhi)
+      continue;
+    auto RecOperands = RecPhi->operands();
+
+    // Repeat above algorithm.
+    if (NonOptPhis.count(RecPhi))
+      continue;
+
+    bool nextone = false;
+    MemoryAccess *RecSame = nullptr;
+    for (auto &Op : RecOperands) {
+      if (Op == RecPhi || Op == RecSame)
+        continue;
+      if (RecSame) {
+        nextone = true;
+        break;
+      }
+      RecSame = cast<MemoryAccess>(&*Op);
+    }
+
+    // Move on to the next Phi in the worklist.
+    if (nextone || RecSame == nullptr)
+      continue;
+
+    if (RecPhi) {
+      RecPhi->replaceAllUsesWith(RecSame);
+      removeMemoryAccess(RecPhi);
+    }
+
+    for (auto U : RecSame->users()) {
+      if (dyn_cast<MemoryPhi>(&*U))
+        Worklist.push_back(TrackingVH<Value>(U));
+    }
+  }
+
+  return Result;
 }
 
 void MemorySSAUpdater::insertUse(MemoryUse *MU, bool RenameUses) {