xref: /llvm-project/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp (revision 1d09925b4a6fd4af0120825132be23be12fb03d6)
1 //===- Scheduler.cpp ------------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
10 
11 namespace llvm::sandboxir {
12 
13 // TODO: Check if we can cache top/bottom to reduce compile-time.
14 DGNode *SchedBundle::getTop() const {
15   DGNode *TopN = Nodes.front();
16   for (auto *N : drop_begin(Nodes)) {
17     if (N->getInstruction()->comesBefore(TopN->getInstruction()))
18       TopN = N;
19   }
20   return TopN;
21 }
22 
23 DGNode *SchedBundle::getBot() const {
24   DGNode *BotN = Nodes.front();
25   for (auto *N : drop_begin(Nodes)) {
26     if (BotN->getInstruction()->comesBefore(N->getInstruction()))
27       BotN = N;
28   }
29   return BotN;
30 }
31 
32 void SchedBundle::cluster(BasicBlock::iterator Where) {
33   for (auto *N : Nodes) {
34     auto *I = N->getInstruction();
35     if (I->getIterator() == Where)
36       ++Where; // Try to maintain bundle order.
37     I->moveBefore(*Where.getNodeParent(), Where);
38   }
39 }
40 
41 #ifndef NDEBUG
42 void SchedBundle::dump(raw_ostream &OS) const {
43   for (auto *N : Nodes)
44     OS << *N;
45 }
46 
47 void SchedBundle::dump() const {
48   dump(dbgs());
49   dbgs() << "\n";
50 }
51 #endif // NDEBUG
52 
53 #ifndef NDEBUG
54 void ReadyListContainer::dump(raw_ostream &OS) const {
55   auto ListCopy = List;
56   while (!ListCopy.empty()) {
57     OS << *ListCopy.top() << "\n";
58     ListCopy.pop();
59   }
60 }
61 
62 void ReadyListContainer::dump() const {
63   dump(dbgs());
64   dbgs() << "\n";
65 }
66 #endif // NDEBUG
67 
68 void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
69   // Find where we should schedule the instructions.
70   assert(ScheduleTopItOpt && "Should have been set by now!");
71   auto Where = *ScheduleTopItOpt;
72   // Move all instructions in `Bndl` to `Where`.
73   Bndl.cluster(Where);
74   // Update the last scheduled bundle.
75   ScheduleTopItOpt = Bndl.getTop()->getInstruction()->getIterator();
76   // Set nodes as "scheduled" and decrement the UnsceduledSuccs counter of all
77   // dependency predecessors.
78   for (DGNode *N : Bndl) {
79     N->setScheduled(true);
80     for (auto *DepN : N->preds(DAG)) {
81       // TODO: preds() should not return nullptr.
82       if (DepN == nullptr)
83         continue;
84       DepN->decrUnscheduledSuccs();
85       if (DepN->ready())
86         ReadyList.insert(DepN);
87     }
88   }
89 }
90 
91 SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
92   SchedBundle::ContainerTy Nodes;
93   Nodes.reserve(Instrs.size());
94   for (auto *I : Instrs)
95     Nodes.push_back(DAG.getNode(I));
96   auto BndlPtr = std::make_unique<SchedBundle>(std::move(Nodes));
97   auto *Bndl = BndlPtr.get();
98   Bndls.push_back(std::move(BndlPtr));
99   return Bndl;
100 }
101 
102 bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
103   // Use a set of instructions, instead of `Instrs` for fast lookups.
104   DenseSet<Instruction *> InstrsToDefer(Instrs.begin(), Instrs.end());
105   // This collects the nodes that correspond to instructions found in `Instrs`
106   // that have just become ready. These nodes won't be scheduled right away.
107   SmallVector<DGNode *, 8> DeferredNodes;
108 
109   // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
110   // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
111   // which are collected in DeferredNodes) are all ready to schedule.
112   while (!ReadyList.empty()) {
113     auto *ReadyN = ReadyList.pop();
114     if (InstrsToDefer.contains(ReadyN->getInstruction())) {
115       // If the ready instruction is one of those in `Instrs`, then we don't
116       // schedule it right away. Instead we defer it until we can schedule it
117       // along with the rest of the instructions in `Instrs`, at the same
118       // time in a single scheduling bundle.
119       DeferredNodes.push_back(ReadyN);
120       bool ReadyToScheduleDeferred = DeferredNodes.size() == Instrs.size();
121       if (ReadyToScheduleDeferred) {
122         scheduleAndUpdateReadyList(*createBundle(Instrs));
123         return true;
124       }
125     } else {
126       // If the ready instruction is not found in `Instrs`, then we wrap it in a
127       // scheduling bundle and schedule it right away.
128       scheduleAndUpdateReadyList(*createBundle({ReadyN->getInstruction()}));
129     }
130   }
131   assert(DeferredNodes.size() != Instrs.size() &&
132          "We should have succesfully scheduled and early-returned!");
133   return false;
134 }
135 
136 bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
137   assert(all_of(drop_begin(Instrs),
138                 [Instrs](Instruction *I) {
139                   return I->getParent() == (*Instrs.begin())->getParent();
140                 }) &&
141          "Instrs not in the same BB!");
142   // Extend the DAG to include Instrs.
143   Interval<Instruction> Extension = DAG.extend(Instrs);
144   // TODO: Set the window of the DAG that we are interested in.
145   // We start scheduling at the bottom instr of Instrs.
146   auto getBottomI = [](ArrayRef<Instruction *> Instrs) -> Instruction * {
147     return *min_element(Instrs,
148                         [](auto *I1, auto *I2) { return I1->comesBefore(I2); });
149   };
150   ScheduleTopItOpt = std::next(getBottomI(Instrs)->getIterator());
151   // Add nodes to ready list.
152   for (auto &I : Extension) {
153     auto *N = DAG.getNode(&I);
154     if (N->ready())
155       ReadyList.insert(N);
156   }
157   // Try schedule all nodes until we can schedule Instrs back-to-back.
158   return tryScheduleUntil(Instrs);
159 }
160 
161 #ifndef NDEBUG
162 void Scheduler::dump(raw_ostream &OS) const {
163   OS << "ReadyList:\n";
164   ReadyList.dump(OS);
165 }
166 void Scheduler::dump() const { dump(dbgs()); }
167 #endif // NDEBUG
168 
169 } // namespace llvm::sandboxir
170