[llvm] [MemorySSA] Make getPreviousDef iterative to avoid stack overflow (PR #205159)

[Antonio Frighetto via llvm-commits]

================
@@ -33,101 +33,232 @@ using namespace llvm;
 // that there are two or more definitions needing to be merged.
 // This still will leave non-minimal form in the case of irreducible control
 // flow, where phi nodes may be in cycles with themselves, but unnecessary.
-MemoryAccess *MemorySSAUpdater::getPreviousDefRecursive(
+//
+// The predecessor walk is driven by an explicit worklist rather than native
+// recursion: it used to be mutually recursive with getPreviousDefFromEnd, and
+// the recursion depth scaled with the length of the walk, so deep CFGs (e.g.
+// long block chains in large generated kernels/shaders) could overflow the
+// native stack. Each StackFrame mirrors one activation of that walk; the
+// ResumePoint records where to resume after a spawned child block has produced
+// its result (returned via Returned). All observable behaviour is preserved:
+// cache lookups/inserts, VisitedBlocks cycle detection (including the
+// single-predecessor insert-without-erase asymmetry), predecessor operand
+// order, and phi simplification/creation.
+MemoryAccess *MemorySSAUpdater::getPreviousDefIterative(
     BasicBlock *BB,
     DenseMap<BasicBlock *, TrackingVH<MemoryAccess>> &CachedPreviousDef) {
-  // First, do a cache lookup. Without this cache, certain CFG structures
-  // (like a series of if statements) take exponential time to visit.
-  auto Cached = CachedPreviousDef.find(BB);
-  if (Cached != CachedPreviousDef.end())
-    return Cached->second;
-
-  // If this method is called from an unreachable block, return LoE.
-  if (!MSSA->DT->isReachableFromEntry(BB))
-    return MSSA->getLiveOnEntryDef();
+  // One frame of the explicit worklist, standing in for a single activation of
+  // the old recursive walk. A frame advances through a small state machine
+  // driven by its ResumePoint, suspending whenever it needs the result of a
+  // child block and resuming once that result is available in Returned:
+  //
+  //   EnterBlock       First visit to the frame's block: do the cache /
+  //                    unreachable / unique-predecessor / cycle checks. May
+  //                    finish the frame outright, or push a child frame and
+  //                    move to ResumeSinglePred or RunPredLoop.
+  //   ResumeSinglePred Resume after the unique predecessor has been resolved:
+  //                    cache its result for this block and finish the frame.
+  //   RunPredLoop      Walk the predecessors, gathering phi operands and
+  //                    suspending into a child frame for each predecessor that
+  //                    has no local definition; once all are gathered, place or
+  //                    simplify the phi and finish the frame.
+  struct StackFrame {
+    enum class ResumePoint { EnterBlock, ResumeSinglePred, RunPredLoop };
+
+    BasicBlock *BB;
+    ResumePoint Resume = ResumePoint::EnterBlock;
+    // Multi-predecessor loop state.
+    SmallVector<TrackingVH<MemoryAccess>, 8> PhiOps;
+    SmallVector<BasicBlock *, 8> Preds;
+    unsigned PredIdx = 0;
+    // When set, `Returned` holds the result for Preds[PredIdx].
+    bool PendingIncoming = false;
+    bool UniqueIncomingAccess = true;
+    MemoryAccess *SingleAccess = nullptr;
 
-  if (BasicBlock *Pred = BB->getUniquePredecessor()) {
-    VisitedBlocks.insert(BB);
-    // Single predecessor case, just recurse, we can only have one definition.
-    MemoryAccess *Result = getPreviousDefFromEnd(Pred, CachedPreviousDef);
-    CachedPreviousDef.insert({BB, Result});
-    return Result;
-  }
+    // Fold an incoming predecessor access into this frame's phi operands,
+    // tracking whether all incoming accesses are identical (so the phi may be
+    // elided). Mirrors the per-predecessor body of the walk.
+    void incorporate(MemoryAccess *Incoming) {
+      if (!SingleAccess)
+        SingleAccess = Incoming;
+      else if (Incoming != SingleAccess)
+        UniqueIncomingAccess = false;
+      PhiOps.push_back(Incoming);
+    }
+  };
 
-  if (VisitedBlocks.count(BB)) {
-    // We hit our node again, meaning we had a cycle, we must insert a phi
-    // node to break it so we have an operand. The only case this will
-    // insert useless phis is if we have irreducible control flow.
-    MemoryAccess *Result = MSSA->createMemoryPhi(BB);
-    CachedPreviousDef.insert({BB, Result});
+  // Non-recursive part of getPreviousDefFromEnd(Pred): if Pred has a local
+  // definition, return its last def and cache it under Pred (exactly as that
+  // helper does). Returns nullptr when Pred has no local def, signalling that
+  // it must instead be visited via its own worklist frame (the part of
+  // getPreviousDefFromEnd that used to recurse).
+  auto GetLocalDefFromEnd = [&](BasicBlock *Pred) -> MemoryAccess * {
+    auto *Defs = MSSA->getBlockDefs(Pred);
+    if (!Defs)
+      return nullptr;
+    MemoryAccess *Result = &*Defs->rbegin();
+    CachedPreviousDef.insert({Pred, Result});
     return Result;
-  }
+  };
 
-  if (VisitedBlocks.insert(BB).second) {
-    // Mark us visited so we can detect a cycle
-    SmallVector<TrackingVH<MemoryAccess>, 8> PhiOps;
+  SmallVector<StackFrame, 8> WorkStack;
+  WorkStack.push_back({BB});
+  // Carries a completed child frame's result back to its parent, the way a
+  // return value would propagate up a call stack.
+  MemoryAccess *Returned = nullptr;
+
+  while (!WorkStack.empty()) {
+    // NOTE: WorkStack.push_back below may reallocate and invalidate this
+    // reference, so every code path that pushes a new frame sets the
+    // ResumePoint on the current frame *before* pushing and then continues the
+    // loop without touching the reference again.
+    StackFrame &F = WorkStack.back();
+    BasicBlock *CurBB = F.BB;
+
+    switch (F.Resume) {
+    case StackFrame::ResumePoint::EnterBlock: {
----------------
antoniofrighetto wrote:

If we added an alias `using ResumePoint = StackFrame::ResumePoint;` on top, then the case would read `case ResumePoint::EnterBlock`, same below for the assignment (`F.Resume = ResumePoint::RunPredLoop;`).

https://github.com/llvm/llvm-project/pull/205159