[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:50 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());
+ for (auto &Info : LoadInfos)
+ Accesses.push_back({Info.Load, Info.BeginOffset, Info.EndOffset,
+ /*IsStore=*/false});
+ for (auto &Info : StoreInfos)
+ Accesses.push_back({Info.Store, Info.BeginOffset, Info.EndOffset,
+ /*IsStore=*/true});
+ llvm::sort(Accesses, [](const Access &A, const Access &B) {
+ return A.Inst->comesBefore(B.Inst);
+ });
- for (auto &StoreInfo : StoreInfos) {
- if (StoreInfo.Store->getParent() != StoreBB)
+ // Ordering constraint 1: InitStore must come before every partial
+ // access — they read/write the RMW state initialised by InitStore.
+ for (auto &Acc : Accesses)
+ if (!InitStore->comesBefore(Acc.Inst))
return std::nullopt;
- if (LoadBB == StoreBB && !StoreInfo.Store->comesBefore(TheLoad))
+ // Ordering constraint 2: when FullLoad shares the block with the
+ // partial accesses, it must come after every one of them — otherwise
+ // it could read a stale value. If FullLoad is in a different block
+ // the later promotion pass resolves the cross-BB ordering.
+ if (FullLoad && FullLoad->getParent() == StoreBB)
+ for (auto &Acc : Accesses)
+ if (!Acc.Inst->comesBefore(FullLoad))
+ return std::nullopt;
+
+ // Coverage check: the distinct (begin, end) ranges touched by the
+ // partial loads and stores must tile the alloca disjointly. That is
+ // the only precondition the per-range SliceValues tracking below
+ // needs — a disjoint tile guarantees the entries don't alias each
+ // other. We don't check per-range load/store counts: a range with
+ // only loads ends with SliceValues[r] = the init extract
+ // (contributed to the final tree-merge), and a range with only
+ // stores ends with SliceValues[r] = its last stored value. Both are
+ // correct.
+ using SliceRange = std::pair<uint64_t, uint64_t>;
+ DenseSet<SliceRange> TouchedRanges;
+ for (auto &Info : LoadInfos)
+ TouchedRanges.insert({Info.BeginOffset, Info.EndOffset});
+ for (auto &Info : StoreInfos)
+ TouchedRanges.insert({Info.BeginOffset, Info.EndOffset});
+ SmallVector<SliceRange, 8> Partition(TouchedRanges.begin(),
+ TouchedRanges.end());
+ llvm::sort(Partition);
+ // Disjoint + contiguous tile of the whole alloca.
+ uint64_t Expected = NewAllocaBeginOffset;
+ for (auto &Range : Partition) {
+ if (Range.first != Expected)
return std::nullopt;
+ Expected = Range.second;
}
+ if (Expected != NewAllocaEndOffset)
+ return std::nullopt;
- // If we reach here, the partition can be merged with a tree structured
- // merge
LLVM_DEBUG({
- dbgs() << "Tree structured merge rewrite:\n Load: " << *TheLoad
- << "\n Ordered stores:\n";
- for (auto [i, Info] : enumerate(StoreInfos))
- dbgs() << " [" << i << "] Range[" << Info.BeginOffset << ", "
- << Info.EndOffset << ") \tStore: " << *Info.Store
- << "\tValue: " << *Info.StoredValue << "\n";
+ dbgs() << "Tree structured merge rewrite (RMW):\n";
+ dbgs() << " Init store: " << *InitStore << "\n";
+ if (FullLoad)
+ dbgs() << " Final load: " << *FullLoad << "\n";
+ dbgs() << " Slice ranges (" << Partition.size() << "):\n";
+ for (auto &Range : Partition)
+ dbgs() << " [" << Range.first << ", " << Range.second << ")\n";
});
- // Instead of having these stores, we merge all the stored values into a
- // vector and store the merged value into the alloca
- std::queue<Value *> VecElements;
- // 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
- // TheLoad 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 ? TheLoad : StoreBB->getTerminator());
- for (const auto &Info : StoreInfos) {
- DeletedValues.push_back(Info.Store);
- VecElements.push(Info.StoredValue);
+ // Initialize SliceValues: one SSA value per slice range, tracking
+ // the value the alloca currently holds at that range. Each entry
+ // starts at the corresponding piece of the init store, obtained by
+ // bitcasting the init value to the alloca's vector type (if needed)
+ // and extracting the slice's sub-range.
+ IRB.SetInsertPoint(InitStore->getNextNode());
+ Value *InitVec = InitValue;
+ if (InitVec->getType() != NewAllocaTy)
+ InitVec = IRB.CreateBitCast(InitVec, NewAllocaTy, "init.cast");
+ DenseMap<SliceRange, Value *> SliceValues;
+ for (auto &Range : Partition) {
+ Value *V = extractVector(IRB, InitVec, getIndex(Range.first),
+ getIndex(Range.second), "init.extract");
+ // extractVector emits an extractelement (scalar) for 1-element
+ // ranges. The tree merge consumes vectors uniformly, so wrap any
+ // scalar into <1 x T> here. Stored values are guaranteed to be
+ // FixedVectorType by IsTypeValidForTreeStructuredMerge, so once
+ // every entry is a vector it stays a vector throughout the walk.
+ if (!isa<FixedVectorType>(V->getType()))
+ V = IRB.CreateBitCast(V, FixedVectorType::get(V->getType(), 1),
+ "init.extract.vec");
----------------
Prince781 wrote:
Could you use a 1-element `shufflevector` instead of `extract` here?
https://github.com/llvm/llvm-project/pull/194441
More information about the llvm-commits
mailing list