[llvm] 2ff80d2 - [SandboxVec][Scheduler] Fix reassignment of SchedBundle to DGNode

[Vasileios Porpodas via llvm-commits]

Author: Vasileios Porpodas
Date: 2025-02-20T15:28:16-08:00
New Revision: 2ff80d2448f87429b1fb55ea4c112eb9a974298f

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

LOG: [SandboxVec][Scheduler] Fix reassignment of SchedBundle to DGNode

When assigning a bundle to a DAG Node that is already assigned to a
SchedBundle we need to remove the node from the old bundle.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index ede51c28fc94d..aff2068b28331 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -101,7 +101,7 @@ class DGNode {
   /// The scheduler bundle that this node belongs to.
   SchedBundle *SB = nullptr;
 
-  void setSchedBundle(SchedBundle &SB) { this->SB = &SB; }
+  void setSchedBundle(SchedBundle &SB);
   void clearSchedBundle() { this->SB = nullptr; }
   friend class SchedBundle; // For setSchedBundle(), clearSchedBundle().
 

diff  --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
index cdd4db36dafad..03788242ba655 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -115,7 +115,8 @@ class SchedBundle {
   void eraseFromBundle(DGNode *N) {
     Nodes.erase(std::remove(Nodes.begin(), Nodes.end(), N), Nodes.end());
   }
-  friend DGNode::~DGNode(); // For eraseFromBundle().
+  friend void DGNode::setSchedBundle(SchedBundle &); // For eraseFromBunde().
+  friend DGNode::~DGNode();                          // For eraseFromBundle().
 
 public:
   SchedBundle() = default;

diff  --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index 3da52b5b4a6f1..9c869dd1bbdca 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -71,6 +71,12 @@ bool PredIterator::operator==(const PredIterator &Other) const {
   return OpIt == Other.OpIt && MemIt == Other.MemIt;
 }
 
+void DGNode::setSchedBundle(SchedBundle &SB) {
+  if (this->SB != nullptr)
+    this->SB->eraseFromBundle(this);
+  this->SB = &SB;
+}
+
 DGNode::~DGNode() {
   if (SB == nullptr)
     return;

diff  --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
index c0d8178f41480..f0a78477f8cd6 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
@@ -186,6 +186,37 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1, i8 %v2, i8 %v3) {
   }
 }
 
+// Check that assigning a bundle to a DAG Node that is already assigned to a
+// bundle, removes the node from the old bundle.
+TEST_F(SchedulerTest, SchedBundleReassign) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr, i8 %v0, i8 %v1, i8 %v2) {
+  store i8 %v0, ptr %ptr
+  store i8 %v1, ptr %ptr
+  store i8 %v2, 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++);
+
+  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);
+  sandboxir::SchedBundle BndlOld({SN0, SN1});
+  sandboxir::SchedBundle BndlNew({SN0, SN2});
+  EXPECT_THAT(BndlOld, testing::ElementsAre(SN1));
+  EXPECT_THAT(BndlNew, testing::ElementsAre(SN0, SN2));
+}
+
 TEST_F(SchedulerTest, Basic) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %v0, i8 %v1) {