[llvm] [SROA] Extend tree-structured merge to handle init + RMW pattern (PR #194441)
via llvm-commits
llvm-commits at lists.llvm.org
Tue May 5 13:47:52 PDT 2026
https://github.com/Chengjunp updated https://github.com/llvm/llvm-project/pull/194441
>From ee893ce621f60877b739b475c81264c6046c0a7e Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Mon, 27 Apr 2026 19:09:22 +0000
Subject: [PATCH 1/2] Extend SROA tree merge to handle RMW pattern
---
llvm/lib/Transforms/Scalar/SROA.cpp | 483 +++++++++++++-----
.../vector-promotion-rmw-cannot-tree-merge.ll | 79 +++
.../SROA/vector-promotion-rmw-tree-merge.ll | 210 ++++++++
3 files changed, 643 insertions(+), 129 deletions(-)
create mode 100644 llvm/test/Transforms/SROA/vector-promotion-rmw-cannot-tree-merge.ll
create mode 100644 llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 8b1bde5f44cb5..6b52873b49cb3 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2856,12 +2856,14 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
/// Attempts to rewrite a partition using tree-structured merge optimization.
///
- /// This function analyzes a partition to determine if it can be optimized
- /// using a tree-structured merge pattern, where multiple non-overlapping
- /// stores completely fill an alloca. And there is no load from the alloca in
- /// the middle of the stores. Such patterns can be optimized by eliminating
- /// the intermediate stores and directly constructing the final vector by
- /// using shufflevectors.
+ /// This function handles two patterns. Both produce an O(log n) tree of
+ /// shufflevectors in place of the linear expand+blend chain that SROA would
+ /// otherwise emit for each partial store.
+ ///
+ /// Pattern 1 (stores-only):
+ /// Multiple non-overlapping partial stores completely fill the alloca and
+ /// there is exactly one full-width load. The stores are tree-merged into
+ /// a single vector and stored once.
///
/// Example transformation:
/// Before: (stores do not have to be in order)
@@ -2870,37 +2872,46 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
/// store <2 x float> %val2, ptr %alloca+16 ; offset 4-5
/// store <2 x float> %val1, ptr %alloca+8 ; offset 2-3
/// store <2 x float> %val3, ptr %alloca+24 ; offset 6-7
+ /// %r = load <8 x float>, ptr %alloca
///
- /// After:
- /// %alloca = alloca <8 x float>
- /// %shuffle0 = shufflevector %val0, %val1, <4 x i32> <i32 0, i32 1, i32 2,
- /// i32 3>
- /// %shuffle1 = shufflevector %val2, %val3, <4 x i32> <i32 0, i32 1, i32 2,
- /// i32 3>
- /// %shuffle2 = shufflevector %shuffle0, %shuffle1, <8 x i32> <i32 0, i32 1,
- /// i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
- /// store %shuffle2, ptr %alloca
+ /// After: tree of shufflevectors producing <8 x float> directly.
+ ///
+ /// Pattern 2 (init + RMW, possibly multi-round):
+ /// A single full-width init store, followed by partial loads and
+ /// partial stores that read-modify-write the alloca one or more
+ /// times, optionally followed by a full-width load. The only
+ /// structural requirement is that the distinct (begin, end) ranges
+ /// touched by the partial loads and stores, taken together, tile
+ /// the alloca disjointly.
+ ///
+ /// We keep a map from each slice range to the SSA value that
+ /// currently lives there, `SliceValues[r] -> Value*`:
+ /// - initialize each entry from the init store via extractvector,
+ /// - walk partial loads and stores in block order,
+ /// - for a partial load at range r: RAUW with `SliceValues[r]`,
+ /// - for a partial store at range r: update `SliceValues[r]` to
+ /// the stored value and drop the store.
+ /// At the end, the final `SliceValues[r]` entries are tree-merged
+ /// (in range order) into a single store to the alloca, and the
+ /// optional full-width load is replaced by a load of the alloca.
///
- /// The optimization looks for partitions that:
- /// 1. Have no overlapping split slice tails
- /// 2. Contain non-overlapping stores that cover the entire alloca
- /// 3. Have exactly one load that reads the complete alloca structure and not
- /// in the middle of the stores (TODO: maybe we can relax the constraint
- /// about reading the entire alloca structure)
+ /// Because the ranges are disjoint by construction, a store at one
+ /// range cannot affect another range's tracked value, so a single
+ /// block-order walk correctly tracks the memory state at each
+ /// range. The algorithm handles multi-round RMW, loads/stores
+ /// interleaved in any order, read-only slices (entry stays at the
+ /// init extract) and write-only slices (entry never flows into a
+ /// load).
///
/// \param P The partition to analyze and potentially rewrite
- /// \return An optional vector of values that were deleted during the rewrite
- /// process, or std::nullopt if the partition cannot be optimized
- /// using tree-structured merge
+ /// \return An optional vector of values that were deleted during the
+ /// rewrite, or std::nullopt if the partition cannot be optimized.
std::optional<SmallVector<Value *, 4>>
rewriteTreeStructuredMerge(Partition &P) {
// No tail slices that overlap with the partition
if (P.splitSliceTails().size() > 0)
return std::nullopt;
- SmallVector<Value *, 4> DeletedValues;
- LoadInst *TheLoad = nullptr;
-
// Structure to hold store information
struct StoreInfo {
StoreInst *Store;
@@ -2910,8 +2921,17 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
StoreInfo(StoreInst *SI, uint64_t Begin, uint64_t End, Value *Val)
: Store(SI), BeginOffset(Begin), EndOffset(End), StoredValue(Val) {}
};
+ struct LoadInfo {
+ LoadInst *Load;
+ uint64_t BeginOffset;
+ uint64_t EndOffset;
+ };
- SmallVector<StoreInfo, 4> StoreInfos;
+ SmallVector<StoreInfo, 4> StoreInfos; // partial stores only
+ SmallVector<LoadInfo, 4> LoadInfos; // partial loads only
+ LoadInst *FullLoad = nullptr; // optional full-width load
+ StoreInst *InitStore = nullptr; // optional full-width init store
+ Value *InitValue = nullptr;
// If the new alloca is a fixed vector type, we use its element type as the
// allocated element type, otherwise we use i8 as the allocated element
@@ -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");
+ SliceValues[Range] = V;
}
-
- LLVM_DEBUG(dbgs() << " Rewrite stores into shufflevectors:\n");
- while (VecElements.size() > 1) {
- const auto NumElts = VecElements.size();
- for ([[maybe_unused]] const auto _ : llvm::seq(NumElts / 2)) {
- Value *V0 = VecElements.front();
- VecElements.pop();
- Value *V1 = VecElements.front();
- VecElements.pop();
- Value *Merged = mergeTwoVectors(V0, V1, DL, AllocatedEltTy, Builder);
- LLVM_DEBUG(dbgs() << " shufflevector: " << *Merged << "\n");
- VecElements.push(Merged);
- }
- if (NumElts % 2 == 1) {
- Value *V = VecElements.front();
- VecElements.pop();
- VecElements.push(V);
+ // The init store itself becomes dead — its value is consumed via the
+ // extracts above.
+ DeletedValues.push_back(InitStore);
+
+ // Walk accesses in block order:
+ // - partial load at range r: replace with SliceValues[r] (bitcast
+ // if the load's type differs from the current tracked value's
+ // type, e.g. because a previous store wrote a vector with a
+ // different element type);
+ // - partial store at range r: update SliceValues[r] to the stored
+ // value and drop the store.
+ for (auto &Acc : Accesses) {
+ SliceRange Range{Acc.BeginOffset, Acc.EndOffset};
+ if (!Acc.IsStore) {
+ Value *V = SliceValues[Range];
+ if (V->getType() != Acc.Inst->getType()) {
+ IRB.SetInsertPoint(cast<LoadInst>(Acc.Inst));
+ V = IRB.CreateBitCast(V, Acc.Inst->getType());
+ }
+ Acc.Inst->replaceAllUsesWith(V);
+ DeletedValues.push_back(Acc.Inst);
+ } else {
+ SliceValues[Range] = cast<StoreInst>(Acc.Inst)->getValueOperand();
+ DeletedValues.push_back(Acc.Inst);
}
}
- // Store the merged value into the alloca
- Value *MergedValue = VecElements.front();
- Builder.CreateAlignedStore(MergedValue, &NewAI, getSliceAlign());
+ // Tree-merge the final per-range values (in range order) into the
+ // alloca's final vector value. Insertion point is just after the
+ // block-order-last partial access — guaranteed to dominate every
+ // SliceValues entry since each one is either an init extract or a
+ // stored value defined before its (now-deleted) store.
+ Instruction *FinalInsert = Accesses.back().Inst->getNextNode();
+ if (!FinalInsert)
+ return std::nullopt;
- IRBuilder<> LoadBuilder(TheLoad);
- TheLoad->replaceAllUsesWith(LoadBuilder.CreateAlignedLoad(
- TheLoad->getType(), &NewAI, getSliceAlign(), TheLoad->isVolatile(),
- TheLoad->getName() + ".sroa.new.load"));
- DeletedValues.push_back(TheLoad);
+ IRBuilder<> Builder(FinalInsert);
+ std::queue<Value *> Q;
+ for (auto &Range : Partition)
+ Q.push(SliceValues[Range]);
+ Value *Merged = TreeMerge(Q, Builder);
+ Builder.CreateAlignedStore(Merged, &NewAI, getSliceAlign());
+
+ // Replace the optional final full-width load with a load of the newly
+ // merged alloca. Later promotion will forward the store above to it.
+ if (FullLoad) {
+ IRBuilder<> LoadBuilder(FullLoad);
+ FullLoad->replaceAllUsesWith(LoadBuilder.CreateAlignedLoad(
+ FullLoad->getType(), &NewAI, getSliceAlign(), FullLoad->isVolatile(),
+ FullLoad->getName() + ".sroa.new.load"));
+ DeletedValues.push_back(FullLoad);
+ }
return DeletedValues;
}
diff --git a/llvm/test/Transforms/SROA/vector-promotion-rmw-cannot-tree-merge.ll b/llvm/test/Transforms/SROA/vector-promotion-rmw-cannot-tree-merge.ll
new file mode 100644
index 0000000000000..1fd8c5dd3fe67
--- /dev/null
+++ b/llvm/test/Transforms/SROA/vector-promotion-rmw-cannot-tree-merge.ll
@@ -0,0 +1,79 @@
+; REQUIRES: asserts
+; RUN: opt < %s -passes='sroa<preserve-cfg>' -disable-output -debug-only=sroa 2>&1 | FileCheck %s
+; RUN: opt < %s -passes='sroa<modify-cfg>' -disable-output -debug-only=sroa 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64"
+
+;
+; Negative tests: patterns that should NOT trigger the init + RMW tree-merge
+; rewrite. Each function violates one precondition; the check below is
+; global, so none of them may emit the debug line "Tree structured merge
+; rewrite (RMW):" that the rewrite prints when it fires.
+;
+
+; CHECK-NOT: Tree structured merge rewrite (RMW):
+
+; Coverage gap: partial loads/stores don't tile the alloca — element 2
+; is neither loaded nor stored, so the partition-tile check fails.
+define <4 x float> @coverage_gap(<4 x float> %init, <2 x float> %a, <1 x float> %b) {
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %r0 = fadd <2 x float> %v0, %a
+ store <2 x float> %r0, ptr %p0, align 8
+
+ ; (nothing touches element 2)
+
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 3
+ %v1 = load <1 x float>, ptr %p1, align 4
+ %r1 = fadd <1 x float> %v1, %b
+ store <1 x float> %r1, ptr %p1, align 4
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Missing init store: the alloca has partial loads/stores but no
+; full-width init, so the RMW path (which needs an init seed) refuses.
+define <4 x float> @missing_init_store(<2 x float> %a, <2 x float> %b) {
+entry:
+ %alloca = alloca [4 x float], align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %r0 = fadd <2 x float> %v0, %a
+ store <2 x float> %r0, ptr %p0, align 8
+
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+ %v1 = load <2 x float>, ptr %p1, align 8
+ %r1 = fadd <2 x float> %v1, %b
+ store <2 x float> %r1, ptr %p1, align 8
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Atomic RMW: the partial load/store at slice 0 carry memory-ordering
+; (monotonic). The rewrite would drop the atomic stores and elide the
+; atomic loads, losing that ordering, so the RMW path must reject it.
+define <4 x float> @atomic_rmw(<4 x float> %init, <2 x float> %a, <2 x float> %b) {
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load atomic <2 x float>, ptr %p0 monotonic, align 8
+ %r0 = fadd <2 x float> %v0, %a
+ store atomic <2 x float> %r0, ptr %p0 monotonic, align 8
+
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+ %v1 = load <2 x float>, ptr %p1, align 8
+ %r1 = fadd <2 x float> %v1, %b
+ store <2 x float> %r1, ptr %p1, align 8
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
diff --git a/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll b/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
new file mode 100644
index 0000000000000..c569dac8353ca
--- /dev/null
+++ b/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
@@ -0,0 +1,210 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes='sroa<preserve-cfg>' -S | FileCheck %s
+; RUN: opt < %s -passes='sroa<modify-cfg>' -S | FileCheck %s
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64"
+
+;
+; Positive tests for the init + RMW tree-merge rewrite in SROA.
+;
+; The pattern this exercises:
+; - a single full-width init store at the top,
+; - one or more rounds of per-slice load-process-store,
+; - a single full-width load at the end.
+;
+; The rewrite should:
+; - replace each partial load with an extractvector from the current
+; per-slice value,
+; - thread each slice's SSA value through the rounds via the user's
+; compute (no selects/blends),
+; - tree-merge the final per-slice values into the final alloca value.
+;
+
+; Canonical single-round RMW: 2 slices, each loaded, processed, stored back.
+define <4 x float> @single_round(<4 x float> %init, <2 x float> %a, <2 x float> %b) {
+; CHECK-LABEL: define <4 x float> @single_round(
+; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A:%.*]], <2 x float> [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT1]], [[B]]
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[R0]], <2 x float> [[R1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: ret <4 x float> [[TMP0]]
+;
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %r0 = fadd <2 x float> %v0, %a
+ store <2 x float> %r0, ptr %p0, align 8
+
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+ %v1 = load <2 x float>, ptr %p1, align 8
+ %r1 = fadd <2 x float> %v1, %b
+ store <2 x float> %r1, ptr %p1, align 8
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Multi-round RMW: the same 2-slice partition is updated twice.
+; Round 2's load at each slice should read round 1's stored value
+; (the per-slice SSA chain), not the init value.
+define <4 x float> @multi_round(<4 x float> %init,
+; CHECK-LABEL: define <4 x float> @multi_round(
+; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A1:%.*]], <2 x float> [[B1:%.*]], <2 x float> [[A2:%.*]], <2 x float> [[B2:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[R0_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A1]]
+; CHECK-NEXT: [[R1_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT1]], [[B1]]
+; CHECK-NEXT: [[R0_2:%.*]] = fadd <2 x float> [[R0_1]], [[A2]]
+; CHECK-NEXT: [[R1_2:%.*]] = fadd <2 x float> [[R1_1]], [[B2]]
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[R0_2]], <2 x float> [[R1_2]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: ret <4 x float> [[TMP0]]
+;
+ <2 x float> %a1, <2 x float> %b1,
+ <2 x float> %a2, <2 x float> %b2) {
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+
+ ; round 1
+ %v0.1 = load <2 x float>, ptr %p0, align 8
+ %r0.1 = fadd <2 x float> %v0.1, %a1
+ store <2 x float> %r0.1, ptr %p0, align 8
+ %v1.1 = load <2 x float>, ptr %p1, align 8
+ %r1.1 = fadd <2 x float> %v1.1, %b1
+ store <2 x float> %r1.1, ptr %p1, align 8
+
+ ; round 2 — loads should see round 1's stores, not init
+ %v0.2 = load <2 x float>, ptr %p0, align 8
+ %r0.2 = fadd <2 x float> %v0.2, %a2
+ store <2 x float> %r0.2, ptr %p0, align 8
+ %v1.2 = load <2 x float>, ptr %p1, align 8
+ %r1.2 = fadd <2 x float> %v1.2, %b2
+ store <2 x float> %r1.2, ptr %p1, align 8
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Uneven slice sizes: 3 + 1 at offsets 0 and 3. The rewrite's tree merge
+; must handle vectors of different widths (via mergeTwoVectors).
+define <4 x float> @uneven_slices(<4 x float> %init, <3 x float> %a, <1 x float> %b) {
+; CHECK-LABEL: define <4 x float> @uneven_slices(
+; CHECK-SAME: <4 x float> [[INIT:%.*]], <3 x float> [[A:%.*]], <1 x float> [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <3 x i32> <i32 0, i32 1, i32 2>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = extractelement <4 x float> [[INIT]], i32 3
+; CHECK-NEXT: [[TMP0:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT1]] to <1 x float>
+; CHECK-NEXT: [[R0:%.*]] = fadd <3 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[TMP0]], [[B]]
+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <1 x float> [[R1]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
+; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <3 x float> [[R0]], <3 x float> [[TMP1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: ret <4 x float> [[TMP2]]
+;
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <3 x float>, ptr %p0, align 4
+ %r0 = fadd <3 x float> %v0, %a
+ store <3 x float> %r0, ptr %p0, align 4
+
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 3
+ %v1 = load <1 x float>, ptr %p1, align 4
+ %r1 = fadd <1 x float> %v1, %b
+ store <1 x float> %r1, ptr %p1, align 4
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Direct load-to-store (slice swap): every partial store's value operand
+; is exactly a previously-loaded slice, with no arithmetic in between. The
+; rewriter must read each store's operand *after* load rewrites have been
+; applied — otherwise the store records the (now-dead) original load and
+; the final tree merge ends up consuming a poison value.
+define <4 x float> @slice_swap(<4 x float> %init) {
+; CHECK-LABEL: define <4 x float> @slice_swap(
+; CHECK-SAME: <4 x float> [[INIT:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[INIT_EXTRACT_EXTRACT1]], <2 x float> [[INIT_EXTRACT_EXTRACT]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: ret <4 x float> [[TMP0]]
+;
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %v1 = load <2 x float>, ptr %p1, align 8
+
+ ; slice0 <- slice1's loaded value; slice1 <- slice0's loaded value.
+ store <2 x float> %v1, ptr %p0, align 8
+ store <2 x float> %v0, ptr %p1, align 8
+
+ %result = load <4 x float>, ptr %alloca, align 16
+ ret <4 x float> %result
+}
+
+; Load-only single-element slice. The init seed for a 1-element range is
+; produced by extractelement (a scalar), and a slice with no partial store
+; never updates that entry.
+define <4 x float> @load_only_single_element_slice(<4 x float> %init, <2 x float> %a, <1 x float> %b) {
+; CHECK-LABEL: define <4 x float> @load_only_single_element_slice(
+; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A:%.*]], <1 x float> [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = extractelement <4 x float> [[INIT]], i32 2
+; CHECK-NEXT: [[INIT_EXTRACT_VEC:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT1]] to <1 x float>
+; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT2:%.*]] = extractelement <4 x float> [[INIT]], i32 3
+; CHECK-NEXT: [[INIT_EXTRACT_VEC3:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT2]] to <1 x float>
+; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[INIT_EXTRACT_VEC]], [[B]]
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <1 x float> [[R1]], <1 x float> poison, <2 x i32> <i32 0, i32 poison>
+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x float> [[R0]], <2 x float> [[TMP0]], <3 x i32> <i32 0, i32 1, i32 2>
+; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <1 x float> [[INIT_EXTRACT_VEC3]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
+; CHECK-NEXT: [[TMP3:%.*]] = shufflevector <3 x float> [[TMP1]], <3 x float> [[TMP2]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[V2_FROM_INIT:%.*]] = bitcast <1 x float> [[INIT_EXTRACT_VEC3]] to float
+; CHECK-NEXT: [[FINAL:%.*]] = insertelement <4 x float> [[TMP3]], float [[V2_FROM_INIT]], i32 0
+; CHECK-NEXT: ret <4 x float> [[FINAL]]
+;
+entry:
+ %alloca = alloca [4 x float], align 16
+ store <4 x float> %init, ptr %alloca, align 16
+
+ ; slice 0: load+store at (0,2)
+ %p0 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %r0 = fadd <2 x float> %v0, %a
+ store <2 x float> %r0, ptr %p0, align 8
+
+ ; slice 1: load+store at (2,3) — single-element with a store
+ %p1 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 2
+ %v1 = load <1 x float>, ptr %p1, align 4
+ %r1 = fadd <1 x float> %v1, %b
+ store <1 x float> %r1, ptr %p1, align 4
+
+ ; slice 2: load-only at (3,4) — single-element, no store (the crash case)
+ %p2 = getelementptr inbounds [4 x float], ptr %alloca, i32 0, i32 3
+ %v2 = load <1 x float>, ptr %p2, align 4
+
+ %result = load <4 x float>, ptr %alloca, align 16
+
+ ; consume %v2 so the load isn't trivially dead before SROA runs
+ %v2_from_init = bitcast <1 x float> %v2 to float
+ %final = insertelement <4 x float> %result, float %v2_from_init, i32 0
+ ret <4 x float> %final
+}
>From 3383b9741934639cec2305de15d728104fbb594d Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Tue, 5 May 2026 20:43:23 +0000
Subject: [PATCH 2/2] Fix issue raised by the review
---
llvm/lib/Transforms/Scalar/SROA.cpp | 96 +++++------
.../SROA/vector-promotion-rmw-tree-merge.ll | 149 ++++++++++++++----
2 files changed, 165 insertions(+), 80 deletions(-)
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 6b52873b49cb3..3552c82571b8c 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2861,9 +2861,9 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
/// otherwise emit for each partial store.
///
/// Pattern 1 (stores-only):
- /// Multiple non-overlapping partial stores completely fill the alloca and
- /// there is exactly one full-width load. The stores are tree-merged into
- /// a single vector and stored once.
+ /// Multiple non-overlapping partial stores completely fill the alloca
+ /// and there is exactly one full-width load coming after the stores.
+ /// The stores are tree-merged into a single vector and stored once.
///
/// Example transformation:
/// Before: (stores do not have to be in order)
@@ -2880,13 +2880,15 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
/// A single full-width init store, followed by partial loads and
/// partial stores that read-modify-write the alloca one or more
/// times, optionally followed by a full-width load. The only
- /// structural requirement is that the distinct (begin, end) ranges
+ /// structural requirement is that the distinct [begin, end) ranges
/// touched by the partial loads and stores, taken together, tile
/// the alloca disjointly.
///
/// We keep a map from each slice range to the SSA value that
/// currently lives there, `SliceValues[r] -> Value*`:
- /// - initialize each entry from the init store via extractvector,
+ /// - initialize each entry to the corresponding piece of the
+ /// init store's value (via a shufflevector picking the
+ /// range's elements out of the init value),
/// - walk partial loads and stores in block order,
/// - for a partial load at range r: RAUW with `SliceValues[r]`,
/// - for a partial store at range r: update `SliceValues[r]` to
@@ -2898,10 +2900,10 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
/// Because the ranges are disjoint by construction, a store at one
/// range cannot affect another range's tracked value, so a single
/// block-order walk correctly tracks the memory state at each
- /// range. The algorithm handles multi-round RMW, loads/stores
- /// interleaved in any order, read-only slices (entry stays at the
- /// init extract) and write-only slices (entry never flows into a
- /// load).
+ /// range. The algorithm handles multi-round RMW, partial loads
+ /// and stores interleaved in any order, read-only slices (the
+ /// tracked value stays at the init extract), and write-only
+ /// slices (the tracked value never flows into a load).
///
/// \param P The partition to analyze and potentially rewrite
/// \return An optional vector of values that were deleted during the
@@ -3100,10 +3102,11 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
LLVM_DEBUG({
dbgs() << "Tree structured merge rewrite (stores-only):\n";
dbgs() << " Load: " << *FullLoad << "\n Ordered stores:\n";
- for (auto [i, Info] : enumerate(StoreInfos))
+ 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
@@ -3140,9 +3143,10 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
// 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;
+ if (any_of(LoadInfos, [&](const LoadInfo &I) {
+ return I.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
@@ -3157,34 +3161,38 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
uint64_t BeginOffset;
uint64_t EndOffset;
bool IsStore;
+ Access(const LoadInfo &L)
+ : Inst(L.Load), BeginOffset(L.BeginOffset), EndOffset(L.EndOffset),
+ IsStore(false) {}
+ Access(const StoreInfo &S)
+ : Inst(S.Store), BeginOffset(S.BeginOffset), EndOffset(S.EndOffset),
+ IsStore(true) {}
};
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});
+ Accesses.append(LoadInfos.begin(), LoadInfos.end());
+ Accesses.append(StoreInfos.begin(), StoreInfos.end());
llvm::sort(Accesses, [](const Access &A, const Access &B) {
return A.Inst->comesBefore(B.Inst);
});
// 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;
+ // Accesses is sorted by block order, so the first element is the
+ // earliest; checking it is enough.
+ if (!InitStore->comesBefore(Accesses.front().Inst))
+ return std::nullopt;
// 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;
+ // it could read a stale value. Accesses is sorted, so the last
+ // element is the latest; checking it is enough. If FullLoad is in a
+ // different block the later promotion pass resolves the cross-BB
+ // ordering.
+ if (FullLoad && FullLoad->getParent() == StoreBB &&
+ !Accesses.back().Inst->comesBefore(FullLoad))
+ return std::nullopt;
- // Coverage check: the distinct (begin, end) ranges touched by the
+ // 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
@@ -3233,17 +3241,15 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
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");
- SliceValues[Range] = V;
+ // Initialize SliceValues[Range] with the piece of InitVec that
+ // lives at this range — a shufflevector picking the
+ // [BeginIdx, EndIdx) elements out of InitVec. This is the value
+ // any subsequent partial load at this range reads, until a
+ // partial store at this range updates the entry below.
+ auto Mask = llvm::to_vector<8>(
+ llvm::seq<int>(getIndex(Range.first), getIndex(Range.second)));
+ SliceValues[Range] =
+ IRB.CreateShuffleVector(InitVec, Mask, "init.extract");
}
// The init store itself becomes dead — its value is consumed via the
// extracts above.
@@ -3276,12 +3282,10 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
// alloca's final vector value. Insertion point is just after the
// block-order-last partial access — guaranteed to dominate every
// SliceValues entry since each one is either an init extract or a
- // stored value defined before its (now-deleted) store.
- Instruction *FinalInsert = Accesses.back().Inst->getNextNode();
- if (!FinalInsert)
- return std::nullopt;
-
- IRBuilder<> Builder(FinalInsert);
+ // stored value defined before its (now-deleted) store. The next
+ // node is guaranteed non-null because every basic block ends in a
+ // terminator and a load/store can never be the terminator.
+ IRBuilder<> Builder(Accesses.back().Inst->getNextNode());
std::queue<Value *> Q;
for (auto &Range : Partition)
Q.push(SliceValues[Range]);
diff --git a/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll b/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
index c569dac8353ca..ab991391bd85d 100644
--- a/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
+++ b/llvm/test/Transforms/SROA/vector-promotion-rmw-tree-merge.ll
@@ -12,10 +12,11 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f3
; - a single full-width load at the end.
;
; The rewrite should:
-; - replace each partial load with an extractvector from the current
-; per-slice value,
-; - thread each slice's SSA value through the rounds via the user's
-; compute (no selects/blends),
+; - initialize each slice's SSA value from the init store via a
+; shufflevector picking out that slice's elements,
+; - replace each partial load with the current per-slice SSA value
+; (RAUW), threading it through the rounds via the user's compute
+; (no selects/blends),
; - tree-merge the final per-slice values into the final alloca value.
;
@@ -24,10 +25,10 @@ define <4 x float> @single_round(<4 x float> %init, <2 x float> %a, <2 x float>
; CHECK-LABEL: define <4 x float> @single_round(
; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A:%.*]], <2 x float> [[B:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
-; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
-; CHECK-NEXT: [[R1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT1]], [[B]]
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <2 x float> [[INIT_EXTRACT1]], [[B]]
; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[R0]], <2 x float> [[R1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
; CHECK-NEXT: ret <4 x float> [[TMP0]]
;
@@ -56,10 +57,10 @@ define <4 x float> @multi_round(<4 x float> %init,
; CHECK-LABEL: define <4 x float> @multi_round(
; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A1:%.*]], <2 x float> [[B1:%.*]], <2 x float> [[A2:%.*]], <2 x float> [[B2:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
-; CHECK-NEXT: [[R0_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A1]]
-; CHECK-NEXT: [[R1_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT1]], [[B1]]
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[R0_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT]], [[A1]]
+; CHECK-NEXT: [[R1_1:%.*]] = fadd <2 x float> [[INIT_EXTRACT1]], [[B1]]
; CHECK-NEXT: [[R0_2:%.*]] = fadd <2 x float> [[R0_1]], [[A2]]
; CHECK-NEXT: [[R1_2:%.*]] = fadd <2 x float> [[R1_1]], [[B2]]
; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[R0_2]], <2 x float> [[R1_2]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
@@ -100,14 +101,13 @@ define <4 x float> @uneven_slices(<4 x float> %init, <3 x float> %a, <1 x float>
; CHECK-LABEL: define <4 x float> @uneven_slices(
; CHECK-SAME: <4 x float> [[INIT:%.*]], <3 x float> [[A:%.*]], <1 x float> [[B:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <3 x i32> <i32 0, i32 1, i32 2>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = extractelement <4 x float> [[INIT]], i32 3
-; CHECK-NEXT: [[TMP0:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT1]] to <1 x float>
-; CHECK-NEXT: [[R0:%.*]] = fadd <3 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
-; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[TMP0]], [[B]]
-; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <1 x float> [[R1]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
-; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <3 x float> [[R0]], <3 x float> [[TMP1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
-; CHECK-NEXT: ret <4 x float> [[TMP2]]
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <3 x i32> <i32 0, i32 1, i32 2>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <1 x i32> <i32 3>
+; CHECK-NEXT: [[R0:%.*]] = fadd <3 x float> [[INIT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[INIT_EXTRACT1]], [[B]]
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <1 x float> [[R1]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <3 x float> [[R0]], <3 x float> [[TMP0]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: ret <4 x float> [[TMP1]]
;
entry:
%alloca = alloca [4 x float], align 16
@@ -136,9 +136,9 @@ define <4 x float> @slice_swap(<4 x float> %init) {
; CHECK-LABEL: define <4 x float> @slice_swap(
; CHECK-SAME: <4 x float> [[INIT:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
-; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[INIT_EXTRACT_EXTRACT1]], <2 x float> [[INIT_EXTRACT_EXTRACT]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[INIT_EXTRACT1]], <2 x float> [[INIT_EXTRACT]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
; CHECK-NEXT: ret <4 x float> [[TMP0]]
;
entry:
@@ -160,24 +160,22 @@ entry:
}
; Load-only single-element slice. The init seed for a 1-element range is
-; produced by extractelement (a scalar), and a slice with no partial store
-; never updates that entry.
+; a 1-element shufflevector, and a slice with no partial store never
+; updates that entry.
define <4 x float> @load_only_single_element_slice(<4 x float> %init, <2 x float> %a, <1 x float> %b) {
; CHECK-LABEL: define <4 x float> @load_only_single_element_slice(
; CHECK-SAME: <4 x float> [[INIT:%.*]], <2 x float> [[A:%.*]], <1 x float> [[B:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT1:%.*]] = extractelement <4 x float> [[INIT]], i32 2
-; CHECK-NEXT: [[INIT_EXTRACT_VEC:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT1]] to <1 x float>
-; CHECK-NEXT: [[INIT_EXTRACT_EXTRACT2:%.*]] = extractelement <4 x float> [[INIT]], i32 3
-; CHECK-NEXT: [[INIT_EXTRACT_VEC3:%.*]] = bitcast float [[INIT_EXTRACT_EXTRACT2]] to <1 x float>
-; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT_EXTRACT]], [[A]]
-; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[INIT_EXTRACT_VEC]], [[B]]
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <1 x i32> <i32 2>
+; CHECK-NEXT: [[INIT_EXTRACT2:%.*]] = shufflevector <4 x float> [[INIT]], <4 x float> poison, <1 x i32> <i32 3>
+; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT]], [[A]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <1 x float> [[INIT_EXTRACT1]], [[B]]
; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <1 x float> [[R1]], <1 x float> poison, <2 x i32> <i32 0, i32 poison>
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x float> [[R0]], <2 x float> [[TMP0]], <3 x i32> <i32 0, i32 1, i32 2>
-; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <1 x float> [[INIT_EXTRACT_VEC3]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
+; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <1 x float> [[INIT_EXTRACT2]], <1 x float> poison, <3 x i32> <i32 0, i32 poison, i32 poison>
; CHECK-NEXT: [[TMP3:%.*]] = shufflevector <3 x float> [[TMP1]], <3 x float> [[TMP2]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
-; CHECK-NEXT: [[V2_FROM_INIT:%.*]] = bitcast <1 x float> [[INIT_EXTRACT_VEC3]] to float
+; CHECK-NEXT: [[V2_FROM_INIT:%.*]] = bitcast <1 x float> [[INIT_EXTRACT2]] to float
; CHECK-NEXT: [[FINAL:%.*]] = insertelement <4 x float> [[TMP3]], float [[V2_FROM_INIT]], i32 0
; CHECK-NEXT: ret <4 x float> [[FINAL]]
;
@@ -208,3 +206,86 @@ entry:
%final = insertelement <4 x float> %result, float %v2_from_init, i32 0
ret <4 x float> %final
}
+
+; Larger merge tree: 8 disjoint <2 x float> slices in a <16 x float> alloca.
+; The tree-merge has 8 leaves, so the final shufflevector tree has height 3.
+define <16 x float> @larger_merge_tree(<16 x float> %init,
+; CHECK-LABEL: define <16 x float> @larger_merge_tree(
+; CHECK-SAME: <16 x float> [[INIT:%.*]], <2 x float> [[A0:%.*]], <2 x float> [[A1:%.*]], <2 x float> [[A2:%.*]], <2 x float> [[A3:%.*]], <2 x float> [[A4:%.*]], <2 x float> [[A5:%.*]], <2 x float> [[A6:%.*]], <2 x float> [[A7:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[INIT_EXTRACT:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT: [[INIT_EXTRACT1:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT: [[INIT_EXTRACT2:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 4, i32 5>
+; CHECK-NEXT: [[INIT_EXTRACT3:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 6, i32 7>
+; CHECK-NEXT: [[INIT_EXTRACT4:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 8, i32 9>
+; CHECK-NEXT: [[INIT_EXTRACT5:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 10, i32 11>
+; CHECK-NEXT: [[INIT_EXTRACT6:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 12, i32 13>
+; CHECK-NEXT: [[INIT_EXTRACT7:%.*]] = shufflevector <16 x float> [[INIT]], <16 x float> poison, <2 x i32> <i32 14, i32 15>
+; CHECK-NEXT: [[R0:%.*]] = fadd <2 x float> [[INIT_EXTRACT]], [[A0]]
+; CHECK-NEXT: [[R1:%.*]] = fadd <2 x float> [[INIT_EXTRACT1]], [[A1]]
+; CHECK-NEXT: [[R2:%.*]] = fadd <2 x float> [[INIT_EXTRACT2]], [[A2]]
+; CHECK-NEXT: [[R3:%.*]] = fadd <2 x float> [[INIT_EXTRACT3]], [[A3]]
+; CHECK-NEXT: [[R4:%.*]] = fadd <2 x float> [[INIT_EXTRACT4]], [[A4]]
+; CHECK-NEXT: [[R5:%.*]] = fadd <2 x float> [[INIT_EXTRACT5]], [[A5]]
+; CHECK-NEXT: [[R6:%.*]] = fadd <2 x float> [[INIT_EXTRACT6]], [[A6]]
+; CHECK-NEXT: [[R7:%.*]] = fadd <2 x float> [[INIT_EXTRACT7]], [[A7]]
+; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x float> [[R0]], <2 x float> [[R1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x float> [[R2]], <2 x float> [[R3]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <2 x float> [[R4]], <2 x float> [[R5]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[TMP3:%.*]] = shufflevector <2 x float> [[R6]], <2 x float> [[R7]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT: [[TMP4:%.*]] = shufflevector <4 x float> [[TMP0]], <4 x float> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT: [[TMP5:%.*]] = shufflevector <4 x float> [[TMP2]], <4 x float> [[TMP3]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT: [[TMP6:%.*]] = shufflevector <8 x float> [[TMP4]], <8 x float> [[TMP5]], <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT: ret <16 x float> [[TMP6]]
+;
+ <2 x float> %a0, <2 x float> %a1,
+ <2 x float> %a2, <2 x float> %a3,
+ <2 x float> %a4, <2 x float> %a5,
+ <2 x float> %a6, <2 x float> %a7) {
+entry:
+ %alloca = alloca [16 x float], align 64
+ store <16 x float> %init, ptr %alloca, align 64
+
+ %p0 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 0
+ %v0 = load <2 x float>, ptr %p0, align 8
+ %r0 = fadd <2 x float> %v0, %a0
+ store <2 x float> %r0, ptr %p0, align 8
+
+ %p1 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 2
+ %v1 = load <2 x float>, ptr %p1, align 8
+ %r1 = fadd <2 x float> %v1, %a1
+ store <2 x float> %r1, ptr %p1, align 8
+
+ %p2 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 4
+ %v2 = load <2 x float>, ptr %p2, align 8
+ %r2 = fadd <2 x float> %v2, %a2
+ store <2 x float> %r2, ptr %p2, align 8
+
+ %p3 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 6
+ %v3 = load <2 x float>, ptr %p3, align 8
+ %r3 = fadd <2 x float> %v3, %a3
+ store <2 x float> %r3, ptr %p3, align 8
+
+ %p4 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 8
+ %v4 = load <2 x float>, ptr %p4, align 8
+ %r4 = fadd <2 x float> %v4, %a4
+ store <2 x float> %r4, ptr %p4, align 8
+
+ %p5 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 10
+ %v5 = load <2 x float>, ptr %p5, align 8
+ %r5 = fadd <2 x float> %v5, %a5
+ store <2 x float> %r5, ptr %p5, align 8
+
+ %p6 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 12
+ %v6 = load <2 x float>, ptr %p6, align 8
+ %r6 = fadd <2 x float> %v6, %a6
+ store <2 x float> %r6, ptr %p6, align 8
+
+ %p7 = getelementptr inbounds [16 x float], ptr %alloca, i32 0, i32 14
+ %v7 = load <2 x float>, ptr %p7, align 8
+ %r7 = fadd <2 x float> %v7, %a7
+ store <2 x float> %r7, ptr %p7, align 8
+
+ %result = load <16 x float>, ptr %alloca, align 64
+ ret <16 x float> %result
+}
More information about the llvm-commits
mailing list