[llvm] [SROA] Extend tree-structured merge to handle init + RMW pattern (PR #194441)
Princeton Ferro via llvm-commits
llvm-commits at lists.llvm.org
Fri May 1 17:05:49 PDT 2026
================
@@ -2935,143 +2955,348 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
for (Slice &S : P) {
auto *User = cast<Instruction>(S.getUse()->getUser());
+ // A "full-width" slice spans the entire alloca; it's either the single
+ // init store (Pattern 2) or the single final load (both patterns).
+ bool IsFullWidth = (S.beginOffset() == NewAllocaBeginOffset &&
+ S.endOffset() == NewAllocaEndOffset);
if (auto *LI = dyn_cast<LoadInst>(User)) {
- // Do not handle the case if
- // 1. There is more than one load
- // 2. The load is volatile
- // 3. The load does not read the entire alloca structure
- // 4. The load does not meet the conditions in the helper function
- if (TheLoad || !IsTypeValidForTreeStructuredMerge(LI->getType()) ||
- S.beginOffset() != NewAllocaBeginOffset ||
- S.endOffset() != NewAllocaEndOffset || LI->isVolatile())
+ // Only handle simple (non-volatile, non-atomic) loads.
+ if (!LI->isSimple() ||
+ !IsTypeValidForTreeStructuredMerge(LI->getType()))
return std::nullopt;
- TheLoad = LI;
+ if (IsFullWidth) {
+ // We accept at most one full-width load (the "final" load, after
+ // all the partial stores).
+ if (FullLoad)
+ return std::nullopt;
+ FullLoad = LI;
+ } else {
+ // Partial load (RMW pattern only).
+ LoadInfos.push_back({LI, S.beginOffset(), S.endOffset()});
+ }
} else if (auto *SI = dyn_cast<StoreInst>(User)) {
// Do not handle the case if
// 1. The store does not meet the conditions in the helper function
- // 2. The store is volatile
- // 3. The total store size is not a multiple of the allocated element
- // type size
- if (!IsTypeValidForTreeStructuredMerge(
- SI->getValueOperand()->getType()) ||
- SI->isVolatile())
+ // 2. The store is not simple — we drop stores as part of the
+ // rewrite, so volatile stores (which must be kept) and atomic
+ // stores (which carry memory-ordering semantics) are unsound
+ // to replace with SSA bookkeeping.
+ // 3. The total store size is not a multiple of the allocated
+ // element type size (required so the tree merge can produce a
+ // vector whose element type matches the alloca).
+ if (!SI->isSimple() || !IsTypeValidForTreeStructuredMerge(
+ SI->getValueOperand()->getType()))
return std::nullopt;
- auto *VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
- unsigned NumElts = VecTy->getNumElements();
- unsigned EltSize = DL.getTypeSizeInBits(VecTy->getElementType());
+ auto *StVecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
+ unsigned NumElts = StVecTy->getNumElements();
+ unsigned EltSize = DL.getTypeSizeInBits(StVecTy->getElementType());
if (NumElts * EltSize % AllocatedEltTySize != 0)
return std::nullopt;
- StoreInfos.emplace_back(SI, S.beginOffset(), S.endOffset(),
- SI->getValueOperand());
+ if (IsFullWidth) {
+ // At most one full-width store is allowed — it's the init store
+ // for the RMW pattern.
+ if (InitStore)
+ return std::nullopt;
+ InitStore = SI;
+ InitValue = SI->getValueOperand();
+ } else {
+ StoreInfos.emplace_back(SI, S.beginOffset(), S.endOffset(),
+ SI->getValueOperand());
+ }
} else {
- // If we have instructions other than load and store, we cannot do the
- // tree structured merge
+ // If we have instructions other than load and store, we cannot do
+ // the tree structured merge.
return std::nullopt;
}
}
- // If we do not have any load, we cannot do the tree structured merge
- if (!TheLoad)
+
+ // Classify the pattern by looking at what we collected:
+ // Pattern 1 (stores-only): only partial stores + exactly one full load.
+ // Pattern 2 (RMW): one full init store + partial loads + partial stores
+ // (+ optional full final load). RMW also needs VecTy to be set
+ // because we use getIndex() to convert byte offsets to element
+ // indices, which requires a promoted vector alloca.
+ bool IsRMWPattern =
+ InitStore != nullptr && !LoadInfos.empty() && VecTy != nullptr;
+ bool IsStoresOnlyPattern =
+ InitStore == nullptr && LoadInfos.empty() && FullLoad != nullptr;
+ if (!IsRMWPattern && !IsStoresOnlyPattern)
return std::nullopt;
- // If we do not have multiple stores, we cannot do the tree structured merge
+ // Need at least two partial stores to benefit from tree-merging; a
+ // single store is already optimal as-is.
if (StoreInfos.size() < 2)
return std::nullopt;
- // Stores should not overlap and should cover the whole alloca
- // Sort by begin offset
- llvm::sort(StoreInfos, [](const StoreInfo &A, const StoreInfo &B) {
- return A.BeginOffset < B.BeginOffset;
- });
+ // All partial stores must live in the same basic block — the tree merge
+ // is built in a single BB using block-order ordering (comesBefore).
+ BasicBlock *StoreBB = StoreInfos[0].Store->getParent();
+ for (auto &Info : StoreInfos)
+ if (Info.Store->getParent() != StoreBB)
+ return std::nullopt;
+
+ SmallVector<Value *, 4> DeletedValues;
- // Check for overlaps and coverage
- uint64_t ExpectedStart = NewAllocaBeginOffset;
- for (auto &StoreInfo : StoreInfos) {
- uint64_t BeginOff = StoreInfo.BeginOffset;
- uint64_t EndOff = StoreInfo.EndOffset;
+ // Helper: run the pairwise tree merge on a queue of vectors and return
+ // the merged result. At each iteration of the while-loop we pop pairs
+ // of vectors, concatenate them with mergeTwoVectors, and push the
+ // result back. An odd leftover is rotated to the end so it pairs up in
+ // the next iteration. Final queue size == 1 is the fully-merged vector.
+ auto TreeMerge = [&](std::queue<Value *> &Q, IRBuilder<> &B) -> Value * {
+ LLVM_DEBUG(dbgs() << " Rewrite stores into shufflevectors:\n");
+ while (Q.size() > 1) {
+ const auto N = Q.size();
+ for ([[maybe_unused]] const auto _ : llvm::seq(N / 2)) {
+ Value *V0 = Q.front();
+ Q.pop();
+ Value *V1 = Q.front();
+ Q.pop();
+ Value *M = mergeTwoVectors(V0, V1, DL, AllocatedEltTy, B);
+ LLVM_DEBUG(dbgs() << " shufflevector: " << *M << "\n");
+ Q.push(M);
+ }
+ if (N % 2 == 1) {
+ Value *V = Q.front();
+ Q.pop();
+ Q.push(V);
+ }
+ }
+ return Q.front();
+ };
- // Check for gap or overlap
- if (BeginOff != ExpectedStart)
+ if (IsStoresOnlyPattern) {
+ // Stores should not overlap and should cover the whole alloca.
+ // Sort by begin offset to verify this with a single linear scan.
+ llvm::sort(StoreInfos, [](const StoreInfo &A, const StoreInfo &B) {
+ return A.BeginOffset < B.BeginOffset;
+ });
+ // Check for gap or overlap: each begin offset must equal the previous
+ // end offset, i.e. the store ranges must tile [NewAllocaBeginOffset,
+ // NewAllocaEndOffset) exactly.
+ uint64_t Expected = NewAllocaBeginOffset;
+ for (auto &Info : StoreInfos) {
+ if (Info.BeginOffset != Expected)
+ return std::nullopt;
+ Expected = Info.EndOffset;
+ }
+ // Stores cover the entire alloca (no trailing gap either).
+ if (Expected != NewAllocaEndOffset)
return std::nullopt;
- ExpectedStart = EndOff;
+ // The load should not be in the middle of the stores.
+ // Note:
+ // If the load is in a different basic block from the stores, we can
+ // still do the tree-structured merge. We don't have store->load
+ // forwarding here — the merged vector is stored back to NewAI and
+ // the new load loads from NewAI. The forwarding will be handled
+ // later when NewAI is promoted.
+ BasicBlock *LoadBB = FullLoad->getParent();
+ if (LoadBB == StoreBB) {
+ for (auto &Info : StoreInfos)
+ if (!Info.Store->comesBefore(FullLoad))
+ return std::nullopt;
+ }
+
+ LLVM_DEBUG({
+ dbgs() << "Tree structured merge rewrite (stores-only):\n";
+ dbgs() << " Load: " << *FullLoad << "\n Ordered stores:\n";
+ for (auto [i, Info] : enumerate(StoreInfos))
+ dbgs() << " [" << i << "] Range[" << Info.BeginOffset << ", "
+ << Info.EndOffset << ") \tStore: " << *Info.Store
+ << "\tValue: " << *Info.StoredValue << "\n";
+ });
+
+ // StoreInfos is sorted by offset, not by block order. Anchoring to
+ // StoreInfos.back().Store (last by offset) can place shuffles before
+ // operands that appear later in the block (invalid SSA). Insert before
+ // FullLoad when it shares the store block (after all stores, before
+ // any later IR in that block). Otherwise insert before the store
+ // block's terminator so the merge runs after every store and any
+ // trailing instructions in that block.
+ IRBuilder<> Builder(LoadBB == StoreBB ? cast<Instruction>(FullLoad)
+ : StoreBB->getTerminator());
+ std::queue<Value *> Q;
+ for (const auto &Info : StoreInfos) {
+ DeletedValues.push_back(Info.Store);
+ Q.push(Info.StoredValue);
+ }
+ // Merge all stored values and store the merged value into the alloca.
+ Value *Merged = TreeMerge(Q, Builder);
+ Builder.CreateAlignedStore(Merged, &NewAI, getSliceAlign());
+
+ // Replace the original load with a load of the newly-merged alloca.
+ // Later promotion will forward the store we just created to this load.
+ IRBuilder<> LoadBuilder(FullLoad);
+ FullLoad->replaceAllUsesWith(LoadBuilder.CreateAlignedLoad(
+ FullLoad->getType(), &NewAI, getSliceAlign(), FullLoad->isVolatile(),
+ FullLoad->getName() + ".sroa.new.load"));
+ DeletedValues.push_back(FullLoad);
+ return DeletedValues;
}
- // Check that stores cover the entire alloca
- if (ExpectedStart != NewAllocaEndOffset)
- return std::nullopt;
- // Stores should be in the same basic block
- // The load should not be in the middle of the stores
- // Note:
- // If the load is in a different basic block with the stores, we can still
- // do the tree structured merge. This is because we do not have the
- // store->load forwarding here. The merged vector will be stored back to
- // NewAI and the new load will load from NewAI. The forwarding will be
- // handled later when we try to promote NewAI.
- BasicBlock *LoadBB = TheLoad->getParent();
- BasicBlock *StoreBB = StoreInfos[0].Store->getParent();
+ // RMW pattern handling starts from here.
+ // Like StoreBB above: keep the init store, all partial loads and all
+ // partial stores in one basic block so we can reason about ordering
+ // with comesBefore and build SSA without PHIs.
+ if (InitStore->getParent() != StoreBB)
+ return std::nullopt;
+ for (auto &Info : LoadInfos)
+ if (Info.Load->getParent() != StoreBB)
+ return std::nullopt;
+ // FullLoad (if any) is allowed to live in a different basic block. See
+ // the note on the stores-only path: we don't do store->load forwarding
+ // directly — the merged vector is stored to NewAI and the new load
+ // loads from NewAI, so cross-BB ordering is resolved later when NewAI
+ // is promoted.
+
+ // Collect the combined partial-load/partial-store accesses sorted
+ // by block order. Used both for ordering checks and for the rewrite
+ // walk below.
+ struct Access {
+ Instruction *Inst;
+ uint64_t BeginOffset;
+ uint64_t EndOffset;
+ bool IsStore;
+ };
+ SmallVector<Access, 16> Accesses;
+ Accesses.reserve(LoadInfos.size() + StoreInfos.size());
----------------
Prince781 wrote:
```suggestion
SmallVector<Access, 16> Accesses(LoadInfos.size() + StoreInfos.size());
```
https://github.com/llvm/llvm-project/pull/194441
More information about the llvm-commits
mailing list