[llvm] [SLP] More OOP to simplify vectorizeStores() (NFC) (PR #134605)
Gaëtan Bossu via llvm-commits
llvm-commits at lists.llvm.org
Wed Apr 16 09:19:26 PDT 2025
https://github.com/gbossu updated https://github.com/llvm/llvm-project/pull/134605
>From ef65d73d1bd6ba36cee89fa597931b2a4d5a821b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Bossu?= <gaetan.bossu at arm.com>
Date: Fri, 4 Apr 2025 08:39:00 +0000
Subject: [PATCH 1/4] [SLP] More OOP to simplify vectorizeStores() (NFC)
This moves more code into the RelatedStoreInsts helper class. The
FillStoresSet lambda is now only a couple of lines and is easier to
read.
---
.../Transforms/Vectorize/SLPVectorizer.cpp | 137 ++++++++++--------
1 file changed, 80 insertions(+), 57 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index f97386159d029..6b85b832daec3 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20856,9 +20856,15 @@ namespace {
/// A group of stores that we'll try to bundle together using vector ops.
/// They are ordered using the signed distance of their address operand to the
/// address of this group's BaseInstr.
-struct RelatedStoreInsts {
- RelatedStoreInsts(unsigned BaseInstrIdx) { reset(BaseInstrIdx); }
+class RelatedStoreInsts {
+public:
+ RelatedStoreInsts(unsigned BaseInstrIdx, ArrayRef<StoreInst *> AllStores)
+ : AllStores(AllStores) {
+ reset(BaseInstrIdx);
+ }
+
void reset(unsigned NewBaseInstr) {
+ assert(NewBaseInstr < AllStores.size());
BaseInstrIdx = NewBaseInstr;
Instrs.clear();
insertOrLookup(NewBaseInstr, 0);
@@ -20873,12 +20879,54 @@ struct RelatedStoreInsts {
return Inserted ? std::nullopt : std::optional<unsigned>(It->second);
}
+ StoreInst &getBaseStore() const { return *AllStores[BaseInstrIdx]; }
+ using DistToInstMap = std::map<int, unsigned>;
+ const DistToInstMap &getStores() const { return Instrs; }
+
+ /// Recompute the pointer distances to be based on \p NewBaseInstIdx.
+ /// Stores whose index is less than \p MinSafeIdx will be dropped.
+ void rebase(unsigned MinSafeIdx, unsigned NewBaseInstIdx,
+ int DistFromCurBase) {
+ DistToInstMap PrevSet = std::move(Instrs);
+ reset(NewBaseInstIdx);
+
+ // Re-insert stores that come after MinSafeIdx to try and vectorize them
+ // again. Their distance will be "rebased" to use NewBaseInstIdx as
+ // reference.
+ for (auto [Dist, InstIdx] : PrevSet) {
+ if (InstIdx >= MinSafeIdx) {
+ insertOrLookup(InstIdx, Dist - DistFromCurBase);
+ }
+ }
+ }
+
+ /// Remove all stores that have been vectorized from this group.
+ void clearVectorizedStores(const BoUpSLP::ValueSet &VectorizedStores) {
+ const auto Begin = Instrs.begin();
+ auto NonVectorizedStore = Instrs.end();
+
+ while (NonVectorizedStore != Begin) {
+ const auto Prev = std::prev(NonVectorizedStore);
+ unsigned InstrIdx = Prev->second;
+ if (VectorizedStores.contains(AllStores[InstrIdx])) {
+ // NonVectorizedStore is the last scalar instruction.
+ // Erase all stores before it so we don't try to vectorize them again.
+ Instrs.erase(Begin, NonVectorizedStore);
+ return;
+ }
+ NonVectorizedStore = Prev;
+ }
+ }
+
+private:
/// The index of the Base instruction, i.e. the one with a 0 pointer distance.
unsigned BaseInstrIdx;
/// Maps a pointer distance from \p BaseInstrIdx to an instruction index.
- using DistToInstMap = std::map<int, unsigned>;
DistToInstMap Instrs;
+
+ /// Reference to all the stores in the BB being analyzed.
+ ArrayRef<StoreInst *> AllStores;
};
} // end anonymous namespace
@@ -21166,14 +21214,7 @@ bool SLPVectorizerPass::vectorizeStores(
}
};
- // Stores pair (first: index of the store into Stores array ref, address of
- // which taken as base, second: sorted set of pairs {index, dist}, which are
- // indices of stores in the set and their store location distances relative to
- // the base address).
-
- // Need to store the index of the very first store separately, since the set
- // may be reordered after the insertion and the first store may be moved. This
- // container allows to reduce number of calls of getPointersDiff() function.
+ /// Groups of stores to vectorize
SmallVector<RelatedStoreInsts> SortedStores;
// Inserts the specified store SI with the given index Idx to the set of the
@@ -21209,52 +21250,34 @@ bool SLPVectorizerPass::vectorizeStores(
// dependencies and no need to waste compile time to try to vectorize them.
// - Try to vectorize the sequence {1, {1, 0}, {3, 2}}.
auto FillStoresSet = [&](unsigned Idx, StoreInst *SI) {
- for (RelatedStoreInsts &StoreSeq : SortedStores) {
- std::optional<int> Diff = getPointersDiff(
- Stores[StoreSeq.BaseInstrIdx]->getValueOperand()->getType(),
- Stores[StoreSeq.BaseInstrIdx]->getPointerOperand(),
- SI->getValueOperand()->getType(), SI->getPointerOperand(), *DL, *SE,
- /*StrictCheck=*/true);
- if (!Diff)
- continue;
- std::optional<unsigned> PrevInst =
- StoreSeq.insertOrLookup(/*InstrIdx=*/Idx, /*PtrDist=*/*Diff);
- if (!PrevInst) {
- // No store was associated to that distance. Keep collecting.
- return;
- }
- // Try to vectorize the first found set to avoid duplicate analysis.
- TryToVectorize(StoreSeq.Instrs);
- RelatedStoreInsts::DistToInstMap PrevSet;
- copy_if(StoreSeq.Instrs, std::inserter(PrevSet, PrevSet.end()),
- [&](const std::pair<int, unsigned> &DistAndIdx) {
- return DistAndIdx.second > *PrevInst;
- });
- StoreSeq.reset(Idx);
- // Insert stores that followed previous match to try to vectorize them
- // with this store.
- unsigned StartIdx = *PrevInst + 1;
- SmallBitVector UsedStores(Idx - StartIdx);
- // Distances to previously found dup store (or this store, since they
- // store to the same addresses).
- SmallVector<int> Dists(Idx - StartIdx, 0);
- for (auto [PtrDist, InstIdx] : reverse(PrevSet)) {
- // Do not try to vectorize sequences, we already tried.
- if (VectorizedStores.contains(Stores[InstIdx]))
- break;
- unsigned BI = InstIdx - StartIdx;
- UsedStores.set(BI);
- Dists[BI] = PtrDist - *Diff;
- }
- for (unsigned I = StartIdx; I < Idx; ++I) {
- unsigned BI = I - StartIdx;
- if (UsedStores.test(BI))
- StoreSeq.insertOrLookup(I, Dists[BI]);
- }
+ std::optional<int> Diff;
+ auto *RelatedStores =
+ find_if(SortedStores, [&](const RelatedStoreInsts &StoreSeq) {
+ StoreInst &BaseStore = StoreSeq.getBaseStore();
+ Diff = getPointersDiff(BaseStore.getValueOperand()->getType(),
+ BaseStore.getPointerOperand(),
+ SI->getValueOperand()->getType(),
+ SI->getPointerOperand(), *DL, *SE,
+ /*StrictCheck=*/true);
+ return Diff.has_value();
+ });
+
+ // We did not find a comparable store, start a new group.
+ if (RelatedStores == SortedStores.end()) {
+ SortedStores.emplace_back(Idx, Stores);
return;
}
- // We did not find a comparable store, start a new sequence.
- SortedStores.emplace_back(Idx);
+
+ // If there is already a store in the group with the same PtrDiff, try to
+ // vectorize the existing instructions before adding the current store.
+ if (std::optional<unsigned> PrevInst =
+ RelatedStores->insertOrLookup(Idx, *Diff)) {
+ TryToVectorize(RelatedStores->getStores());
+ RelatedStores->clearVectorizedStores(VectorizedStores);
+ RelatedStores->rebase(/*MinSafeIdx=*/*PrevInst + 1,
+ /*NewBaseInstIdx=*/Idx,
+ /*DistFromCurBase=*/*Diff);
+ }
};
Type *PrevValTy = nullptr;
for (auto [I, SI] : enumerate(Stores)) {
@@ -21265,7 +21288,7 @@ bool SLPVectorizerPass::vectorizeStores(
// Check that we do not try to vectorize stores of different types.
if (PrevValTy != SI->getValueOperand()->getType()) {
for (RelatedStoreInsts &StoreSeq : SortedStores)
- TryToVectorize(StoreSeq.Instrs);
+ TryToVectorize(StoreSeq.getStores());
SortedStores.clear();
PrevValTy = SI->getValueOperand()->getType();
}
@@ -21274,7 +21297,7 @@ bool SLPVectorizerPass::vectorizeStores(
// Final vectorization attempt.
for (RelatedStoreInsts &StoreSeq : SortedStores)
- TryToVectorize(StoreSeq.Instrs);
+ TryToVectorize(StoreSeq.getStores());
return Changed;
}
>From c1850923e5c58797598d6e196d03ecf582771897 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Bossu?= <gaetan.bossu at arm.com>
Date: Fri, 11 Apr 2025 14:30:26 +0000
Subject: [PATCH 2/4] fixup! [SLP] More OOP to simplify vectorizeStores() (NFC)
---
llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 6b85b832daec3..fc798b2e3f3bf 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20864,7 +20864,8 @@ class RelatedStoreInsts {
}
void reset(unsigned NewBaseInstr) {
- assert(NewBaseInstr < AllStores.size());
+ assert(NewBaseInstr < AllStores.size() &&
+ "Instruction index out of bounds");
BaseInstrIdx = NewBaseInstr;
Instrs.clear();
insertOrLookup(NewBaseInstr, 0);
@@ -20894,9 +20895,8 @@ class RelatedStoreInsts {
// again. Their distance will be "rebased" to use NewBaseInstIdx as
// reference.
for (auto [Dist, InstIdx] : PrevSet) {
- if (InstIdx >= MinSafeIdx) {
+ if (InstIdx >= MinSafeIdx)
insertOrLookup(InstIdx, Dist - DistFromCurBase);
- }
}
}
@@ -21270,6 +21270,7 @@ bool SLPVectorizerPass::vectorizeStores(
// If there is already a store in the group with the same PtrDiff, try to
// vectorize the existing instructions before adding the current store.
+ // Otherwise, insert this store and keep collecting.
if (std::optional<unsigned> PrevInst =
RelatedStores->insertOrLookup(Idx, *Diff)) {
TryToVectorize(RelatedStores->getStores());
>From f935439e8d249575e8f46381d215bc2cf8263b25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Bossu?= <gaetan.bossu at arm.com>
Date: Mon, 14 Apr 2025 09:38:17 +0000
Subject: [PATCH 3/4] fixup! [SLP] More OOP to simplify vectorizeStores() (NFC)
---
.../Transforms/Vectorize/SLPVectorizer.cpp | 23 ++++++++-----------
1 file changed, 9 insertions(+), 14 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index fc798b2e3f3bf..aabce015d2ec8 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20902,20 +20902,15 @@ class RelatedStoreInsts {
/// Remove all stores that have been vectorized from this group.
void clearVectorizedStores(const BoUpSLP::ValueSet &VectorizedStores) {
- const auto Begin = Instrs.begin();
- auto NonVectorizedStore = Instrs.end();
-
- while (NonVectorizedStore != Begin) {
- const auto Prev = std::prev(NonVectorizedStore);
- unsigned InstrIdx = Prev->second;
- if (VectorizedStores.contains(AllStores[InstrIdx])) {
- // NonVectorizedStore is the last scalar instruction.
- // Erase all stores before it so we don't try to vectorize them again.
- Instrs.erase(Begin, NonVectorizedStore);
- return;
- }
- NonVectorizedStore = Prev;
- }
+ DistToInstMap::reverse_iterator LastVectorizedStore = find_if(
+ reverse(Instrs), [&](const std::pair<int, unsigned> &DistAndIdx) {
+ return VectorizedStores.contains(AllStores[DistAndIdx.second]);
+ });
+
+ // Get a forward iterator pointing after the last vectorized store and erase
+ // all stores before it so we don't try to vectorize them again.
+ DistToInstMap::iterator VectorizedStoresEnd = LastVectorizedStore.base();
+ Instrs.erase(Instrs.begin(), VectorizedStoresEnd);
}
private:
>From d03b68bfdedae1503f02a859d37e1502c75a0b61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Bossu?= <gaetan.bossu at arm.com>
Date: Wed, 16 Apr 2025 15:38:44 +0000
Subject: [PATCH 4/4] fixup! [SLP] More OOP to simplify vectorizeStores() (NFC)
---
.../Transforms/Vectorize/SLPVectorizer.cpp | 31 +++++++++++--------
1 file changed, 18 insertions(+), 13 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index aabce015d2ec8..bb85f8be92163 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20880,10 +20880,20 @@ class RelatedStoreInsts {
return Inserted ? std::nullopt : std::optional<unsigned>(It->second);
}
- StoreInst &getBaseStore() const { return *AllStores[BaseInstrIdx]; }
using DistToInstMap = std::map<int, unsigned>;
const DistToInstMap &getStores() const { return Instrs; }
+ /// If \p SI is related to this group of stores, return the distance of its
+ /// pointer operand to the one the group's BaseInstr.
+ std::optional<int> getPointerDiff(StoreInst &SI, const DataLayout &DL,
+ ScalarEvolution &SE) const {
+ StoreInst &BaseStore = *AllStores[BaseInstrIdx];
+ return getPointersDiff(
+ BaseStore.getValueOperand()->getType(), BaseStore.getPointerOperand(),
+ SI.getValueOperand()->getType(), SI.getPointerOperand(), DL, SE,
+ /*StrictCheck=*/true);
+ }
+
/// Recompute the pointer distances to be based on \p NewBaseInstIdx.
/// Stores whose index is less than \p MinSafeIdx will be dropped.
void rebase(unsigned MinSafeIdx, unsigned NewBaseInstIdx,
@@ -21245,16 +21255,11 @@ bool SLPVectorizerPass::vectorizeStores(
// dependencies and no need to waste compile time to try to vectorize them.
// - Try to vectorize the sequence {1, {1, 0}, {3, 2}}.
auto FillStoresSet = [&](unsigned Idx, StoreInst *SI) {
- std::optional<int> Diff;
- auto *RelatedStores =
- find_if(SortedStores, [&](const RelatedStoreInsts &StoreSeq) {
- StoreInst &BaseStore = StoreSeq.getBaseStore();
- Diff = getPointersDiff(BaseStore.getValueOperand()->getType(),
- BaseStore.getPointerOperand(),
- SI->getValueOperand()->getType(),
- SI->getPointerOperand(), *DL, *SE,
- /*StrictCheck=*/true);
- return Diff.has_value();
+ std::optional<int> PtrDist;
+ auto *RelatedStores = find_if(
+ SortedStores, [&PtrDist, SI, this](const RelatedStoreInsts &StoreSeq) {
+ PtrDist = StoreSeq.getPointerDiff(*SI, *DL, *SE);
+ return PtrDist.has_value();
});
// We did not find a comparable store, start a new group.
@@ -21267,12 +21272,12 @@ bool SLPVectorizerPass::vectorizeStores(
// vectorize the existing instructions before adding the current store.
// Otherwise, insert this store and keep collecting.
if (std::optional<unsigned> PrevInst =
- RelatedStores->insertOrLookup(Idx, *Diff)) {
+ RelatedStores->insertOrLookup(Idx, *PtrDist)) {
TryToVectorize(RelatedStores->getStores());
RelatedStores->clearVectorizedStores(VectorizedStores);
RelatedStores->rebase(/*MinSafeIdx=*/*PrevInst + 1,
/*NewBaseInstIdx=*/Idx,
- /*DistFromCurBase=*/*Diff);
+ /*DistFromCurBase=*/*PtrDist);
}
};
Type *PrevValTy = nullptr;
More information about the llvm-commits
mailing list