[SandboxVec][DAG] Update DAG whenever a Use is set

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

@@ -313,6 +313,7 @@ class DependencyGraph {
   std::optional<Context::CallbackID> CreateInstrCB;
   std::optional<Context::CallbackID> EraseInstrCB;
   std::optional<Context::CallbackID> MoveInstrCB;
+  std::optional<Context::CallbackID> SetUseCB;
 
   std::unique_ptr<BatchAAResults> BatchAA;
 
@@ -368,6 +369,8 @@ class DependencyGraph {
   /// Called by the callbacks when instruction \p I is about to be moved to
   /// \p To.
   void notifyMoveInstr(Instruction *I, const BBIterator &To);
+  /// Called by the callbacks when \p U's source is about to be set to \p NewSrc
+  void notifySetUse(const Use &U, Value *NewSrc);
 
 public:
   /// This constructor also registers callbacks.
@@ -381,6 +384,8 @@ class DependencyGraph {
         [this](Instruction *I, const BBIterator &To) {
           notifyMoveInstr(I, To);
         });
+    SetUseCB = Ctx.registerSetUseCallback(
+        [this](const Use &U, Value *NewSrc) { notifySetUse(U, NewSrc); });
   }
   ~DependencyGraph() {
     if (CreateInstrCB)
@@ -389,6 +394,8 @@ class DependencyGraph {
       Ctx->unregisterEraseInstrCallback(*EraseInstrCB);
     if (MoveInstrCB)
       Ctx->unregisterMoveInstrCallback(*MoveInstrCB);
+    if (SetUseCB)
+      Ctx->unregisterSetUseCallback(*SetUseCB);
   }
 
   DGNode *getNode(Instruction *I) const {

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

@@ -498,6 +498,21 @@ void DependencyGraph::notifyEraseInstr(Instruction *I) {
   // TODO: Update the dependencies.
 }
 
+void DependencyGraph::notifySetUse(const Use &U, Value *NewSrc) {
+  // Update the UnscheduledSuccs counter for both the current source and NewSrc
+  // if needed.
+  if (auto *CurrSrcI = dyn_cast<Instruction>(U.get())) {
+    if (auto *CurrSrcN = getNode(CurrSrcI)) {
+      CurrSrcN->decrUnscheduledSuccs();
+    }
+  }
+  if (auto *NewSrcI = dyn_cast<Instruction>(NewSrc)) {
+    if (auto *NewSrcN = getNode(NewSrcI)) {
+      ++NewSrcN->UnscheduledSuccs;
+    }
+  }
+}
+
 Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
   if (Instrs.empty())
     return {};

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

@@ -1052,3 +1052,55 @@ define void @foo(ptr %ptr, i8 %v, i8 %v0, i8 %v1, i8 %v2, i8 %v3) {
   DAG.extend({S0, S1});
   EXPECT_TRUE(memDependency(DAG.getNode(S0), DAG.getNode(S1)));
 }
+
+// Setting a Use with a setOperand(), RUW, RAUW etc. can add/remove use-def
+// edges. This needs to maintain the UnscheduledSuccs counter.
+TEST_F(DependencyGraphTest, MaintainUnscheduledSuccsOnUseSet) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v0, i8 %v1) {
+  %add0 = add i8 %v0, %v1
+  %add1 = add i8 %add0, %v1
+  ret void
+}
+)IR");
+  llvm::Function *LLVMF = &*M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(LLVMF);
+  auto *Arg0 = F->getArg(0);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);
+  sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({Add0, Add1});
+  auto *N0 = DAG.getNode(Add0);
+
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+  // Now change %add1 operand to not use %add0.
+  Add1->setOperand(0, Arg0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+  // Restore it: %add0 is now used by %add1.
+  Add1->setOperand(0, Add0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+
+  // RAUW
+  Add0->replaceAllUsesWith(Arg0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+  // Restore it: %add0 is now used by %add1.
+  Add1->setOperand(0, Add0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+
+  // RUWIf
+  Add0->replaceUsesWithIf(Arg0, [](const auto &U) { return true; });
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+  // Restore it: %add0 is now used by %add1.
+  Add1->setOperand(0, Add0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+
+  // RUOW
+  Add1->replaceUsesOfWith(Add0, Arg0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+  // Restore it: %add0 is now used by %add1.
+  Add1->setOperand(0, Add0);
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+}