[llvm] [SROA] Enable alloca promotion for stores through pointer PHIs (PR #222893)
Yonah Goldberg via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 11 02:11:36 PDT 2026
https://github.com/YonahGoldberg created https://github.com/llvm/llvm-project/pull/222893
None
>From 8b7a4cb542c56bd728777be17415c9e4ba4be178 Mon Sep 17 00:00:00 2001
From: Yonah Goldberg <ygoldberg at nvidia.com>
Date: Fri, 11 Sep 2026 09:09:24 +0000
Subject: [PATCH] fix
---
llvm/lib/Transforms/Scalar/SROA.cpp | 161 ++++++++++++++--
llvm/test/Transforms/SROA/phi-store.ll | 249 +++++++++++++++++++++++++
2 files changed, 392 insertions(+), 18 deletions(-)
create mode 100644 llvm/test/Transforms/SROA/phi-store.ll
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index a0d7a0c921796..8ba88949a700b 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -113,7 +113,7 @@ STATISTIC(NumLoadsPredicated,
"Number of loads rewritten into predicated loads to allow promotion");
STATISTIC(
NumStoresPredicated,
- "Number of stores rewritten into predicated loads to allow promotion");
+ "Number of stores rewritten into predicated stores to allow promotion");
STATISTIC(NumDeleted, "Number of instructions deleted");
STATISTIC(NumVectorized, "Number of vectorized aggregates");
@@ -215,6 +215,13 @@ class SROA {
/// queue.
SmallSetVector<PHINode *, 8> SpeculatablePHIs;
+ /// A worklist of stores through PHIs to predicate onto the incoming edges.
+ ///
+ /// All of these stores have been checked so that rewriting them will allow
+ /// the corresponding allocas to be promoted. Edges which require splitting
+ /// are only present when CFG modification is allowed.
+ SmallMapVector<PHINode *, StoreInst *, 8> PHIStoresToRewrite;
+
/// A worklist of select instructions to rewrite prior to promoting
/// allocas.
SmallMapVector<SelectInst *, RewriteableMemOps, 8> SelectsToRewrite;
@@ -1569,8 +1576,8 @@ findCommonType(AllocaSlices::const_iterator B, AllocaSlices::const_iterator E,
/// ...
/// %V = phi [i32 %V1, i32 %V2]
///
-/// We can do this to a select if its only uses are loads and if the operands
-/// to the select can be loaded unconditionally.
+/// We can do this to a PHI if its only uses are loads and if any loads moved
+/// across other outgoing edges can be executed unconditionally.
///
/// FIXME: This should be hoisted into a generic utility, likely in
/// Transforms/Util/Local.h
@@ -1580,7 +1587,6 @@ static bool isSafePHIToSpeculate(PHINode &PN) {
// For now, we can only do this promotion if the load is in the same block
// as the PHI, and if there are no stores between the phi and load.
// TODO: Allow recursive phi users.
- // TODO: Allow stores.
BasicBlock *BB = PN.getParent();
Align MaxAlign;
uint64_t APWidth = DL.getIndexTypeSizeInBits(PN.getType());
@@ -1649,6 +1655,61 @@ static bool isSafePHIToSpeculate(PHINode &PN) {
return true;
}
+/// Check whether a single store through PN can be moved onto each incoming
+/// edge.
+static StoreInst *getPHIStoreToRewrite(PHINode &PN, bool PreserveCFG,
+ DominatorTree &DT) {
+ // TODO: Support multiple stores and mixed load/store users.
+ // TODO: Look through other instructions, such as other phis or addrspacecasts
+ if (!PN.hasOneUse())
+ return nullptr;
+
+ auto *SI = dyn_cast<StoreInst>(PN.user_back());
+ if (!SI || SI->getPointerOperand() != &PN)
+ return nullptr;
+
+ if (SI->isVolatile())
+ return nullptr;
+
+ BasicBlock *BB = PN.getParent();
+ // TODO: Allow a harmless prefix between the PHIs and the store.
+ if (&*BB->getFirstNonPHIOrDbg() != SI)
+ return nullptr;
+
+ Value *StoredValue = SI->getValueOperand();
+ SmallPtrSet<BasicBlock *, 4> SeenPreds;
+ for (unsigned Idx = 0, Num = PN.getNumIncomingValues(); Idx != Num; ++Idx) {
+ BasicBlock *Pred = PN.getIncomingBlock(Idx);
+
+ // Only one store is needed for duplicate edges from the same predecessor.
+ if (!SeenPreds.insert(Pred).second)
+ continue;
+
+ // TODO: Support other terminators.
+ Instruction *TI = Pred->getTerminator();
+ if (!isa<CondBrInst, UncondBrInst>(TI))
+ return nullptr;
+
+ if (Pred == BB)
+ return nullptr;
+
+ // TODO: If StoredValue is another PHI in BB, use its corresponding
+ // incoming value instead of requiring it to dominate every predecessor.
+ if (!DT.dominates(StoredValue, TI))
+ return nullptr;
+
+ if (TI->getNumSuccessors() == 1)
+ continue;
+
+ // If the predecessor has more than one successor, then we will need to
+ // split it so that we can insert the store only on the path to this BB.
+ if (PreserveCFG || !BB->canSplitPredecessors())
+ return nullptr;
+ }
+
+ return SI;
+}
+
static void speculatePHINodeLoads(IRBuilderTy &IRB, PHINode &PN) {
LLVM_DEBUG(dbgs() << " original: " << PN << "\n");
@@ -1702,6 +1763,51 @@ static void speculatePHINodeLoads(IRBuilderTy &IRB, PHINode &PN) {
PN.eraseFromParent();
}
+/// Move a store through a pointer PHI onto each of the PHI's incoming edges.
+/// Returns whether this required modifying the CFG.
+static bool rewritePHINodeStore(PHINode &PN, StoreInst &SI, DomTreeUpdater &DTU,
+ SmallSetVector<AllocaInst *, 16> &Worklist) {
+ LLVM_DEBUG(dbgs() << " original: " << PN << "\n"
+ << " " << SI << "\n");
+
+ // Splitting one edge rewrites all PHIs in the destination block. Snapshot
+ // the original predecessor/value pairs before making any CFG changes.
+ SmallVector<std::pair<BasicBlock *, Value *>, 4> IncomingValues;
+ SmallPtrSet<BasicBlock *, 4> SeenPreds;
+ for (unsigned Idx = 0, Num = PN.getNumIncomingValues(); Idx != Num; ++Idx) {
+ BasicBlock *Pred = PN.getIncomingBlock(Idx);
+ if (!SeenPreds.insert(Pred).second)
+ continue;
+ Value *InVal = PN.getIncomingValue(Idx);
+ IncomingValues.emplace_back(Pred, InVal);
+
+ // Revisit every alloca exposed by removing the pointer PHI,
+ if (auto *AI = dyn_cast<AllocaInst>(getUnderlyingObject(InVal)))
+ Worklist.insert(AI);
+ }
+
+ bool CFGChanged = false;
+ BasicBlock *BB = PN.getParent();
+ for (auto [Pred, InVal] : IncomingValues) {
+ BasicBlock *StoreBB = Pred;
+ if (Pred->getTerminator()->getNumSuccessors() != 1) {
+ StoreBB = SplitBlockPredecessors(BB, {Pred}, ".sroa.store", &DTU);
+ assert(StoreBB && "store edge was not checked for splitting");
+ CFGChanged = true;
+ }
+
+ auto *NewStore = cast<StoreInst>(SI.clone());
+ NewStore->setOperand(StoreInst::getPointerOperandIndex(), InVal);
+ NewStore->insertBefore(StoreBB->getTerminator()->getIterator());
+ ++NumStoresPredicated;
+ LLVM_DEBUG(dbgs() << " to: " << *NewStore << "\n");
+ }
+
+ SI.eraseFromParent();
+ PN.eraseFromParent();
+ return CFGChanged;
+}
+
SelectHandSpeculativity &
SelectHandSpeculativity::setAsSpeculatable(bool isTrueVal) {
if (isTrueVal)
@@ -5630,14 +5736,24 @@ SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS, Partition &P) {
// Now that we've processed all the slices in the new partition, check if any
// PHIs or Selects would block promotion.
- for (PHINode *PHI : PHIUsers)
- if (!isSafePHIToSpeculate(*PHI)) {
- Promotable = false;
- PHIUsers.clear();
- SelectUsers.clear();
- break;
+ SmallVector<PHINode *, 8> NewSpeculatablePHIs;
+ SmallVector<std::pair<PHINode *, StoreInst *>, 2> NewPHIStoresToRewrite;
+ for (PHINode *PHI : PHIUsers) {
+ if (isSafePHIToSpeculate(*PHI)) {
+ NewSpeculatablePHIs.push_back(PHI);
+ continue;
+ }
+ if (StoreInst *SI =
+ getPHIStoreToRewrite(*PHI, PreserveCFG, DTU->getDomTree())) {
+ NewPHIStoresToRewrite.emplace_back(PHI, SI);
+ continue;
}
+ Promotable = false;
+ SelectUsers.clear();
+ break;
+ }
+
SmallVector<std::pair<SelectInst *, RewriteableMemOps>, 2>
NewSelectsToRewrite;
NewSelectsToRewrite.reserve(SelectUsers.size());
@@ -5646,9 +5762,6 @@ SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS, Partition &P) {
isSafeSelectToSpeculate(*Sel, PreserveCFG);
if (!Ops) {
Promotable = false;
- PHIUsers.clear();
- SelectUsers.clear();
- NewSelectsToRewrite.clear();
break;
}
NewSelectsToRewrite.emplace_back(std::make_pair(Sel, *Ops));
@@ -5662,14 +5775,19 @@ SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS, Partition &P) {
if (isInstructionTriviallyDead(OldInst))
DeadInsts.push_back(OldInst);
}
- if (PHIUsers.empty() && SelectUsers.empty()) {
+ if (NewSpeculatablePHIs.empty() && NewPHIStoresToRewrite.empty() &&
+ SelectUsers.empty()) {
// Promote the alloca.
PromotableAllocas.insert(NewAI);
} else {
- // If we have either PHIs or Selects to speculate, add them to those
- // worklists and re-queue the new alloca so that we promote in on the
- // next iteration.
- SpeculatablePHIs.insert_range(PHIUsers);
+ // If we have either PHIs or Selects to rewrite, add them to those
+ // worklists and re-queue the new alloca so that we promote it on the next
+ // iteration.
+ SpeculatablePHIs.insert_range(NewSpeculatablePHIs);
+ PHIStoresToRewrite.reserve(PHIStoresToRewrite.size() +
+ NewPHIStoresToRewrite.size());
+ for (auto [PN, SI] : NewPHIStoresToRewrite)
+ PHIStoresToRewrite.insert({PN, SI});
SelectsToRewrite.reserve(SelectsToRewrite.size() +
NewSelectsToRewrite.size());
for (auto &&KV : llvm::make_range(
@@ -6230,6 +6348,13 @@ SROA::runOnAlloca(AllocaInst &AI) {
while (!SpeculatablePHIs.empty())
speculatePHINodeLoads(IRB, *SpeculatablePHIs.pop_back_val());
+ LLVM_DEBUG(dbgs() << " Rewriting stores through PHIs\n");
+ auto RemainingPHIStoresToRewrite = PHIStoresToRewrite.takeVector();
+ while (!RemainingPHIStoresToRewrite.empty()) {
+ const auto [PN, SI] = RemainingPHIStoresToRewrite.pop_back_val();
+ CFGChanged |= rewritePHINodeStore(*PN, *SI, *DTU, Worklist);
+ }
+
LLVM_DEBUG(dbgs() << " Rewriting Selects\n");
auto RemainingSelectsToRewrite = SelectsToRewrite.takeVector();
while (!RemainingSelectsToRewrite.empty()) {
diff --git a/llvm/test/Transforms/SROA/phi-store.ll b/llvm/test/Transforms/SROA/phi-store.ll
new file mode 100644
index 0000000000000..df5104ccf5fa5
--- /dev/null
+++ b/llvm/test/Transforms/SROA/phi-store.ll
@@ -0,0 +1,249 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
+; RUN: opt -S -sroa-skip-mem2reg -passes='sroa<preserve-cfg>' < %s | FileCheck %s --check-prefixes=CHECK,CHECK-PRESERVE-CFG
+; RUN: opt -S -sroa-skip-mem2reg -passes='sroa<modify-cfg>' < %s | FileCheck %s --check-prefixes=CHECK,CHECK-MODIFY-CFG
+
+%pair = type { i32, i32 }
+
+define void @store_to_phi(i1 %cond, i32 %value) {
+; CHECK-LABEL: @store_to_phi(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT: [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: store i32 [[VALUE:%.*]], ptr [[A]], align 4
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: store i32 [[VALUE]], ptr [[B]], align 4
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ %b = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %b, %right ]
+ store i32 %value, ptr %ptr, align 4
+ ret void
+}
+
+define void @store_to_phi_alloca_and_argument(i1 %cond, i32 %value, ptr %other) {
+; CHECK-LABEL: @store_to_phi_alloca_and_argument(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: store i32 [[VALUE:%.*]], ptr [[A]], align 4
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: store i32 [[VALUE]], ptr [[OTHER:%.*]], align 4
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %other, %right ]
+ store i32 %value, ptr %ptr, align 4
+ ret void
+}
+
+define void @store_to_phi_split_edge(i1 %cond, i1 %other_cond, i32 %value) {
+; CHECK-PRESERVE-CFG-LABEL: @store_to_phi_split_edge(
+; CHECK-PRESERVE-CFG-NEXT: entry:
+; CHECK-PRESERVE-CFG-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-PRESERVE-CFG-NEXT: [[B:%.*]] = alloca i32, align 4
+; CHECK-PRESERVE-CFG-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK-PRESERVE-CFG: left:
+; CHECK-PRESERVE-CFG-NEXT: br label [[MERGE:%.*]]
+; CHECK-PRESERVE-CFG: right:
+; CHECK-PRESERVE-CFG-NEXT: br i1 [[OTHER_COND:%.*]], label [[MERGE]], label [[EXIT:%.*]]
+; CHECK-PRESERVE-CFG: merge:
+; CHECK-PRESERVE-CFG-NEXT: [[PTR:%.*]] = phi ptr [ [[A]], [[LEFT]] ], [ [[B]], [[RIGHT]] ]
+; CHECK-PRESERVE-CFG-NEXT: store i32 [[VALUE:%.*]], ptr [[PTR]], align 4
+; CHECK-PRESERVE-CFG-NEXT: br label [[EXIT]]
+; CHECK-PRESERVE-CFG: exit:
+; CHECK-PRESERVE-CFG-NEXT: ret void
+;
+; CHECK-MODIFY-CFG-LABEL: @store_to_phi_split_edge(
+; CHECK-MODIFY-CFG-NEXT: entry:
+; CHECK-MODIFY-CFG-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-MODIFY-CFG-NEXT: [[B:%.*]] = alloca i32, align 4
+; CHECK-MODIFY-CFG-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK-MODIFY-CFG: left:
+; CHECK-MODIFY-CFG-NEXT: store i32 [[VALUE:%.*]], ptr [[A]], align 4
+; CHECK-MODIFY-CFG-NEXT: br label [[MERGE:%.*]]
+; CHECK-MODIFY-CFG: right:
+; CHECK-MODIFY-CFG-NEXT: br i1 [[OTHER_COND:%.*]], label [[MERGE_SROA_STORE:%.*]], label [[EXIT:%.*]]
+; CHECK-MODIFY-CFG: merge.sroa.store:
+; CHECK-MODIFY-CFG-NEXT: store i32 [[VALUE]], ptr [[B]], align 4
+; CHECK-MODIFY-CFG-NEXT: br label [[MERGE]]
+; CHECK-MODIFY-CFG: merge:
+; CHECK-MODIFY-CFG-NEXT: br label [[EXIT]]
+; CHECK-MODIFY-CFG: exit:
+; CHECK-MODIFY-CFG-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ %b = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br i1 %other_cond, label %merge, label %exit
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %b, %right ]
+ store i32 %value, ptr %ptr, align 4
+ br label %exit
+
+exit:
+ ret void
+}
+
+define void @store_to_phi_value_phi(i1 %cond) {
+; CHECK-LABEL: @store_to_phi_value_phi(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT: [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: [[PTR:%.*]] = phi ptr [ [[A]], [[LEFT]] ], [ [[B]], [[RIGHT]] ]
+; CHECK-NEXT: [[VALUE:%.*]] = phi i32 [ 1, [[LEFT]] ], [ 2, [[RIGHT]] ]
+; CHECK-NEXT: store i32 [[VALUE]], ptr [[PTR]], align 4
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ %b = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %b, %right ]
+ %value = phi i32 [ 1, %left ], [ 2, %right ]
+ store i32 %value, ptr %ptr, align 4
+ ret void
+}
+
+define void @store_to_phi_of_geps(i1 %cond, i32 %value) {
+; CHECK-LABEL: @store_to_phi_of_geps(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A_SROA_0:%.*]] = alloca i32, align 4
+; CHECK-NEXT: [[B_SROA_0:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: store i32 [[VALUE:%.*]], ptr [[A_SROA_0]], align 4
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: store i32 [[VALUE]], ptr [[B_SROA_0]], align 4
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca %pair, align 4
+ %b = alloca %pair, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %b, %right ]
+ %field = getelementptr inbounds %pair, ptr %ptr, i32 0, i32 1
+ store i32 %value, ptr %field, align 4
+ ret void
+}
+
+define void @store_atomic_to_phi(i1 %cond, i32 %value, ptr %other) {
+; CHECK-LABEL: @store_atomic_to_phi(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: store i32 [[VALUE:%.*]], ptr [[A]], align 4
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: store atomic i32 [[VALUE]], ptr [[OTHER:%.*]] seq_cst, align 4
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %other, %right ]
+ store atomic i32 %value, ptr %ptr seq_cst, align 4
+ ret void
+}
+
+define void @store_atomic_volatile_to_phi(i1 %cond, i32 %value, ptr %other) {
+; CHECK-LABEL: @store_atomic_volatile_to_phi(
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT: br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK: left:
+; CHECK-NEXT: br label [[MERGE:%.*]]
+; CHECK: right:
+; CHECK-NEXT: br label [[MERGE]]
+; CHECK: merge:
+; CHECK-NEXT: [[PTR:%.*]] = phi ptr [ [[A]], [[LEFT]] ], [ [[OTHER:%.*]], [[RIGHT]] ]
+; CHECK-NEXT: store atomic volatile i32 [[VALUE:%.*]], ptr [[PTR]] seq_cst, align 4
+; CHECK-NEXT: ret void
+;
+entry:
+ %a = alloca i32, align 4
+ br i1 %cond, label %left, label %right
+
+left:
+ br label %merge
+
+right:
+ br label %merge
+
+merge:
+ %ptr = phi ptr [ %a, %left ], [ %other, %right ]
+ store atomic volatile i32 %value, ptr %ptr seq_cst, align 4
+ ret void
+}
More information about the llvm-commits
mailing list