[SandboxVec][Scheduler] Notify scheduler about instruction creation

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

@@ -121,6 +121,7 @@ class DGNode {
     assert(UnscheduledSuccs > 0 && "Counting error!");
     --UnscheduledSuccs;
   }
+  void incrUnscheduledSuccs() { ++UnscheduledSuccs; }
   /// \Returns true if all dependent successors have been scheduled.
   bool ready() const { return UnscheduledSuccs == 0; }
   /// \Returns true if this node has been scheduled.

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

@@ -54,6 +54,22 @@ class ReadyListContainer {
   }
   bool empty() const { return List.empty(); }
   void clear() { List = {}; }
+  /// \Removes \p N if found in the ready list.
+  void remove(DGNode *N) {
+    // TODO: Use a more efficient data-structure for the ready list because the
+    // priority queue does not support fast removals.
+    SmallVector<DGNode *, 8> Keep;
+    Keep.reserve(List.size());
+    while (!List.empty()) {
+      auto *Top = List.top();
+      List.pop();
+      if (Top == N)
+        break;
+      Keep.push_back(Top);
+    }
+    for (auto *KeepN : Keep)
+      List.push(KeepN);
+  }
 #ifndef NDEBUG
   void dump(raw_ostream &OS) const;
   LLVM_DUMP_METHOD void dump() const;
@@ -116,6 +132,8 @@ class Scheduler {
   /// The dependency graph is used by the scheduler to determine the legal
   /// ordering of instructions.
   DependencyGraph DAG;
+  friend class SchedulerInternalsAttorney; // For DAG.
+  Context &Ctx;
   /// This is the top of the schedule, i.e. the location where the scheduler
   /// is about to place the scheduled instructions. It gets updated as we
   /// schedule.
@@ -124,6 +142,12 @@ class Scheduler {
   DenseMap<SchedBundle *, std::unique_ptr<SchedBundle>> Bndls;
   /// The BB that we are currently scheduling.
   BasicBlock *ScheduledBB = nullptr;
+  /// The ID of the callback we register with Sandbox IR.
+  std::optional<Context::CallbackID> CreateInstrCB;
+  /// Called by Sandbox IR's callback system, after \p I has been created.
+  /// NOTE: This should run after DAG's callback has run.
+  // TODO: Perhaps call DAG's notify function from within this one?
+  void notifyCreateInstr(Instruction *I);
 
   /// \Returns a scheduling bundle containing \p Instrs.
   SchedBundle *createBundle(ArrayRef<Instruction *> Instrs);
@@ -153,8 +177,16 @@ class Scheduler {
   Scheduler &operator=(const Scheduler &) = delete;
 
 public:
-  Scheduler(AAResults &AA, Context &Ctx) : DAG(AA, Ctx) {}
-  ~Scheduler() {}
+  Scheduler(AAResults &AA, Context &Ctx) : DAG(AA, Ctx), Ctx(Ctx) {
+    // NOTE: The scheduler's callback depends on the DAG's callback running
+    // before it and updating the DAG accordingly.
+    CreateInstrCB = Ctx.registerCreateInstrCallback(
+        [this](Instruction *I) { notifyCreateInstr(I); });
+  }
+  ~Scheduler() {
+    if (CreateInstrCB)
+      Ctx.unregisterCreateInstrCallback(*CreateInstrCB);
+  }
   /// Tries to build a schedule that includes all of \p Instrs scheduled at the
   /// same scheduling cycle. This essentially checks that there are no
   /// dependencies among \p Instrs. This function may involve scheduling
@@ -180,6 +212,13 @@ class Scheduler {
 #endif
 };
 
+/// A client-attorney class for accessing the Scheduler's internals (used for
+/// unit tests).
+class SchedulerInternalsAttorney {
+public:
+  static DependencyGraph &getDAG(Scheduler &Sched) { return Sched.DAG; }
+};
+
 } // namespace llvm::sandboxir
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H

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

@@ -86,6 +86,31 @@ void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
   }
 }
 
+void Scheduler::notifyCreateInstr(Instruction *I) {
+  // The DAG notifier should have run by now.
+  auto *N = DAG.getNode(I);
+  // If there is no DAG node for `I` it means that this is out of scope for the
+  // DAG and as such out of scope for the scheduler too, so nothing to do.
+  if (N == nullptr)
+    return;
+  // If the instruction is inserted below the top-of-schedule then we mark it as
+  // "scheduled".
+  bool IsScheduled = ScheduleTopItOpt &&
+                     *ScheduleTopItOpt != I->getParent()->end() &&
+                     (*ScheduleTopItOpt.value()).comesBefore(I);
+  if (IsScheduled)
+    N->setScheduled(true);
+  // If the new instruction is above the top of schedule we need to remove its
+  // dependency predecessors from the ready list and increment their
+  // `UnscheduledSuccs` counters.
+  if (!IsScheduled) {
+    for (auto *PredN : N->preds(DAG)) {
+      ReadyList.remove(PredN);
+      PredN->incrUnscheduledSuccs();
+    }
+  }
+}
+
 SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
   SchedBundle::ContainerTy Nodes;
   Nodes.reserve(Instrs.size());

llvm/test/Transforms/SandboxVectorizer/scheduler.ll

@@ -0,0 +1,51 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="bottom-up-vec<>" %s -S | FileCheck %s
+
+; This used to crash because the newly added pack   instructions would not update
+; the DAG and scheduler, leading to def-after-use.
+define void @check_dag_scheduler_update(ptr noalias %p, ptr noalias %p1) {
+; CHECK-LABEL: define void @check_dag_scheduler_update(
+; CHECK-SAME: ptr noalias [[P:%.*]], ptr noalias [[P1:%.*]]) {
+; CHECK-NEXT:    [[I:%.*]] = load i32, ptr [[P]], align 4
+; CHECK-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr i32, ptr [[P]], i64 32
+; CHECK-NEXT:    [[I2:%.*]] = load i32, ptr [[ARRAYIDX4]], align 4
+; CHECK-NEXT:    [[ARRAYIDX11:%.*]] = getelementptr i32, ptr [[P]], i64 33
+; CHECK-NEXT:    [[I4:%.*]] = load i32, ptr [[ARRAYIDX11]], align 4
+; CHECK-NEXT:    [[ARRAYIDX18:%.*]] = getelementptr i32, ptr [[P]], i64 34
+; CHECK-NEXT:    [[I6:%.*]] = load i32, ptr [[ARRAYIDX18]], align 4
+; CHECK-NEXT:    [[PACK:%.*]] = insertelement <4 x i32> poison, i32 [[I]], i32 0
+; CHECK-NEXT:    [[PACK1:%.*]] = insertelement <4 x i32> [[PACK]], i32 [[I2]], i32 1
+; CHECK-NEXT:    [[PACK2:%.*]] = insertelement <4 x i32> [[PACK1]], i32 [[I4]], i32 2
+; CHECK-NEXT:    [[PACK3:%.*]] = insertelement <4 x i32> [[PACK2]], i32 [[I6]], i32 3
+; CHECK-NEXT:    [[VECL:%.*]] = load <4 x i32>, ptr [[P]], align 4
+; CHECK-NEXT:    [[VEC:%.*]] = add nsw <4 x i32> [[PACK3]], [[VECL]]
+; CHECK-NEXT:    store <4 x i32> [[VEC]], ptr [[P1]], align 4
+; CHECK-NEXT:    ret void
+;
+  %i = load i32, ptr %p
+  %i1 = load i32, ptr %p
+  %add = add nsw i32 %i, %i1
+  store i32 %add, ptr %p1
+  %arrayidx4 = getelementptr i32, ptr %p, i64 32
+  %i2 = load i32, ptr %arrayidx4
+  %arrayidx6 = getelementptr i32, ptr %p, i64 1
+  %i3 = load i32, ptr %arrayidx6
+  %add7 = add nsw i32 %i2, %i3
+  %arrayidx9 = getelementptr i32, ptr %p1, i64 1
+  store i32 %add7, ptr %arrayidx9
+  %arrayidx11 = getelementptr i32, ptr %p, i64 33
+  %i4 = load i32, ptr %arrayidx11
+  %arrayidx13 = getelementptr i32, ptr %p, i64 2
+  %i5 = load i32, ptr %arrayidx13
+  %add14 = add nsw i32 %i4, %i5
+  %arrayidx16 = getelementptr i32, ptr %p1, i64 2
+  store i32 %add14, ptr %arrayidx16
+  %arrayidx18 = getelementptr i32, ptr %p, i64 34
+  %i6 = load i32, ptr %arrayidx18
+  %arrayidx19 = getelementptr i32, ptr %p, i64 3
+  %i7 = load i32, ptr %arrayidx19
+  %add21 = add nsw i32 %i6, %i7
+  %arrayidx23 = getelementptr i32, ptr %p1, i64 3
+  store i32 %add21, ptr %arrayidx23
+  ret void
+}

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

@@ -289,3 +289,106 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %v0, i8 %v1) {
     EXPECT_FALSE(Sched.trySchedule({S0, S1}));
   }
 }
+
+TEST_F(SchedulerTest, NotifyCreateInst) {
+  parseIR(C, R"IR(
+define void @foo(ptr noalias %ptr, ptr noalias %ptr1, ptr noalias %ptr2) {
+  %ld0 = load i8, ptr %ptr
+  store i8 %ld0, 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 *L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *S0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+  auto *Ptr1 = F->getArg(1);
+  auto *Ptr2 = F->getArg(2);
+
+  sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);
+  // Schedule Ret and S0. The top of schedule should be at S0.
+  EXPECT_TRUE(Sched.trySchedule({Ret}));
+  EXPECT_TRUE(Sched.trySchedule({S0}));
+  auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);
+  DAG.extend({L0});
+  auto *L0N = DAG.getNode(L0);
+  EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 0u);
+  // We should have DAG nodes for all instructions at this point
+
+  // Now create a new instruction below S0.
+  sandboxir::StoreInst *NewS1 =
+      sandboxir::StoreInst::create(L0, Ptr1, Align(8), Ret->getIterator(),
+                                   /*IsVolatile=*/false, Ctx);
+  // Check that it is marked as "scheduled".
+  auto *NewS1N = DAG.getNode(NewS1);
+  EXPECT_TRUE(NewS1N->scheduled());
+  // Check that L0's UnscheduledSuccs are still == 0 since NewS1 is "scheduled".
+  EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 0u);
+
+  // Now create a new instruction above S0.
+  sandboxir::StoreInst *NewS2 =
+      sandboxir::StoreInst::create(L0, Ptr2, Align(8), S0->getIterator(),
+                                   /*IsVolatile=*/false, Ctx);
+  // Check that it is not marked as "scheduled".
+  auto *NewS2N = DAG.getNode(NewS2);
+  EXPECT_FALSE(NewS2N->scheduled());
+  // Check that L0's UnscheduledSuccs got updated because of NewS2.
+  EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 1u);
+
+  sandboxir::ReadyListContainer ReadyList;
+  // Check empty().
+  EXPECT_TRUE(ReadyList.empty());
+}
+
+TEST_F(SchedulerTest, ReadyList) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %ld0 = load i8, ptr %ptr
+  store i8 %ld0, 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 *L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *S0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+
+  sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({&*BB->begin(), BB->getTerminator()});
+  auto *L0N = DAG.getNode(L0);
+  auto *S0N = DAG.getNode(S0);
+  auto *RetN = DAG.getNode(Ret);
+
+  sandboxir::ReadyListContainer ReadyList;
+  // Check empty().
+  EXPECT_TRUE(ReadyList.empty());
+  // Check insert(), pop().
+  ReadyList.insert(L0N);
+  EXPECT_FALSE(ReadyList.empty());
+  EXPECT_EQ(ReadyList.pop(), L0N);
+  // Check clear().
+  ReadyList.insert(L0N);
+  EXPECT_FALSE(ReadyList.empty());
+  ReadyList.clear();
+  EXPECT_TRUE(ReadyList.empty());
+  // Check remove().
+  EXPECT_TRUE(ReadyList.empty());
+  ReadyList.remove(L0N); // Removing a non-existing node should be valid.
+  ReadyList.insert(L0N);
+  ReadyList.insert(S0N);
+  ReadyList.insert(RetN);
+  ReadyList.remove(S0N);
+  DenseSet<sandboxir::DGNode *> Nodes;
+  Nodes.insert(ReadyList.pop());
+  Nodes.insert(ReadyList.pop());
+  EXPECT_TRUE(ReadyList.empty());
+  EXPECT_THAT(Nodes, testing::UnorderedElementsAre(L0N, RetN));
+}