[llvm] [SandboxVec][DAG] Fix notifyEraseInstr to skip scheduled neighbors (PR #212868)

[Anshil Gandhi via llvm-commits]

https://github.com/gandhi56 updated https://github.com/llvm/llvm-project/pull/212868

>From 528ee3726d0c363c9c8c37c1f6154e678790f852 Mon Sep 17 00:00:00 2001
From: Anshil Gandhi <gandhi21299 at gmail.com>
Date: Wed, 29 Jul 2026 16:45:40 -0400
Subject: [PATCH] [SandboxIR] Fix notifyEraseInstr to skip scheduled neighbors

Guard both loops with !PredN->scheduled() / !SuccN->scheduled() so
scheduled neighbors are left untouched, and add a unit test that erases
a node with one scheduled and one unscheduled predecessor to cover the
fix.
---
 .../SandboxVectorizer/DependencyGraph.cpp     |  3 +-
 .../SandboxVectorizer/DependencyGraphTest.cpp | 43 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index 11149a16b044f..55d5419bd2f59 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -583,7 +583,8 @@ void DependencyGraph::notifyEraseInstr(Instruction *I) {
     // If this is a non-mem node we only need to update UnscheduledSuccs.
     if (!N->scheduled()) {
       for (auto *PredN : N->preds(*this))
-        PredN->decrUnscheduledSuccs();
+        if (!PredN->scheduled())
+          PredN->decrUnscheduledSuccs();
       for (auto *SuccN : N->succs(*this))
         SuccN->decrUnscheduledPreds();
     }
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index 831976875f8d3..b789b421d1563 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -1516,3 +1516,46 @@ define void @foo(i8 %v0) {
   Add0->setOperand(0, Sched);
   EXPECT_EQ(Add0N->getNumUnscheduledPreds(), 0u);
 }
+
+// When erasing a non-mem instruction we must not touch the UnscheduledSuccs
+// of an already-scheduled predecessor, since that counter is set to
+// std::nullopt once a node is scheduled.
+TEST_F(DependencyGraphTest, EraseInstrCallbackNonMemWithScheduledPred) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v0) {
+  %predSched = add i8 %v0, 0
+  %predUnsched = add i8 %v0, 1
+  %n = add i8 %predSched, %predUnsched
+  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 *PredSched = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *PredUnsched = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *N = cast<sandboxir::BinaryOperator>(&*It++);
+
+  sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({PredSched, N});
+  auto *PredSchedN = DAG.getNode(PredSched);
+  auto *PredUnschedN = DAG.getNode(PredUnsched);
+  EXPECT_EQ(PredSchedN->getNumUnscheduledSuccs(), 1u);
+  EXPECT_EQ(PredUnschedN->getNumUnscheduledSuccs(), 1u);
+
+  // Mark one of N's predecessors as scheduled. Its UnscheduledSuccs becomes
+  // std::nullopt.
+  PredSchedN->setScheduled();
+
+  // Erase N, which is *not* scheduled. This must not attempt to decrement
+  // the (now invalid) UnscheduledSuccs of PredSchedN, but should still
+  // update the counter of the unscheduled predecessor.
+  N->eraseFromParent();
+  EXPECT_EQ(DAG.getNode(N), nullptr);
+  EXPECT_EQ(PredUnschedN->getNumUnscheduledSuccs(), 0u);
+#ifndef NDEBUG
+  EXPECT_FALSE(PredSchedN->validUnscheduledSuccs());
+#endif
+}