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

vporpo · web-flow · commit 1d09925b4a6f · 2024-10-18T16:18:43.000-07:00
This patch implements a ready-list-based scheduler that operates on
DependencyGraph.
It is used by the sandbox vectorizer to test the legality of vectorizing
a group of instrs.

SchedBundle is a helper container, containing all DGNodes that
correspond to the instructions that we are attempting to schedule with
trySchedule(Instrs).
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -113,8 +113,15 @@ class DGNode {
   virtual ~DGNode() = default;
   /// \Returns the number of unscheduled successors.
   unsigned getNumUnscheduledSuccs() const { return UnscheduledSuccs; }
+  void decrUnscheduledSuccs() {
+    assert(UnscheduledSuccs > 0 && "Counting error!");
+    --UnscheduledSuccs;
+  }
+  /// \Returns true if all dependent successors have been scheduled.
+  bool ready() const { return UnscheduledSuccs == 0; }
   /// \Returns true if this node has been scheduled.
   bool scheduled() const { return Scheduled; }
+  void setScheduled(bool NewVal) { Scheduled = NewVal; }
   /// \Returns true if this is before \p Other in program order.
   bool comesBefore(const DGNode *Other) { return I->comesBefore(Other->I); }
   using iterator = PredIterator;
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -0,0 +1,126 @@
+//===- Scheduler.h ----------------------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the bottom-up list scheduler used by the vectorizer. It is used for
+// checking the legality of vectorization and for scheduling instructions in
+// such a way that makes vectorization possible, if legal.
+//
+// The legality check is performed by `trySchedule(Instrs)`, which will try to
+// schedule the IR until all instructions in `Instrs` can be scheduled together
+// back-to-back. If this fails then it is illegal to vectorize `Instrs`.
+//
+// Internally the scheduler uses the vectorizer-specific DependencyGraph class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H
+
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h"
+#include <queue>
+
+namespace llvm::sandboxir {
+
+class PriorityCmp {
+public:
+  bool operator()(const DGNode *N1, const DGNode *N2) {
+    // TODO: This should be a hierarchical comparator.
+    return N1->getInstruction()->comesBefore(N2->getInstruction());
+  }
+};
+
+/// The list holding nodes that are ready to schedule. Used by the scheduler.
+class ReadyListContainer {
+  PriorityCmp Cmp;
+  /// Control/Other dependencies are not modeled by the DAG to save memory.
+  /// These have to be modeled in the ready list for correctness.
+  /// This means that the list will hold back nodes that need to meet such
+  /// unmodeled dependencies.
+  std::priority_queue<DGNode *, std::vector<DGNode *>, PriorityCmp> List;
+
+public:
+  ReadyListContainer() : List(Cmp) {}
+  void insert(DGNode *N) { List.push(N); }
+  DGNode *pop() {
+    auto *Back = List.top();
+    List.pop();
+    return Back;
+  }
+  bool empty() const { return List.empty(); }
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+#endif // NDEBUG
+};
+
+/// The nodes that need to be scheduled back-to-back in a single scheduling
+/// cycle form a SchedBundle.
+class SchedBundle {
+public:
+  using ContainerTy = SmallVector<DGNode *, 4>;
+
+private:
+  ContainerTy Nodes;
+
+public:
+  SchedBundle() = default;
+  SchedBundle(ContainerTy &&Nodes) : Nodes(std::move(Nodes)) {}
+  using iterator = ContainerTy::iterator;
+  using const_iterator = ContainerTy::const_iterator;
+  iterator begin() { return Nodes.begin(); }
+  iterator end() { return Nodes.end(); }
+  const_iterator begin() const { return Nodes.begin(); }
+  const_iterator end() const { return Nodes.end(); }
+  /// \Returns the bundle node that comes before the others in program order.
+  DGNode *getTop() const;
+  /// \Returns the bundle node that comes after the others in program order.
+  DGNode *getBot() const;
+  /// Move all bundle instructions to \p Where back-to-back.
+  void cluster(BasicBlock::iterator Where);
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+};
+
+/// The list scheduler.
+class Scheduler {
+  ReadyListContainer ReadyList;
+  DependencyGraph DAG;
+  std::optional<BasicBlock::iterator> ScheduleTopItOpt;
+  SmallVector<std::unique_ptr<SchedBundle>> Bndls;
+
+  /// \Returns a scheduling bundle containing \p Instrs.
+  SchedBundle *createBundle(ArrayRef<Instruction *> Instrs);
+  /// Schedule nodes until we can schedule \p Instrs back-to-back.
+  bool tryScheduleUntil(ArrayRef<Instruction *> Instrs);
+  /// Schedules all nodes in \p Bndl, marks them as scheduled, updates the
+  /// UnscheduledSuccs counter of all dependency predecessors, and adds any of
+  /// them that become ready to the ready list.
+  void scheduleAndUpdateReadyList(SchedBundle &Bndl);
+
+  /// Disable copies.
+  Scheduler(const Scheduler &) = delete;
+  Scheduler &operator=(const Scheduler &) = delete;
+
+public:
+  Scheduler(AAResults &AA) : DAG(AA) {}
+  ~Scheduler() {}
+
+  bool trySchedule(ArrayRef<Instruction *> Instrs);
+
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -9,6 +9,7 @@ add_llvm_component_library(LLVMVectorize
   SandboxVectorizer/Passes/RegionsFromMetadata.cpp
   SandboxVectorizer/SandboxVectorizer.cpp
   SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
+  SandboxVectorizer/Scheduler.cpp
   SandboxVectorizer/SeedCollector.cpp
   SLPVectorizer.cpp
   Vectorize.cpp
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -60,7 +60,7 @@ bool PredIterator::operator==(const PredIterator &Other) const {
 
 #ifndef NDEBUG
 void DGNode::print(raw_ostream &OS, bool PrintDeps) const {
-  OS << *I << " USuccs:" << UnscheduledSuccs << "\n";
+  OS << *I << " USuccs:" << UnscheduledSuccs << " Sched:" << Scheduled << "\n";
 }
 void DGNode::dump() const { print(dbgs()); }
 void MemDGNode::print(raw_ostream &OS, bool PrintDeps) const {
@@ -249,6 +249,10 @@ void DependencyGraph::setDefUseUnscheduledSuccs(
   // Walk over all instructions in "BotInterval" and update the counter
   // of operands that are in "TopInterval".
   for (Instruction &BotI : BotInterval) {
+    auto *BotN = getNode(&BotI);
+    // Skip scheduled nodes.
+    if (BotN->scheduled())
+      continue;
     for (Value *Op : BotI.operands()) {
       auto *OpI = dyn_cast<Instruction>(Op);
       if (OpI == nullptr)
@@ -286,7 +290,9 @@ void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
         MemDGNodeIntervalBuilder::getBotMemDGNode(TopInterval, *this);
     MemDGNode *LinkBotN =
         MemDGNodeIntervalBuilder::getTopMemDGNode(BotInterval, *this);
-    assert(LinkTopN->comesBefore(LinkBotN) && "Wrong order!");
+    assert((LinkTopN == nullptr || LinkBotN == nullptr ||
+            LinkTopN->comesBefore(LinkBotN)) &&
+           "Wrong order!");
     if (LinkTopN != nullptr && LinkBotN != nullptr) {
       LinkTopN->setNextNode(LinkBotN);
       LinkBotN->setPrevNode(LinkTopN);
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
@@ -0,0 +1,169 @@
+//===- 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 of instructions, instead of `Instrs` for fast lookups.
+  DenseSet<Instruction *> InstrsToDefer(Instrs.begin(), Instrs.end());
+  // This collects the nodes that correspond to instructions found in `Instrs`
+  // that have just become ready. These nodes won't be scheduled right away.
+  SmallVector<DGNode *, 8> DeferredNodes;
+
+  // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
+  // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
+  // which are collected in DeferredNodes) are all ready to schedule.
+  while (!ReadyList.empty()) {
+    auto *ReadyN = ReadyList.pop();
+    if (InstrsToDefer.contains(ReadyN->getInstruction())) {
+      // If the ready instruction is one of those in `Instrs`, then we don't
+      // schedule it right away. Instead we defer it until we can schedule it
+      // along with the rest of the instructions in `Instrs`, at the same
+      // time in a single scheduling bundle.
+      DeferredNodes.push_back(ReadyN);
+      bool ReadyToScheduleDeferred = DeferredNodes.size() == Instrs.size();
+      if (ReadyToScheduleDeferred) {
+        scheduleAndUpdateReadyList(*createBundle(Instrs));
+        return true;
+      }
+    } else {
+      // If the ready instruction is not found in `Instrs`, then we wrap it in a
+      // scheduling bundle and schedule it right away.
+      scheduleAndUpdateReadyList(*createBundle({ReadyN->getInstruction()}));
+    }
+  }
+  assert(DeferredNodes.size() != Instrs.size() &&
+         "We should have succesfully scheduled and early-returned!");
+  return false;
+}
+
+bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
+  assert(all_of(drop_begin(Instrs),
+                [Instrs](Instruction *I) {
+                  return I->getParent() == (*Instrs.begin())->getParent();
+                }) &&
+         "Instrs not in the same BB!");
+  // Extend the DAG to include Instrs.
+  Interval<Instruction> Extension = DAG.extend(Instrs);
+  // TODO: Set the window of the DAG that we are interested in.
+  // We start scheduling at the bottom instr of Instrs.
+  auto getBottomI = [](ArrayRef<Instruction *> Instrs) -> Instruction * {
+    return *min_element(Instrs,
+                        [](auto *I1, auto *I2) { return I1->comesBefore(I2); });
+  };
+  ScheduleTopItOpt = std::next(getBottomI(Instrs)->getIterator());
+  // Add nodes to ready list.
+  for (auto &I : Extension) {
+    auto *N = DAG.getNode(&I);
+    if (N->ready())
+      ReadyList.insert(N);
+  }
+  // Try schedule all nodes until we can schedule Instrs back-to-back.
+  return tryScheduleUntil(Instrs);
+}
+
+#ifndef NDEBUG
+void Scheduler::dump(raw_ostream &OS) const {
+  OS << "ReadyList:\n";
+  ReadyList.dump(OS);
+}
+void Scheduler::dump() const { dump(dbgs()); }
+#endif // NDEBUG
+
+} // namespace llvm::sandboxir
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
@@ -11,5 +11,6 @@ add_llvm_unittest(SandboxVectorizerTests
   DependencyGraphTest.cpp
   IntervalTest.cpp
   LegalityTest.cpp
+  SchedulerTest.cpp
   SeedCollectorTest.cpp	
 )
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -254,6 +254,18 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
   EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u); // N1
   EXPECT_EQ(N1->getNumUnscheduledSuccs(), 0u);
   EXPECT_EQ(N2->getNumUnscheduledSuccs(), 0u);
+
+  // Check decrUnscheduledSuccs.
+  N0->decrUnscheduledSuccs();
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+#ifndef NDEBUG
+  EXPECT_DEATH(N0->decrUnscheduledSuccs(), ".*Counting.*");
+#endif // NDEBUG
+
+  // Check scheduled(), setScheduled().
+  EXPECT_FALSE(N0->scheduled());
+  N0->setScheduled(true);
+  EXPECT_TRUE(N0->scheduled());
 }
 
 TEST_F(DependencyGraphTest, Preds) {
@@ -773,4 +785,16 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
     EXPECT_EQ(S4N->getNumUnscheduledSuccs(), 1u); // S5N
     EXPECT_EQ(S5N->getNumUnscheduledSuccs(), 0u);
   }
+
+  {
+    // Check UnscheduledSuccs when a node is scheduled
+    sandboxir::DependencyGraph DAG(getAA(*LLVMF));
+    DAG.extend({S2, S2});
+    auto *S2N = cast<sandboxir::MemDGNode>(DAG.getNode(S2));
+    S2N->setScheduled(true);
+
+    DAG.extend({S1, S1});
+    auto *S1N = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
+    EXPECT_EQ(S1N->getNumUnscheduledSuccs(), 0u); // S1 is scheduled
+  }
 }
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp

Original file line number	Diff line number	Diff line change
`@@ -11,5 +11,6 @@ add_llvm_unittest(SandboxVectorizerTests`
`11`	`11`	`DependencyGraphTest.cpp`
`12`	`12`	`IntervalTest.cpp`
`13`	`13`	`LegalityTest.cpp`
	`14`	`+ SchedulerTest.cpp`
`14`	`15`	`SeedCollectorTest.cpp`
`15`	`16`	`)`