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

via llvm-commits llvm-commits at lists.llvm.org
Fri Jul 17 05:42:20 PDT 2026


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

>From d2502ace26b24f3d5b41173deefa1c79eed0e609 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        | 281 ++++++++++++------
 llvm/unittests/Analysis/MemorySSATest.cpp     |  48 +++
 3 files changed, 247 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..37cea287284cd 100644
--- a/llvm/lib/Analysis/MemorySSAUpdater.cpp
+++ b/llvm/lib/Analysis/MemorySSAUpdater.cpp
@@ -33,101 +33,216 @@ 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 so that its depth does not scale with the length of the walk;
+// otherwise deep CFGs (e.g. long block chains in large generated
+// kernels/shaders) could overflow the native stack.
+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. Each frame runs a small state machine
+  // driven by its ResumePoint, suspending when it needs a child block's result
+  // (delivered via Returned) and resuming once that result is available:
+  //   EnterBlock       Initial cache / unreachable / unique-predecessor / cycle
+  //                    checks. May finish the frame or push a child frame.
+  //   ResumeSinglePred Resume after the unique predecessor is resolved.
+  //   RunPredLoop      Gather phi operands from predecessors, then place or
+  //                    simplify the phi.
+  struct StackFrame {
+    enum class ResumePoint { EnterBlock, ResumeSinglePred, RunPredLoop };
+
+    BasicBlock *BB;
+    explicit StackFrame(BasicBlock *BB) : BB(BB), PredIt(pred_begin(BB)) {}
+    ResumePoint Resume = ResumePoint::EnterBlock;
+    // Multi-predecessor loop state.
+    SmallVector<TrackingVH<MemoryAccess>, 4> PhiOps;
+    // Cursor over BB's predecessors for the resumable RunPredLoop walk. This
+    // stays valid across suspend/resume because the walk only modifies
+    // MemorySSA, never terminators or CFG edges, so BB's predecessor list does
+    // not change.
+    pred_iterator PredIt;
+    // When set, `Returned` holds the result for the predecessor at PredIt.
+    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).
+    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, cache and return its last def. Returns nullptr when Pred has no
+  // local def, so it must instead be visited via its own worklist frame.
+  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.emplace_back(BB);
+  // Carries a completed child frame's result back to its parent frame.
+  MemoryAccess *Returned = nullptr;
+
+  while (!WorkStack.empty()) {
+    // NOTE: emplace_back below may reallocate and invalidate this reference, so
+    // every path that pushes a new frame sets the ResumePoint first 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. If Pred
+        // has a local def take it, otherwise descend into Pred.
+        if (MemoryAccess *Result = GetLocalDefFromEnd(Pred)) {
+          CachedPreviousDef.insert({CurBB, Result});
+          Returned = Result;
+          WorkStack.pop_back();
+          continue;
+        }
+        F.Resume = ResumePoint::ResumeSinglePred;
+        WorkStack.emplace_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.
+      // PredIt was initialized to pred_begin(CurBB) when the frame was created.
+      VisitedBlocks.insert(CurBB);
+      F.Resume = ResumePoint::RunPredLoop;
+      continue;
+    }
+
+    case ResumePoint::ResumeSinglePred: {
+      // The single predecessor's 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 the result for the predecessor at PredIt.
+        F.incorporate(Returned);
+        F.PendingIncoming = false;
+        ++F.PredIt;
+      }
+
+      bool Suspended = false;
+      for (; F.PredIt != pred_end(CurBB); ++F.PredIt) {
+        BasicBlock *Pred = *F.PredIt;
+        if (MSSA->DT->isReachableFromEntry(Pred)) {
+          // Local def resolves now, otherwise descend into Pred.
+          if (MemoryAccess *IncomingAccess = GetLocalDefFromEnd(Pred)) {
+            F.incorporate(IncomingAccess);
+          } else {
+            F.PendingIncoming = true;
+            WorkStack.emplace_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 +253,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 +294,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..7780c57501465 100644
--- a/llvm/unittests/Analysis/MemorySSATest.cpp
+++ b/llvm/unittests/Analysis/MemorySSATest.cpp
@@ -1801,3 +1801,51 @@ TEST_F(MemorySSATest, TestNoDbgInsts) {
   ASSERT_EQ(MSSA.getMemoryAccess(DbgDeclare), nullptr);
   ASSERT_EQ(MSSA.getMemoryAccess(DbgValue), nullptr);
 }
+
+// getPreviousDefIterative walks the predecessors with an explicit worklist so
+// its stack usage does not scale with the length of the walk. This builds a
+// long single-predecessor chain with a store in the entry block and a load in
+// the last block, forcing the walk over the whole chain, and checks it
+// completes without exhausting the native stack.
+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, with a chain long enough that a
+  // per-block stack frame would exhaust the native stack.
+  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, forcing the walk back over the whole
+  // chain.
+  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.
+  EXPECT_EQ(LoadAccess->getDefiningAccess(), MSSA.getMemoryAccess(SI));
+  MSSA.verifyMemorySSA();
+}



More information about the llvm-commits mailing list