[llvm] a1a1073 - [SandboxVec][DAG] Remove all matching nodes from bundle

[Vasileios Porpodas via llvm-commits]

Author: Vasileios Porpodas
Date: 2025-02-20T15:13:52-08:00
New Revision: a1a10732e417020269b1cafcbdaba7bc217c5299

URL: https://github.com/llvm/llvm-project/commit/a1a10732e417020269b1cafcbdaba7bc217c5299
DIFF: https://github.com/llvm/llvm-project/commit/a1a10732e417020269b1cafcbdaba7bc217c5299.diff

LOG: [SandboxVec][DAG] Remove all matching nodes from bundle

This fixes a bug where `SchedBundle::eraseFromBundle()` would not erase all
matching nodes but just the first one.

Added: 
    

Modified: 
    llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
    llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
index 6b56f348f328c..cdd4db36dafad 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -112,7 +112,9 @@ class SchedBundle {
   ContainerTy Nodes;
 
   /// Called by the DGNode destructor to avoid accessing freed memory.
-  void eraseFromBundle(DGNode *N) { Nodes.erase(find(Nodes, N)); }
+  void eraseFromBundle(DGNode *N) {
+    Nodes.erase(std::remove(Nodes.begin(), Nodes.end(), N), Nodes.end());
+  }
   friend DGNode::~DGNode(); // For eraseFromBundle().
 
 public:

diff  --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
index 97724100ba341..c0d8178f41480 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
@@ -145,6 +145,47 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
               testing::ElementsAre(SN0, SN1));
 }
 
+// Check that when we erase a DAG node its SchedBundle gets updated.
+TEST_F(SchedulerTest, SchedBundleEraseDGNode) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr, i8 %v0, i8 %v1, i8 %v2, i8 %v3) {
+  store i8 %v0, ptr %ptr
+  store i8 %v1, ptr %ptr
+  store i8 %v2, ptr %ptr
+  store i8 %v3, ptr %ptr
+  ret void
+}
+)IR");
+  llvm::Function *LLVMF = &*M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(LLVMF);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  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::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({&*BB->begin(), BB->getTerminator()});
+  auto *SN0 = DAG.getNode(S0);
+  auto *SN1 = DAG.getNode(S1);
+  auto *SN2 = DAG.getNode(S2);
+  auto *SN3 = DAG.getNode(S3);
+  {
+    // Check the common case, when the bundle contains unique nodes.
+    sandboxir::SchedBundle Bndl({SN0, SN1});
+    S0->eraseFromParent();
+    EXPECT_THAT(Bndl, testing::ElementsAre(SN1));
+  }
+  {
+    // Check corner case when the node appears more than once.
+    sandboxir::SchedBundle Bndl({SN2, SN3, SN2});
+    S2->eraseFromParent();
+    EXPECT_THAT(Bndl, testing::ElementsAre(SN3));
+  }
+}
+
 TEST_F(SchedulerTest, Basic) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %v0, i8 %v1) {