[SandboxVectorizer] Register erase callback for seed collection #115951

Sterling-Augustine · 2024-11-12T22:35:42Z

No description provided.

…lection

llvmbot · 2024-11-12T22:36:16Z

@llvm/pr-subscribers-llvm-transforms

Author: None (Sterling-Augustine)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/115951.diff

3 Files Affected:

(modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h (+6-1)
(modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp (+9-4)
(modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp (+15)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index ed1cb8488c29eb..958ec673b60164 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 1dbdd80117563c..17544afcc185fd 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>(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<StoreInst>(I))
+      StoreSeeds.erase(SI);
+    else if (auto LI = dyn_cast<LoadInst>(I))
+      LoadSeeds.erase(LI);
+  });
+
   // Actually collect the seeds.
   for (auto &I : *BB) {
     if (StoreInst *SI = dyn_cast<StoreInst>(&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 b1b0156a4fe31a..4fc44f9fddc9a8 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) {

llvmbot · 2024-11-12T22:36:16Z

@llvm/pr-subscribers-vectorizers

Author: None (Sterling-Augustine)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/115951.diff

3 Files Affected:

(modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h (+6-1)
(modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp (+9-4)
(modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp (+15)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index ed1cb8488c29eb..958ec673b60164 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 1dbdd80117563c..17544afcc185fd 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>(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<StoreInst>(I))
+      StoreSeeds.erase(SI);
+    else if (auto LI = dyn_cast<LoadInst>(I))
+      LoadSeeds.erase(LI);
+  });
+
   // Actually collect the seeds.
   for (auto &I : *BB) {
     if (StoreInst *SI = dyn_cast<StoreInst>(&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 b1b0156a4fe31a..4fc44f9fddc9a8 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) {

vporpo

Just some minor comments.

vporpo · 2024-11-12T22:44:01Z

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


nit: TODO: in a separate line.
also double whitespace before "So".

vporpo · 2024-11-12T22:48:28Z

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h

+
+    // 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


The sentence "so for now just check that removing all the seeds from a bundle also empties the bundle" seems to refer to the test.

[SandboxVectorizer] Register erase instruction callbacks for seed col…

f35a467

…lection

Sterling-Augustine requested a review from vporpo November 12, 2024 22:35

llvmbot added vectorizers llvm:transforms labels Nov 12, 2024

Sterling-Augustine requested a review from slackito November 12, 2024 22:36

vporpo approved these changes Nov 12, 2024

View reviewed changes

Address comments

ed034ea

Sterling-Augustine merged commit 7ba864b into llvm:main Nov 13, 2024
5 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[SandboxVectorizer] Register erase callback for seed collection #115951

[SandboxVectorizer] Register erase callback for seed collection #115951

Uh oh!

Sterling-Augustine commented Nov 12, 2024

Uh oh!

llvmbot commented Nov 12, 2024

Uh oh!

llvmbot commented Nov 12, 2024

Uh oh!

vporpo left a comment

Uh oh!

vporpo Nov 12, 2024

Uh oh!

vporpo Nov 12, 2024

Uh oh!

Uh oh!

Uh oh!

[SandboxVectorizer] Register erase callback for seed collection #115951

[SandboxVectorizer] Register erase callback for seed collection #115951

Uh oh!

Conversation

Sterling-Augustine commented Nov 12, 2024

Uh oh!

llvmbot commented Nov 12, 2024

Uh oh!

llvmbot commented Nov 12, 2024

Uh oh!

vporpo left a comment

Choose a reason for hiding this comment

Uh oh!

vporpo Nov 12, 2024

Choose a reason for hiding this comment

Uh oh!

vporpo Nov 12, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!