[SandboxVec][Scheduler] Don't insert scheduled instrs into the ready list (#127688)

vporpo · web-flow · commit 0cc7381543e3 · 2025-02-18T16:17:46.000-08:00
In a particular scenario (see test) we used to insert scheduled
instructions into the ready list. This patch fixes this by fixing the
trimSchedule() function.
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -123,6 +123,10 @@ class DGNode {
     --UnscheduledSuccs;
   }
   void incrUnscheduledSuccs() { ++UnscheduledSuccs; }
+  void resetScheduleState() {
+    UnscheduledSuccs = 0;
+    Scheduled = false;
+  }
   /// \Returns true if all dependent successors have been scheduled.
   bool ready() const { return UnscheduledSuccs == 0; }
   /// \Returns true if this node has been scheduled.
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -61,7 +61,18 @@ class ReadyListContainer {
 
 public:
   ReadyListContainer() : List(Cmp) {}
-  void insert(DGNode *N) { List.push(N); }
+  void insert(DGNode *N) {
+#ifndef NDEBUG
+    assert(!N->scheduled() && "Don't insert a scheduled node!");
+    auto ListCopy = List;
+    while (!ListCopy.empty()) {
+      DGNode *Top = ListCopy.top();
+      ListCopy.pop();
+      assert(Top != N && "Node already exists in ready list!");
+    }
+#endif
+    List.push(N);
+  }
   DGNode *pop() {
     auto *Back = List.top();
     List.pop();
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
@@ -77,12 +77,12 @@ void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
   // Set nodes as "scheduled" and decrement the UnsceduledSuccs counter of all
   // dependency predecessors.
   for (DGNode *N : Bndl) {
-    N->setScheduled(true);
     for (auto *DepN : N->preds(DAG)) {
       DepN->decrUnscheduledSuccs();
-      if (DepN->ready())
+      if (DepN->ready() && !DepN->scheduled())
         ReadyList.insert(DepN);
     }
+    N->setScheduled(true);
   }
 }
 
@@ -188,6 +188,19 @@ Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
 }
 
 void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
+  //                                |  Legend: N: DGNode
+  //  N <- DAGInterval.top()        |          B: SchedBundle
+  //  N                             |          *: Contains instruction in Instrs
+  //  B <- TopI (Top of schedule)   +-------------------------------------------
+  //  B
+  //  B *
+  //  B
+  //  B * <- LowestI (Lowest in Instrs)
+  //  B
+  //  N
+  //  N
+  //  N <- DAGInterval.bottom()
+  //
   Instruction *TopI = &*ScheduleTopItOpt.value();
   Instruction *LowestI = VecUtils::getLowest(Instrs);
   // Destroy the schedule bundles from LowestI all the way to the top.
@@ -199,13 +212,28 @@ void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
     if (auto *SB = N->getSchedBundle())
       eraseBundle(SB);
   }
-  // TODO: For now we clear the DAG. Trim view once it gets implemented.
-  Bndls.clear();
-  DAG.clear();
-
-  // Since we are scheduling NewRegion from scratch, we clear the ready lists.
-  // The nodes currently in the list may not be ready after clearing the View.
+  // The DAG Nodes contain state like the number of UnscheduledSuccs and the
+  // Scheduled flag. We need to reset their state. We need to do this for all
+  // nodes from LowestI to the top of the schedule. DAG Nodes that are above the
+  // top of schedule that depend on nodes that got reset need to have their
+  // UnscheduledSuccs adjusted.
+  Interval<Instruction> ResetIntvl(TopI, LowestI);
+  for (Instruction &I : ResetIntvl) {
+    auto *N = DAG.getNode(&I);
+    N->resetScheduleState();
+    // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even for
+    // nodes above the top of schedule.
+    for (auto *PredN : N->preds(DAG))
+      PredN->incrUnscheduledSuccs();
+  }
+  // Refill the ready list by visiting all nodes from the top of DAG to LowestI.
   ReadyList.clear();
+  Interval<Instruction> RefillIntvl(DAG.getInterval().top(), LowestI);
+  for (Instruction &I : RefillIntvl) {
+    auto *N = DAG.getNode(&I);
+    if (N->ready())
+      ReadyList.insert(N);
+  }
 }
 
 bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
@@ -214,6 +242,12 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
                   return I->getParent() == (*Instrs.begin())->getParent();
                 }) &&
          "Instrs not in the same BB, should have been rejected by Legality!");
+  // TODO: For now don't cross BBs.
+  if (!DAG.getInterval().empty()) {
+    auto *BB = DAG.getInterval().top()->getParent();
+    if (any_of(Instrs, [BB](auto *I) { return I->getParent() != BB; }))
+      return false;
+  }
   if (ScheduledBB == nullptr)
     ScheduledBB = Instrs[0]->getParent();
   // We don't support crossing BBs for now.
@@ -230,21 +264,13 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
     // top-most part of the schedule that includes the instrs in the bundle and
     // re-schedule.
     trimSchedule(Instrs);
-    ScheduleTopItOpt = std::nullopt;
-    [[fallthrough]];
+    ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
+    return tryScheduleUntil(Instrs);
   case BndlSchedState::NoneScheduled: {
     // TODO: Set the window of the DAG that we are interested in.
     if (!ScheduleTopItOpt)
       // We start scheduling at the bottom instr of Instrs.
       ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
-
-    // TODO: For now don't cross BBs.
-    if (!DAG.getInterval().empty()) {
-      auto *BB = DAG.getInterval().top()->getParent();
-      if (any_of(Instrs, [BB](auto *I) { return I->getParent() != BB; }))
-        return false;
-    }
-
     // Extend the DAG to include Instrs.
     Interval<Instruction> Extension = DAG.extend(Instrs);
     // Add nodes to ready list.
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
@@ -246,6 +246,77 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) {
   EXPECT_TRUE(Sched.trySchedule({L0, L1}));
 }
 
+// Check scheduling in the following order: {A0,A1},{B0,B1},{C0,C1},{D0,D1}
+// assuming program order: B0,B1,C0,C1,D0,D1,E0,D1.
+// This will effectively schedule nodes below already scheduled nodes, which
+// can expose issues in the code that adds nodes to the ready list.
+// For example, we schedule {D0,D1} while {C0,C1} are scheduled and there is
+// a dependency D0->C0 and D1->C1.
+//
+//                   {A0,A1}  {B0,B1}  {C0,C1}  {D0,D1}
+//   B0,B1                    | S
+//   |\                       |
+//   | C0,C1                  |        | S      | S
+//   |  | \                   |                 |
+//   |  |  D0,D1              |                 | S
+//   | /                      |
+//   A0,A1             | S    | S
+//                 +------------------------+
+//                 | Legend   |: DAG        |
+//                 |          S: Scheduled  |
+TEST_F(SchedulerTest, ScheduledPredecessors) {
+  parseIR(C, R"IR(
+define void @foo(ptr noalias %ptrA0, ptr noalias %ptrA1,
+                 ptr noalias %ptrB0, ptr noalias %ptrB1,
+                 ptr noalias %ptrD0, ptr noalias %ptrD1) {
+  %B0 = load i8, ptr %ptrB0
+  %B1 = load i8, ptr %ptrB1
+  %C0 = add i8 %B0, 0
+  %C1 = add i8 %B1, 1
+  store i8 %C0, ptr %ptrD0
+  store i8 %C1, ptr %ptrD1
+  store i8 %B0, ptr %ptrA0
+  store i8 %B1, ptr %ptrA1
+  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 *B1 = cast<sandboxir::LoadInst>(&*It++);
+  auto *B0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *C1 = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *C0 = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *D1 = cast<sandboxir::StoreInst>(&*It++);
+  auto *D0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *A1 = cast<sandboxir::StoreInst>(&*It++);
+  auto *A0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+
+  sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);
+  EXPECT_TRUE(Sched.trySchedule({A0, A1}));
+  // NOTE: We schedule the intermediate nodes between {A0,A1} and {B0,B1} by
+  // hand one by one to make sure they are scheduled in that order because
+  // the scheduler may reorder them a bit if we let it do it.
+  EXPECT_TRUE(Sched.trySchedule(D0));
+  EXPECT_TRUE(Sched.trySchedule(D1));
+  EXPECT_TRUE(Sched.trySchedule(C0));
+  EXPECT_TRUE(Sched.trySchedule(C1));
+  EXPECT_TRUE(Sched.trySchedule({B0, B1}));
+  // At this point all nodes must have been scheduled from B0,B1 to A0,A1.
+  // The ones in between are scheduled as single-instruction nodes.
+  // So when we attempt to schedule {C0,C1} we will need to reschedule.
+  // At this point we will trim the schedule from {C0,C1} upwards.
+  EXPECT_TRUE(Sched.trySchedule({C0, C1}));
+  // Now the schedule should only contain {C0,C1} which should be marked as
+  // "scheduled".
+  // {D0,D1} are below {C0,C1}, so we grow the DAG downwards, while
+  // {C0,C1} are marked as "scheduled" above them.
+  EXPECT_TRUE(Sched.trySchedule({D0, D1}));
+}
+
 TEST_F(SchedulerTest, DontCrossBBs) {
   parseIR(C, R"IR(
 define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %v0, i8 %v1) {
