[llvm] [SandboxVec][Scheduler] Boilerplate and initial implementation. (PR #112449)

[Sriraman Tallam via llvm-commits]

================
@@ -0,0 +1,161 @@
+//===- Scheduler.cpp ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
+
+namespace llvm::sandboxir {
+
+// TODO: Check if we can cache top/bottom to reduce compile-time.
+DGNode *SchedBundle::getTop() const {
+  DGNode *TopN = Nodes.front();
+  for (auto *N : drop_begin(Nodes)) {
+    if (N->getInstruction()->comesBefore(TopN->getInstruction()))
+      TopN = N;
+  }
+  return TopN;
+}
+
+DGNode *SchedBundle::getBot() const {
+  DGNode *BotN = Nodes.front();
+  for (auto *N : drop_begin(Nodes)) {
+    if (BotN->getInstruction()->comesBefore(N->getInstruction()))
+      BotN = N;
+  }
+  return BotN;
+}
+
+void SchedBundle::cluster(BasicBlock::iterator Where) {
+  for (auto *N : Nodes) {
+    auto *I = N->getInstruction();
+    if (I->getIterator() == Where)
+      ++Where; // Try to maintain bundle order.
+    I->moveBefore(*Where.getNodeParent(), Where);
+  }
+}
+
+#ifndef NDEBUG
+void SchedBundle::dump(raw_ostream &OS) const {
+  for (auto *N : Nodes)
+    OS << *N;
+}
+
+void SchedBundle::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+#ifndef NDEBUG
+void ReadyListContainer::dump(raw_ostream &OS) const {
+  auto ListCopy = List;
+  while (!ListCopy.empty()) {
+    OS << *ListCopy.top() << "\n";
+    ListCopy.pop();
+  }
+}
+
+void ReadyListContainer::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
+  // Find where we should schedule the instructions.
+  assert(ScheduleTopItOpt && "Should have been set by now!");
+  auto Where = *ScheduleTopItOpt;
+  // Move all instructions in `Bndl` to `Where`.
+  Bndl.cluster(Where);
+  // Update the last scheduled bundle.
+  ScheduleTopItOpt = Bndl.getTop()->getInstruction()->getIterator();
+  // 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)) {
+      // TODO: preds() should not return nullptr.
+      if (DepN == nullptr)
+        continue;
+      DepN->decrUnscheduledSuccs();
+      if (DepN->ready())
+        ReadyList.insert(DepN);
+    }
+  }
+}
+
+SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
+  SchedBundle::ContainerTy Nodes;
+  Nodes.reserve(Instrs.size());
+  for (auto *I : Instrs)
+    Nodes.push_back(DAG.getNode(I));
+  auto BndlPtr = std::make_unique<SchedBundle>(std::move(Nodes));
+  auto *Bndl = BndlPtr.get();
+  Bndls.push_back(std::move(BndlPtr));
+  return Bndl;
+}
+
+bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
+  // Use a set for fast lookups.
+  DenseSet<Instruction *> InstrsToDefer(Instrs.begin(), Instrs.end());
+  SmallVector<DGNode *, 8> DeferredNodes;
+
+  // Keep scheduling ready nodes.
+  while (!ReadyList.empty()) {
+    auto *ReadyN = ReadyList.pop();
+    // We defer scheduling of instructions in `Instrs` until we can schedule all
+    // of them at the same time in a single scheduling bundle.
+    if (InstrsToDefer.contains(ReadyN->getInstruction())) {
+      DeferredNodes.push_back(ReadyN);
+      bool ReadyToScheduleDeferred = DeferredNodes.size() == Instrs.size();
+      if (ReadyToScheduleDeferred) {
+        scheduleAndUpdateReadyList(*createBundle(Instrs));
+        return true;
+      }
+    } else {
+      scheduleAndUpdateReadyList(*createBundle({ReadyN->getInstruction()}));
+    }
----------------
tmsri wrote:

I feel like this can be structured differently. First, my understanding:

-  You are trying to schedule Instrs and ReadyList already is populated with the DGNodes for those Instrs.
-  In the while loop, you check if the ReadyNode's instruction is part of Instrs.  If not, you create a singleton bundle and schedule it.

I am asking why would a DGNode's->getInstruction() not be contained in InstrsToDefer?  I believe that is because when we schedule a node, it's preds() could become ready and get added to the list, in function "scheduleAndUpdateReadyList".  But, that would get added only to the end since Ready list is a queue.  Is this understanding correct?

If that is the case, you can split the if into two parts.  You will bundle all the DeferredNodes first and schedule it and then schedule the remaining singletons.  



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