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

[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()}));
+    }
----------------
vporpo wrote:

I think this is not entirely correct. The problem is that some or all instructions in `Instrs` won't be ready to get scheduled. So what we are trying to do is keep scheduling program instructions individually as they become ready, until either we get all instructions in `Instrs` ready at the same time, or we run out of instructions to schedule. As program instructions become ready, some of them may be in `Instrs`, and these are the ones that are "deferred", because we don't want to schedule them (which would make their predecessors ready etc.). Once the deferred ones match `Instrs`, then we proceed to schedule them.

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