[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.

Author: Vasileios Porpodas

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().
 

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;

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;

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) {