From f35a46729c2d2f2c4caf82222eedfffc96621e23 Mon Sep 17 00:00:00 2001 From: Sterling Augustine Date: Tue, 12 Nov 2024 14:32:57 -0800 Subject: [PATCH 1/2] [SandboxVectorizer] Register erase instruction callbacks for seed collection --- .../Vectorize/SandboxVectorizer/SeedCollector.h | 7 ++++++- .../Vectorize/SandboxVectorizer/SeedCollector.cpp | 13 +++++++++---- .../SandboxVectorizer/SeedCollectorTest.cpp | 15 +++++++++++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h index ed1cb8488c29e..958ec673b6016 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h @@ -223,6 +223,11 @@ class SeedContainer { /// Note that the bundles themselves may have additional ordering, created /// by the subclasses by insertAt. The bundles themselves may also have used /// instructions. + + // TODO: Range_size counts fully used-bundles. Further, iterating over + // anything other than the Bundles in a SeedContainer includes used seeds, + // so for now just check that removing all the seeds from a bundle also + // empties the bundle. Rework the iterator logic to clean this up. iterator(BundleMapT &Map, BundleMapT::iterator MapIt, ValT *Vec, int VecIdx) : Map(&Map), MapIt(MapIt), Vec(Vec), VecIdx(VecIdx) {} value_type &operator*() { @@ -288,7 +293,7 @@ class SeedCollector { SeedContainer StoreSeeds; SeedContainer LoadSeeds; Context &Ctx; - + Context::CallbackID EraseCallbackID; /// \Returns the number of SeedBundle groups for all seed types. /// This is to be used for limiting compilation time. unsigned totalNumSeedGroups() const { diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp index 1dbdd80117563..17544afcc185f 100644 --- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp +++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp @@ -159,13 +159,19 @@ template bool isValidMemSeed(StoreInst *LSI); SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE) : StoreSeeds(SE), LoadSeeds(SE), Ctx(BB->getContext()) { - // TODO: Register a callback for updating the Collector data structures upon - // instr removal bool CollectStores = CollectSeeds.find(StoreSeedsDef) != std::string::npos; bool CollectLoads = CollectSeeds.find(LoadSeedsDef) != std::string::npos; if (!CollectStores && !CollectLoads) return; + + EraseCallbackID = Ctx.registerEraseInstrCallback([this](Instruction *I) { + if (auto SI = dyn_cast(I)) + StoreSeeds.erase(SI); + else if (auto LI = dyn_cast(I)) + LoadSeeds.erase(LI); + }); + // Actually collect the seeds. for (auto &I : *BB) { if (StoreInst *SI = dyn_cast(&I)) @@ -181,8 +187,7 @@ SeedCollector::SeedCollector(BasicBlock *BB, ScalarEvolution &SE) } SeedCollector::~SeedCollector() { - // TODO: Unregister the callback for updating the seed datastructures upon - // instr removal + Ctx.unregisterEraseInstrCallback(EraseCallbackID); } #ifndef NDEBUG diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp index b1b0156a4fe31..4fc44f9fddc9a 100644 --- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp @@ -374,6 +374,21 @@ define void @foo(ptr noalias %ptr, float %val) { // Expect just one vector of store seeds EXPECT_EQ(range_size(StoreSeedsRange), 1u); ExpectThatElementsAre(SB, {St0, St1, St2, St3}); + // Check that the EraseInstr callback works. TODO: Range_size counts fully + // used-bundles even though the iterator skips them. Further, iterating over + // anything other than the Bundles in a SeedContainer includes used seeds. So + // for now just check that removing all the seeds from a bundle also empties + // the bundle. + St0->eraseFromParent(); + St1->eraseFromParent(); + St2->eraseFromParent(); + St3->eraseFromParent(); + size_t nonEmptyBundleCount = 0; + for (auto &B : SC.getStoreSeeds()) { + (void)B; + nonEmptyBundleCount++; + } + EXPECT_EQ(nonEmptyBundleCount, 0u); } TEST_F(SeedBundleTest, VectorStores) { From ed034eaa9b42fc0a3723625a51402e095a39a56f Mon Sep 17 00:00:00 2001 From: Sterling Augustine Date: Tue, 12 Nov 2024 15:42:09 -0800 Subject: [PATCH 2/2] Address comments --- .../Vectorize/SandboxVectorizer/SeedCollector.h | 5 ++--- .../Vectorize/SandboxVectorizer/SeedCollectorTest.cpp | 11 ++++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h index 958ec673b6016..6e16a84d832e5 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h @@ -225,9 +225,8 @@ class SeedContainer { /// instructions. // TODO: Range_size counts fully used-bundles. Further, iterating over - // anything other than the Bundles in a SeedContainer includes used seeds, - // so for now just check that removing all the seeds from a bundle also - // empties the bundle. Rework the iterator logic to clean this up. + // anything other than the Bundles in a SeedContainer includes used + // seeds. Rework the iterator logic to clean this up. iterator(BundleMapT &Map, BundleMapT::iterator MapIt, ValT *Vec, int VecIdx) : Map(&Map), MapIt(MapIt), Vec(Vec), VecIdx(VecIdx) {} value_type &operator*() { diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp index 4fc44f9fddc9a..95a8f66ac1e66 100644 --- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp @@ -374,11 +374,12 @@ define void @foo(ptr noalias %ptr, float %val) { // Expect just one vector of store seeds EXPECT_EQ(range_size(StoreSeedsRange), 1u); ExpectThatElementsAre(SB, {St0, St1, St2, St3}); - // Check that the EraseInstr callback works. TODO: Range_size counts fully - // used-bundles even though the iterator skips them. Further, iterating over - // anything other than the Bundles in a SeedContainer includes used seeds. So - // for now just check that removing all the seeds from a bundle also empties - // the bundle. + // Check that the EraseInstr callback works. + + // TODO: Range_size counts fully used-bundles even though the iterator skips + // them. Further, iterating over anything other than the Bundles in a + // SeedContainer includes used seeds. So for now just check that removing all + // the seeds from a bundle also empties the bundle. St0->eraseFromParent(); St1->eraseFromParent(); St2->eraseFromParent();