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

via llvm-commits llvm-commits at lists.llvm.org
Wed Jun 24 07:47:45 PDT 2026


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

>From ec53d481beda94aef53a63bc684d49bd7d081ba0 Mon Sep 17 00:00:00 2001
From: "Stachowiak, Urszula" <urszula.stachowiak at intel.com>
Date: Mon, 22 Jun 2026 19:31:00 +0200
Subject: [PATCH] [MemorySSA] Make getPreviousDef iterative to avoid stack
 overflow

MemorySSAUpdater::getPreviousDefRecursive walked the CFG predecessor-first
using native recursion (mutually recursive with getPreviousDefFromEnd). On
functions with very deep control flow - e.g. long chains of blocks produced
by large auto-generated kernels/shaders - the recursion depth scales with the
number of blocks and can overflow the native stack.

Rewrite the walk as an explicit worklist of frames on the heap. Each frame
tracks the block being processed and a resume stage so a child block's result
can be folded back into its parent, exactly mirroring the previous recursive
control flow. Behaviour is otherwise unchanged: cache lookups, VisitedBlocks
cycle detection (including the single-predecessor insert-without-erase
asymmetry), predecessor operand order, and phi placement/simplification are
all preserved. The helper is renamed getPreviousDefRecursive ->
getPreviousDefIterative since it no longer recurses.

Add a unit test that builds a long single-predecessor chain and incrementally
updates MemorySSA; the previous recursive implementation overflowed the stack
on such input.

Fixes: https://github.com/llvm/llvm-project/issues/121279
---
 llvm/include/llvm/Analysis/MemorySSAUpdater.h |   2 +-
 llvm/lib/Analysis/MemorySSAUpdater.cpp        | 298 +++++++++++++-----
 llvm/unittests/Analysis/MemorySSATest.cpp     |  54 ++++
 3 files changed, 270 insertions(+), 84 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MemorySSAUpdater.h b/llvm/include/llvm/Analysis/MemorySSAUpdater.h
index 96bf99922d848..5b81deee5987f 100644
--- a/llvm/include/llvm/Analysis/MemorySSAUpdater.h
+++ b/llvm/include/llvm/Analysis/MemorySSAUpdater.h
@@ -255,7 +255,7 @@ class MemorySSAUpdater {
   getPreviousDefFromEnd(BasicBlock *,
                         DenseMap<BasicBlock *, TrackingVH<MemoryAccess>> &);
   MemoryAccess *
-  getPreviousDefRecursive(BasicBlock *,
+  getPreviousDefIterative(BasicBlock *,
                           DenseMap<BasicBlock *, TrackingVH<MemoryAccess>> &);
   MemoryAccess *recursePhi(MemoryAccess *Phi);
   MemoryAccess *tryRemoveTrivialPhi(MemoryPhi *Phi);
diff --git a/llvm/lib/Analysis/MemorySSAUpdater.cpp b/llvm/lib/Analysis/MemorySSAUpdater.cpp
index 4492a74a2bd7e..6848f17d5c6e5 100644
--- a/llvm/lib/Analysis/MemorySSAUpdater.cpp
+++ b/llvm/lib/Analysis/MemorySSAUpdater.cpp
@@ -33,101 +33,233 @@ 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});
+    // 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);
+    }
+  };
+  using ResumePoint = StackFrame::ResumePoint;
+
+  // 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.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});
-    return Result;
-  }
+  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 ResumePoint::EnterBlock: {
+      // 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(CurBB);
+      if (Cached != CachedPreviousDef.end()) {
+        Returned = Cached->second;
+        WorkStack.pop_back();
+        continue;
+      }
 
-  if (VisitedBlocks.insert(BB).second) {
-    // Mark us visited so we can detect a cycle
-    SmallVector<TrackingVH<MemoryAccess>, 8> PhiOps;
+      // If this method is called from an unreachable block, return LoE.
+      if (!MSSA->DT->isReachableFromEntry(CurBB)) {
+        Returned = MSSA->getLiveOnEntryDef();
+        WorkStack.pop_back();
+        continue;
+      }
 
-    // Recurse to get the values in our predecessors for placement of a
-    // potential phi node. This will insert phi nodes if we cycle in order to
-    // break the cycle and have an operand.
-    bool UniqueIncomingAccess = true;
-    MemoryAccess *SingleAccess = nullptr;
-    for (auto *Pred : predecessors(BB)) {
-      if (MSSA->DT->isReachableFromEntry(Pred)) {
-        auto *IncomingAccess = getPreviousDefFromEnd(Pred, CachedPreviousDef);
-        if (!SingleAccess)
-          SingleAccess = IncomingAccess;
-        else if (IncomingAccess != SingleAccess)
-          UniqueIncomingAccess = false;
-        PhiOps.push_back(IncomingAccess);
-      } else
-        PhiOps.push_back(MSSA->getLiveOnEntryDef());
+      if (BasicBlock *Pred = CurBB->getUniquePredecessor()) {
+        VisitedBlocks.insert(CurBB);
+        // Single predecessor case, there can be only one definition. Resolve
+        // getPreviousDefFromEnd(Pred): if Pred has a local def, take it;
+        // otherwise descend into Pred by pushing a new frame.
+        if (MemoryAccess *Result = GetLocalDefFromEnd(Pred)) {
+          CachedPreviousDef.insert({CurBB, Result});
+          Returned = Result;
+          WorkStack.pop_back();
+          continue;
+        }
+        F.Resume = ResumePoint::ResumeSinglePred;
+        WorkStack.push_back({Pred});
+        continue;
+      }
+
+      if (VisitedBlocks.count(CurBB)) {
+        // 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(CurBB);
+        CachedPreviousDef.insert({CurBB, Result});
+        Returned = Result;
+        WorkStack.pop_back();
+        continue;
+      }
+
+      // Mark us visited so we can detect a cycle, then walk the predecessors.
+      VisitedBlocks.insert(CurBB);
+      for (BasicBlock *Pred : predecessors(CurBB))
+        F.Preds.push_back(Pred);
+      F.Resume = ResumePoint::RunPredLoop;
+      continue;
+    }
+
+    case ResumePoint::ResumeSinglePred: {
+      // The single predecessor's getPreviousDefFromEnd result is in Returned.
+      CachedPreviousDef.insert({CurBB, Returned});
+      WorkStack.pop_back();
+      continue;
     }
 
-    // Now try to simplify the ops to avoid placing a phi.
-    // This may return null if we never created a phi yet, that's okay
-    MemoryPhi *Phi = dyn_cast_or_null<MemoryPhi>(MSSA->getMemoryAccess(BB));
-
-    // See if we can avoid the phi by simplifying it.
-    auto *Result = tryRemoveTrivialPhi(Phi, PhiOps);
-    // If we couldn't simplify, we may have to create a phi
-    if (Result == Phi && UniqueIncomingAccess && SingleAccess) {
-      // A concrete Phi only exists if we created an empty one to break a cycle.
-      if (Phi) {
-        assert(Phi->operands().empty() && "Expected empty Phi");
-        Phi->replaceAllUsesWith(SingleAccess);
-        removeMemoryAccess(Phi);
+    case ResumePoint::RunPredLoop: {
+      // Get the values in our predecessors for placement of a potential phi
+      // node. This will insert phi nodes if we cycle in order to break the
+      // cycle and have an operand.
+      if (F.PendingIncoming) {
+        // Returned holds getPreviousDefFromEnd(Preds[PredIdx]) for the
+        // reachable predecessor we just descended into.
+        F.incorporate(Returned);
+        F.PendingIncoming = false;
+        ++F.PredIdx;
+      }
+
+      bool Suspended = false;
+      for (; F.PredIdx < F.Preds.size(); ++F.PredIdx) {
+        BasicBlock *Pred = F.Preds[F.PredIdx];
+        if (MSSA->DT->isReachableFromEntry(Pred)) {
+          // Resolve getPreviousDefFromEnd(Pred): local def resolves now,
+          // otherwise descend into Pred by pushing a new frame.
+          if (MemoryAccess *IncomingAccess = GetLocalDefFromEnd(Pred)) {
+            F.incorporate(IncomingAccess);
+          } else {
+            F.PendingIncoming = true;
+            WorkStack.push_back({Pred});
+            Suspended = true;
+            break;
+          }
+        } else
+          F.PhiOps.push_back(MSSA->getLiveOnEntryDef());
       }
-      Result = SingleAccess;
-    } else if (Result == Phi && !(UniqueIncomingAccess && SingleAccess)) {
-      if (!Phi)
-        Phi = MSSA->createMemoryPhi(BB);
-
-      // See if the existing phi operands match what we need.
-      // Unlike normal SSA, we only allow one phi node per block, so we can't just
-      // create a new one.
-      if (Phi->getNumOperands() != 0) {
-        // FIXME: Figure out whether this is dead code and if so remove it.
-        if (!std::equal(Phi->op_begin(), Phi->op_end(), PhiOps.begin())) {
-          // These will have been filled in by the recursive read we did above.
-          llvm::copy(PhiOps, Phi->op_begin());
-          std::copy(pred_begin(BB), pred_end(BB), Phi->block_begin());
+      if (Suspended)
+        continue;
+
+      // Now try to simplify the ops to avoid placing a phi.
+      // This may return null if we never created a phi yet, that's okay
+      MemoryPhi *Phi =
+          dyn_cast_or_null<MemoryPhi>(MSSA->getMemoryAccess(CurBB));
+
+      // See if we can avoid the phi by simplifying it.
+      auto *Result = tryRemoveTrivialPhi(Phi, F.PhiOps);
+      // If we couldn't simplify, we may have to create a phi
+      if (Result == Phi && F.UniqueIncomingAccess && F.SingleAccess) {
+        // A concrete Phi only exists if we created an empty one to break a
+        // cycle.
+        if (Phi) {
+          assert(Phi->operands().empty() && "Expected empty Phi");
+          Phi->replaceAllUsesWith(F.SingleAccess);
+          removeMemoryAccess(Phi);
         }
-      } else {
-        unsigned i = 0;
-        for (auto *Pred : predecessors(BB))
-          Phi->addIncoming(&*PhiOps[i++], Pred);
-        InsertedPHIs.push_back(Phi);
+        Result = F.SingleAccess;
+      } else if (Result == Phi && !(F.UniqueIncomingAccess && F.SingleAccess)) {
+        if (!Phi)
+          Phi = MSSA->createMemoryPhi(CurBB);
+
+        // See if the existing phi operands match what we need.
+        // Unlike normal SSA, we only allow one phi node per block, so we can't
+        // just create a new one.
+        if (Phi->getNumOperands() != 0) {
+          // FIXME: Figure out whether this is dead code and if so remove it.
+          if (!std::equal(Phi->op_begin(), Phi->op_end(), F.PhiOps.begin())) {
+            // These will have been filled in by the predecessor walk above.
+            llvm::copy(F.PhiOps, Phi->op_begin());
+            llvm::copy(predecessors(CurBB), Phi->block_begin());
+          }
+        } else {
+          unsigned I = 0;
+          for (auto *Pred : predecessors(CurBB))
+            Phi->addIncoming(&*F.PhiOps[I++], Pred);
+          InsertedPHIs.push_back(Phi);
+        }
+        Result = Phi;
       }
-      Result = Phi;
-    }
 
-    // Set ourselves up for the next variable by resetting visited state.
-    VisitedBlocks.erase(BB);
-    CachedPreviousDef.insert({BB, Result});
-    return Result;
+      // Set ourselves up for the next variable by resetting visited state.
+      VisitedBlocks.erase(CurBB);
+      CachedPreviousDef.insert({CurBB, Result});
+      Returned = Result;
+      WorkStack.pop_back();
+      continue;
+    }
+    }
   }
-  llvm_unreachable("Should have hit one of the three cases above");
+
+  return Returned;
 }
 
 // This starts at the memory access, and goes backwards in the block to find the
@@ -138,7 +270,7 @@ MemoryAccess *MemorySSAUpdater::getPreviousDef(MemoryAccess *MA) {
   if (auto *LocalResult = getPreviousDefInBlock(MA))
     return LocalResult;
   DenseMap<BasicBlock *, TrackingVH<MemoryAccess>> CachedPreviousDef;
-  return getPreviousDefRecursive(MA->getBlock(), CachedPreviousDef);
+  return getPreviousDefIterative(MA->getBlock(), CachedPreviousDef);
 }
 
 // This starts at the memory access, and goes backwards in the block to the find
@@ -179,7 +311,7 @@ MemoryAccess *MemorySSAUpdater::getPreviousDefFromEnd(
     return &*Defs->rbegin();
   }
 
-  return getPreviousDefRecursive(BB, CachedPreviousDef);
+  return getPreviousDefIterative(BB, CachedPreviousDef);
 }
 // Recurse over a set of phi uses to eliminate the trivial ones
 MemoryAccess *MemorySSAUpdater::recursePhi(MemoryAccess *Phi) {
diff --git a/llvm/unittests/Analysis/MemorySSATest.cpp b/llvm/unittests/Analysis/MemorySSATest.cpp
index 6e4b7af2495b6..aa887499d2b7d 100644
--- a/llvm/unittests/Analysis/MemorySSATest.cpp
+++ b/llvm/unittests/Analysis/MemorySSATest.cpp
@@ -1801,3 +1801,57 @@ TEST_F(MemorySSATest, TestNoDbgInsts) {
   ASSERT_EQ(MSSA.getMemoryAccess(DbgDeclare), nullptr);
   ASSERT_EQ(MSSA.getMemoryAccess(DbgValue), nullptr);
 }
+
+// getPreviousDefIterative (formerly getPreviousDefRecursive) walks the CFG
+// backwards to find the previous memory definition for a use. The original
+// implementation used native recursion, nesting one call per block along the
+// predecessor walk (mutually recursive with getPreviousDefFromEnd), so a
+// sufficiently long block chain overflowed the native stack. This test builds
+// a long single-predecessor chain with a store in the entry block and inserts
+// a load in the last block; updating MemorySSA forces the walk over the whole
+// chain. The recursive implementation overflowed the stack on this input,
+// while the iterative one completes and wires the load to the entry store.
+TEST_F(MemorySSATest, DeepChainDoesNotRecurse) {
+  F = Function::Create(FunctionType::get(B.getVoidTy(), {B.getPtrTy()}, false),
+                       GlobalValue::ExternalLinkage, "F", &M);
+  Argument *PointerArg = &*F->arg_begin();
+
+  // Build entry -> bb0 -> bb1 -> ... -> exit. The chain is far deeper than the
+  // native stack could tolerate under the old recursive walk (~2 frames per
+  // block), but is processed in O(1) stack space by the iterative version.
+  const unsigned Depth = 16 * 1024;
+  SmallVector<BasicBlock *, 16> Blocks;
+  Blocks.push_back(BasicBlock::Create(C, "entry", F));
+  for (unsigned I = 0; I < Depth; ++I)
+    Blocks.push_back(BasicBlock::Create(C, "bb" + std::to_string(I), F));
+  Blocks.push_back(BasicBlock::Create(C, "exit", F));
+
+  // Entry block: a single store, then branch into the chain.
+  B.SetInsertPoint(Blocks.front());
+  StoreInst *SI = B.CreateStore(B.getInt8(0), PointerArg);
+  B.CreateBr(Blocks[1]);
+  // Remaining blocks: unconditional branch to the next, exit returns.
+  for (unsigned I = 1; I + 1 < Blocks.size(); ++I) {
+    B.SetInsertPoint(Blocks[I]);
+    B.CreateBr(Blocks[I + 1]);
+  }
+  B.SetInsertPoint(Blocks.back());
+  B.CreateRetVoid();
+
+  setupAnalyses();
+  MemorySSA &MSSA = *Analyses->MSSA;
+  MemorySSAUpdater Updater(&MSSA);
+
+  // Insert a load in the exit block. getPreviousDefIterative must walk the
+  // entire predecessor chain back to the entry store without recursing.
+  B.SetInsertPoint(&Blocks.back()->front());
+  LoadInst *LI = B.CreateLoad(B.getInt8Ty(), PointerArg);
+  MemoryUse *LoadAccess = cast<MemoryUse>(Updater.createMemoryAccessInBB(
+      LI, nullptr, Blocks.back(), MemorySSA::Beginning));
+  Updater.insertUse(LoadAccess, /*RenameUses=*/false);
+
+  // The load must be defined by the single store in the entry block, found at
+  // the far end of the chain.
+  EXPECT_EQ(LoadAccess->getDefiningAccess(), MSSA.getMemoryAccess(SI));
+  MSSA.verifyMemorySSA();
+}



More information about the llvm-commits mailing list