[SandboxVec][DAG] Implement DAG maintainance on Instruction removal (#127361)

This patch implements dependency maintenance upon receiveing the
notification that an instruction gets deleted.

Author: vporpo

llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp

@@ -483,19 +483,37 @@ void DependencyGraph::notifyEraseInstr(Instruction *I) {
   if (Ctx->getTracker().getState() == Tracker::TrackerState::Reverting)
     // We don't maintain the DAG while reverting.
     return;
-  // Update the MemDGNode chain if this is a memory node.
-  if (auto *MemN = dyn_cast_or_null<MemDGNode>(getNodeOrNull(I))) {
+  auto *N = getNode(I);
+  if (N == nullptr)
+    // Early return if there is no DAG node for `I`.
+    return;
+  if (auto *MemN = dyn_cast<MemDGNode>(getNode(I))) {
+    // Update the MemDGNode chain if this is a memory node.
     auto *PrevMemN = getMemDGNodeBefore(MemN, /*IncludingN=*/false);
     auto *NextMemN = getMemDGNodeAfter(MemN, /*IncludingN=*/false);
     if (PrevMemN != nullptr)
       PrevMemN->NextMemN = NextMemN;
     if (NextMemN != nullptr)
       NextMemN->PrevMemN = PrevMemN;
-  }
 
+    // Drop the memory dependencies from both predecessors and successors.
+    while (!MemN->memPreds().empty()) {
+      auto *PredN = *MemN->memPreds().begin();
+      MemN->removeMemPred(PredN);
+    }
+    while (!MemN->memSuccs().empty()) {
+      auto *SuccN = *MemN->memSuccs().begin();
+      SuccN->removeMemPred(MemN);
+    }
+    // NOTE: The unscheduled succs for MemNodes get updated be setMemPred().
+  } else {
+    // If this is a non-mem node we only need to update UnscheduledSuccs.
+    if (!N->scheduled())
+      for (auto *PredN : N->preds(*this))
+        PredN->decrUnscheduledSuccs();
+  }
+  // Finally erase the Node.
   InstrToNodeMap.erase(I);
-
-  // TODO: Update the dependencies.
 }
 
 void DependencyGraph::notifySetUse(const Use &U, Value *NewSrc) {

llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp

@@ -940,10 +940,11 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %new1, i8 %new2) {
 
 TEST_F(DependencyGraphTest, EraseInstrCallback) {
   parseIR(C, R"IR(
-define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %arg) {
+define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %arg) {
   store i8 %v1, ptr %ptr
   store i8 %v2, ptr %ptr
   store i8 %v3, ptr %ptr
+  store i8 %v4, ptr %ptr
   ret void
 }
 )IR");
@@ -955,25 +956,67 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %arg) {
   auto *S1 = cast<sandboxir::StoreInst>(&*It++);
   auto *S2 = cast<sandboxir::StoreInst>(&*It++);
   auto *S3 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S4NotInDAG = cast<sandboxir::StoreInst>(&*It++);
 
   // Check erase instruction callback.
   sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
   DAG.extend({S1, S3});
+  auto *S1MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
+  auto *S2MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S2));
+  auto *S3MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S3));
+  EXPECT_EQ(S1MemN->getNumUnscheduledSuccs(), 2u);
+  EXPECT_EQ(S2MemN->getNumUnscheduledSuccs(), 1u);
+  EXPECT_EQ(S3MemN->getNumUnscheduledSuccs(), 0u);
   S2->eraseFromParent();
-  auto *DeletedN = DAG.getNodeOrNull(S2);
+  // Check that the DAG Node for S2 no longer exists.
+  auto *DeletedN = DAG.getNode(S2);
   EXPECT_TRUE(DeletedN == nullptr);
+  // Check that dependencies are maintained.
+  EXPECT_THAT(S3MemN->preds(DAG), testing::UnorderedElementsAre(S1MemN));
+  // Also check that UnscheduledSuccs was updated for S1.
+  EXPECT_EQ(S1MemN->getNumUnscheduledSuccs(), 1u);
 
   // Check the MemDGNode chain.
-  auto *S1MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
-  auto *S3MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S3));
   EXPECT_EQ(S1MemN->getNextNode(), S3MemN);
   EXPECT_EQ(S3MemN->getPrevNode(), S1MemN);
 
   // Check the chain when we erase the top node.
   S1->eraseFromParent();
   EXPECT_EQ(S3MemN->getPrevNode(), nullptr);
 
-  // TODO: Check the dependencies to/from NewSN after they land.
+  // Check that we don't crash if we erase a node not in the DAG.
+  S4NotInDAG->eraseFromParent();
+}
+
+// Same but check that we don't update UnscheduledSuccs when Node is scheduled.
+TEST_F(DependencyGraphTest, EraseInstrCallbackScheduled) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %arg) {
+  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 *S1 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S2 = cast<sandboxir::StoreInst>(&*It++);
+
+  sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({S1, S2});
+  auto *S1MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
+  auto *S2MemN = cast<sandboxir::MemDGNode>(DAG.getNode(S2));
+  EXPECT_EQ(S1MemN->getNumUnscheduledSuccs(), 1u);
+  EXPECT_EQ(S2MemN->getNumUnscheduledSuccs(), 0u);
+  // Mark S2 as scheduled and erase it.
+  S2MemN->setScheduled(true);
+  S2->eraseFromParent();
+  EXPECT_EQ(DAG.getNode(S2), nullptr);
+  // Check that we did not update S1's UnscheduledSuccs
+  EXPECT_EQ(S1MemN->getNumUnscheduledSuccs(), 1u);
 }
 
 TEST_F(DependencyGraphTest, MoveInstrCallback) {