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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Anshil Gandhi (gandhi56)

<details>
<summary>Changes</summary>

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.


---
Full diff: https://github.com/llvm/llvm-project/pull/212868.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp (+4-2) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp (+43) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index 11149a16b044f..123b4f038371b 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -583,9 +583,11 @@ 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();
+        if (!SuccN->scheduled())
+          SuccN->decrUnscheduledPreds();
     }
   }
   // Finally erase the Node.
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index 831976875f8d3..20bd959c1b107 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
+}

``````````

</details>


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