[llvm] b26c514 - [SandboxVectorizer] Add container class to track and manage SeedBundles (#112048)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Oct 14 16:36:08 PDT 2024
Author: Sterling-Augustine
Date: 2024-10-14T16:36:05-07:00
New Revision: b26c514b2c7e089fa6f31c9433841521f34ae37f
URL: https://github.com/llvm/llvm-project/commit/b26c514b2c7e089fa6f31c9433841521f34ae37f
DIFF: https://github.com/llvm/llvm-project/commit/b26c514b2c7e089fa6f31c9433841521f34ae37f.diff
LOG: [SandboxVectorizer] Add container class to track and manage SeedBundles (#112048)
Added:
Modified:
llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index 619c2147f2e5c4..626fcb5afcb90b 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
@@ -54,6 +54,8 @@ class SeedBundle {
NumUnusedBits += Utils::getNumBits(I);
}
+ virtual void insert(Instruction *I, ScalarEvolution &SE) = 0;
+
unsigned getFirstUnusedElementIdx() const {
for (unsigned ElmIdx : seq<unsigned>(0, Seeds.size()))
if (!isUsed(ElmIdx))
@@ -96,6 +98,9 @@ class SeedBundle {
MutableArrayRef<Instruction *>
getSlice(unsigned StartIdx, unsigned MaxVecRegBits, bool ForcePowOf2);
+ /// \Returns the number of seed elements in the bundle.
+ std::size_t size() const { return Seeds.size(); }
+
protected:
SmallVector<Instruction *> Seeds;
/// The lanes that we have already vectorized.
@@ -148,7 +153,7 @@ template <typename LoadOrStoreT> class MemSeedBundle : public SeedBundle {
"Expected LoadInst or StoreInst!");
assert(isa<LoadOrStoreT>(MemI) && "Expected Load or Store!");
}
- void insert(sandboxir::Instruction *I, ScalarEvolution &SE) {
+ void insert(sandboxir::Instruction *I, ScalarEvolution &SE) override {
assert(isa<LoadOrStoreT>(I) && "Expected a Store or a Load!");
auto Cmp = [&SE](Instruction *I0, Instruction *I1) {
return Utils::atLowerAddress(cast<LoadOrStoreT>(I0),
@@ -162,5 +167,123 @@ template <typename LoadOrStoreT> class MemSeedBundle : public SeedBundle {
using StoreSeedBundle = MemSeedBundle<sandboxir::StoreInst>;
using LoadSeedBundle = MemSeedBundle<sandboxir::LoadInst>;
+/// Class to conveniently track Seeds within SeedBundles. Saves newly collected
+/// seeds in the proper bundle. Supports constant-time removal, as seeds and
+/// entire bundles are vectorized and marked used to signify removal. Iterators
+/// skip bundles that are completely used.
+class SeedContainer {
+ // Use the same key for
diff erent seeds if they are the same type and
+ // reference the same pointer, even if at
diff erent offsets. This directs
+ // potentially vectorizable seeds into the same bundle.
+ using KeyT = std::tuple<Value *, Type *, Instruction::Opcode>;
+ // Trying to vectorize too many seeds at once is expensive in
+ // compilation-time. Use a vector of bundles (all with the same key) to
+ // partition the candidate set into more manageable units. Each bundle is
+ // size-limited by sbvec-seed-bundle-size-limit. TODO: There might be a
+ // better way to divide these than by simple insertion order.
+ using ValT = SmallVector<std::unique_ptr<SeedBundle>>;
+ using BundleMapT = MapVector<KeyT, ValT>;
+ // Map from {pointer, Type, Opcode} to a vector of bundles.
+ BundleMapT Bundles;
+ // Allows finding a particular Instruction's bundle.
+ DenseMap<Instruction *, SeedBundle *> SeedLookupMap;
+
+ ScalarEvolution &SE;
+
+ template <typename LoadOrStoreT> KeyT getKey(LoadOrStoreT *LSI) const;
+
+public:
+ SeedContainer(ScalarEvolution &SE) : SE(SE) {}
+
+ class iterator {
+ BundleMapT *Map = nullptr;
+ BundleMapT::iterator MapIt;
+ ValT *Vec = nullptr;
+ size_t VecIdx;
+
+ public:
+ using
diff erence_type = std::ptr
diff _t;
+ using value_type = SeedBundle;
+ using pointer = value_type *;
+ using reference = value_type &;
+ using iterator_category = std::input_iterator_tag;
+
+ /// Iterates over the \p Map of SeedBundle Vectors, starting at \p MapIt,
+ /// and \p Vec at \p VecIdx, skipping vectors that are completely
+ /// used. Iteration order over the keys {Pointer, Type, Opcode} follows
+ /// DenseMap iteration order. For a given key, the vectors of
+ /// SeedBundles will be returned in insertion order. As in the
+ /// pseudo code below:
+ ///
+ /// for Key,Value in Bundles
+ /// for SeedBundleVector in Value
+ /// for SeedBundle in SeedBundleVector
+ /// if !SeedBundle.allUsed() ...
+ ///
+ /// Note that the bundles themselves may have additional ordering, created
+ /// by the subclasses by insertAt. The bundles themselves may also have used
+ /// instructions.
+ iterator(BundleMapT &Map, BundleMapT::iterator MapIt, ValT *Vec, int VecIdx)
+ : Map(&Map), MapIt(MapIt), Vec(Vec), VecIdx(VecIdx) {}
+ value_type &operator*() {
+ assert(Vec != nullptr && "Already at end!");
+ return *(*Vec)[VecIdx];
+ }
+ // Skip completely used bundles by repeatedly calling operator++().
+ void skipUsed() {
+ while (Vec && VecIdx < Vec->size() && this->operator*().allUsed())
+ ++(*this);
+ }
+ // Iterators iterate over the bundles
+ iterator &operator++() {
+ assert(VecIdx >= 0 && "Already at end!");
+ ++VecIdx;
+ if (VecIdx >= Vec->size()) {
+ assert(MapIt != Map->end() && "Already at end!");
+ VecIdx = 0;
+ ++MapIt;
+ if (MapIt != Map->end())
+ Vec = &MapIt->second;
+ else {
+ Vec = nullptr;
+ }
+ }
+ skipUsed();
+ return *this;
+ }
+ iterator operator++(int) {
+ auto Copy = *this;
+ ++(*this);
+ return Copy;
+ }
+ bool operator==(const iterator &Other) const {
+ assert(Map == Other.Map && "Iterator of
diff erent objects!");
+ return MapIt == Other.MapIt && VecIdx == Other.VecIdx;
+ }
+ bool operator!=(const iterator &Other) const { return !(*this == Other); }
+ };
+ using const_iterator = BundleMapT::const_iterator;
+ template <typename LoadOrStoreT> void insert(LoadOrStoreT *LSI);
+ // To support constant-time erase, these just mark the element used, rather
+ // than actually removing them from the bundle.
+ bool erase(Instruction *I);
+ bool erase(const KeyT &Key) { return Bundles.erase(Key); }
+ iterator begin() {
+ if (Bundles.empty())
+ return end();
+ auto BeginIt =
+ iterator(Bundles, Bundles.begin(), &Bundles.begin()->second, 0);
+ BeginIt.skipUsed();
+ return BeginIt;
+ }
+ iterator end() { return iterator(Bundles, Bundles.end(), nullptr, 0); }
+ unsigned size() const { return Bundles.size(); }
+
+#ifndef NDEBUG
+ LLVM_DUMP_METHOD void dump() const;
+#endif // NDEBUG
+};
+
} // namespace llvm::sandboxir
+
#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
index 00a7dc3fcec93e..20df9e344b61cd 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
@@ -19,6 +19,10 @@
using namespace llvm;
namespace llvm::sandboxir {
+cl::opt<unsigned> SeedBundleSizeLimit(
+ "sbvec-seed-bundle-size-limit", cl::init(32), cl::Hidden,
+ cl::desc("Limit the size of the seed bundle to cap compilation time."));
+
MutableArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
unsigned MaxVecRegBits,
bool ForcePowerOf2) {
@@ -61,4 +65,68 @@ MutableArrayRef<Instruction *> SeedBundle::getSlice(unsigned StartIdx,
return {};
}
+template <typename LoadOrStoreT>
+SeedContainer::KeyT SeedContainer::getKey(LoadOrStoreT *LSI) const {
+ assert((isa<LoadInst>(LSI) || isa<StoreInst>(LSI)) &&
+ "Expected Load or Store!");
+ Value *Ptr = Utils::getMemInstructionBase(LSI);
+ Instruction::Opcode Op = LSI->getOpcode();
+ Type *Ty = Utils::getExpectedType(LSI);
+ if (auto *VTy = dyn_cast<VectorType>(Ty))
+ Ty = VTy->getElementType();
+ return {Ptr, Ty, Op};
+}
+
+// Explicit instantiations
+template SeedContainer::KeyT
+SeedContainer::getKey<LoadInst>(LoadInst *LSI) const;
+template SeedContainer::KeyT
+SeedContainer::getKey<StoreInst>(StoreInst *LSI) const;
+
+bool SeedContainer::erase(Instruction *I) {
+ assert((isa<LoadInst>(I) || isa<StoreInst>(I)) && "Expected Load or Store!");
+ auto It = SeedLookupMap.find(I);
+ if (It == SeedLookupMap.end())
+ return false;
+ SeedBundle *Bndl = It->second;
+ Bndl->setUsed(I);
+ return true;
+}
+
+template <typename LoadOrStoreT> void SeedContainer::insert(LoadOrStoreT *LSI) {
+ // Find the bundle containing seeds for this symbol and type-of-access.
+ auto &BundleVec = Bundles[getKey(LSI)];
+ // Fill this vector of bundles front to back so that only the last bundle in
+ // the vector may have available space. This avoids iteration to find one with
+ // space.
+ if (BundleVec.empty() || BundleVec.back()->size() == SeedBundleSizeLimit)
+ BundleVec.emplace_back(std::make_unique<MemSeedBundle<LoadOrStoreT>>(LSI));
+ else
+ BundleVec.back()->insert(LSI, SE);
+
+ SeedLookupMap[LSI] = BundleVec.back().get();
+}
+
+// Explicit instantiations
+template void SeedContainer::insert<LoadInst>(LoadInst *);
+template void SeedContainer::insert<StoreInst>(StoreInst *);
+
+#ifndef NDEBUG
+void SeedContainer::dump() const {
+ for (const auto &Pair : Bundles) {
+ auto [I, Ty, Opc] = Pair.first;
+ const auto &SeedsVec = Pair.second;
+ std::string RefType = dyn_cast<LoadInst>(I) ? "Load"
+ : dyn_cast<StoreInst>(I) ? "Store"
+ : "Other";
+ dbgs() << "[Inst=" << *I << " Ty=" << Ty << " " << RefType << "]\n";
+ for (const auto &SeedPtr : SeedsVec) {
+ SeedPtr->dump(dbgs());
+ dbgs() << "\n";
+ }
+ }
+ dbgs() << "\n";
+}
+#endif // NDEBUG
+
} // namespace llvm::sandboxir
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
index dd41b0a6605095..82b230d50c4ec9 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
@@ -40,6 +40,15 @@ struct SeedBundleTest : public testing::Test {
}
};
+// Stub class to make the abstract base class testable.
+class SeedBundleForTest : public sandboxir::SeedBundle {
+public:
+ using sandboxir::SeedBundle::SeedBundle;
+ void insert(sandboxir::Instruction *I, ScalarEvolution &SE) override {
+ insertAt(Seeds.end(), I);
+ }
+};
+
TEST_F(SeedBundleTest, SeedBundle) {
parseIR(C, R"IR(
define void @foo(float %v0, i32 %i0, i16 %i1, i8 %i2) {
@@ -66,7 +75,7 @@ define void @foo(float %v0, i32 %i0, i16 %i1, i8 %i2) {
// Assume first two instructions are identical in the number of bits.
const unsigned IOBits = sandboxir::Utils::getNumBits(I0, DL);
// Constructor
- sandboxir::SeedBundle SBO(I0);
+ SeedBundleForTest SBO(I0);
EXPECT_EQ(*SBO.begin(), I0);
// getNumUnusedBits after constructor
EXPECT_EQ(SBO.getNumUnusedBits(), IOBits);
@@ -103,7 +112,7 @@ define void @foo(float %v0, i32 %i0, i16 %i1, i8 %i2) {
EXPECT_EQ(BundleBits, 88u);
auto Seeds = Insts;
// Constructor
- sandboxir::SeedBundle SB1(std::move(Seeds));
+ SeedBundleForTest SB1(std::move(Seeds));
// getNumUnusedBits after constructor
EXPECT_EQ(SB1.getNumUnusedBits(), BundleBits);
// setUsed with index
@@ -196,3 +205,66 @@ define void @foo(ptr %ptrA, float %val, ptr %ptr) {
sandboxir::LoadSeedBundle LB(std::move(Loads), SE);
EXPECT_THAT(LB, testing::ElementsAre(L0, L1, L2, L3));
}
+
+TEST_F(SeedBundleTest, Container) {
+ parseIR(C, R"IR(
+define void @foo(ptr %ptrA, float %val, ptr %ptrB) {
+bb:
+ %gepA0 = getelementptr float, ptr %ptrA, i32 0
+ %gepA1 = getelementptr float, ptr %ptrA, i32 1
+ %gepB0 = getelementptr float, ptr %ptrB, i32 0
+ %gepB1 = getelementptr float, ptr %ptrB, i32 1
+ store float %val, ptr %gepA0
+ store float %val, ptr %gepA1
+ store float %val, ptr %gepB0
+ store float %val, ptr %gepB1
+ ret void
+}
+)IR");
+ Function &LLVMF = *M->getFunction("foo");
+
+ DominatorTree DT(LLVMF);
+ TargetLibraryInfoImpl TLII;
+ TargetLibraryInfo TLI(TLII);
+ DataLayout DL(M->getDataLayout());
+ LoopInfo LI(DT);
+ AssumptionCache AC(LLVMF);
+ ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
+
+ sandboxir::Context Ctx(C);
+ auto &F = *Ctx.createFunction(&LLVMF);
+ auto &BB = *F.begin();
+ auto It = std::next(BB.begin(), 4);
+ auto *S0 = cast<sandboxir::StoreInst>(&*It++);
+ auto *S1 = cast<sandboxir::StoreInst>(&*It++);
+ auto *S2 = cast<sandboxir::StoreInst>(&*It++);
+ auto *S3 = cast<sandboxir::StoreInst>(&*It++);
+ sandboxir::SeedContainer SC(SE);
+ // Check begin() end() when empty.
+ EXPECT_EQ(SC.begin(), SC.end());
+
+ SC.insert(S0);
+ SC.insert(S1);
+ SC.insert(S2);
+ SC.insert(S3);
+ unsigned Cnt = 0;
+ SmallVector<sandboxir::SeedBundle *> Bndls;
+ for (auto &SeedBndl : SC) {
+ EXPECT_EQ(SeedBndl.size(), 2u);
+ ++Cnt;
+ Bndls.push_back(&SeedBndl);
+ }
+ EXPECT_EQ(Cnt, 2u);
+
+ // Mark them "Used" to check if operator++ skips them in the next loop.
+ for (auto *SeedBndl : Bndls)
+ for (auto Lane : seq<unsigned>(SeedBndl->size()))
+ SeedBndl->setUsed(Lane);
+ // Check if iterator::operator++ skips used lanes.
+ Cnt = 0;
+ for (auto &SeedBndl : SC) {
+ (void)SeedBndl;
+ ++Cnt;
+ }
+ EXPECT_EQ(Cnt, 0u);
+}
More information about the llvm-commits
mailing list