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#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"
11
12namespace llvm::sandboxir {
13
14// TODO: Check if we can cache top/bottom to reduce compile-time.
15DGNode *SchedBundle::getTop() const {
16 DGNode *TopN = Nodes.front();
17 for (auto *N : drop_begin(RangeOrContainer: Nodes)) {
18 if (N->getInstruction()->comesBefore(Other: TopN->getInstruction()))
19 TopN = N;
20 }
21 return TopN;
22}
23
24DGNode *SchedBundle::getBot() const {
25 DGNode *BotN = Nodes.front();
26 for (auto *N : drop_begin(RangeOrContainer: Nodes)) {
27 if (BotN->getInstruction()->comesBefore(Other: N->getInstruction()))
28 BotN = N;
29 }
30 return BotN;
31}
32
33void SchedBundle::cluster(BasicBlock::iterator Where) {
34 for (auto *N : Nodes) {
35 auto *I = N->getInstruction();
36 if (I->getIterator() == Where)
37 ++Where; // Try to maintain bundle order.
38 I->moveBefore(BB&: *Where.getNodeParent(), WhereIt: Where);
39 }
40}
41
42#ifndef NDEBUG
43void SchedBundle::dump(raw_ostream &OS) const {
44 for (auto *N : Nodes)
45 OS << *N;
46}
47
48void SchedBundle::dump() const {
49 dump(dbgs());
50 dbgs() << "\n";
51}
52#endif // NDEBUG
53
54#ifndef NDEBUG
55void ReadyListContainer::dump(raw_ostream &OS) const {
56 auto ListCopy = List;
57 while (!ListCopy.empty()) {
58 OS << *ListCopy.top() << "\n";
59 ListCopy.pop();
60 }
61}
62
63void ReadyListContainer::dump() const {
64 dump(dbgs());
65 dbgs() << "\n";
66}
67#endif // NDEBUG
68
69void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
70 // Find where we should schedule the instructions.
71 assert(ScheduleTopItOpt && "Should have been set by now!");
72 auto Where = *ScheduleTopItOpt;
73 // Move all instructions in `Bndl` to `Where`.
74 Bndl.cluster(Where);
75 // Update the last scheduled bundle.
76 ScheduleTopItOpt = Bndl.getTop()->getInstruction()->getIterator();
77 // Set nodes as "scheduled" and decrement the UnsceduledSuccs counter of all
78 // dependency predecessors.
79 for (DGNode *N : Bndl) {
80 for (auto *DepN : N->preds(DAG)) {
81 DepN->decrUnscheduledSuccs();
82 if (DepN->ready() && !DepN->scheduled())
83 ReadyList.insert(N: DepN);
84 }
85 N->setScheduled(true);
86 }
87}
88
89void Scheduler::notifyCreateInstr(Instruction *I) {
90 // The DAG notifier should have run by now.
91 auto *N = DAG.getNode(I);
92 // If there is no DAG node for `I` it means that this is out of scope for the
93 // DAG and as such out of scope for the scheduler too, so nothing to do.
94 if (N == nullptr)
95 return;
96 // If the instruction is inserted below the top-of-schedule then we mark it as
97 // "scheduled".
98 bool IsScheduled = ScheduleTopItOpt &&
99 *ScheduleTopItOpt != I->getParent()->end() &&
100 (*ScheduleTopItOpt.value()).comesBefore(Other: I);
101 if (IsScheduled)
102 N->setScheduled(true);
103 // If the new instruction is above the top of schedule we need to remove its
104 // dependency predecessors from the ready list and increment their
105 // `UnscheduledSuccs` counters.
106 if (!IsScheduled) {
107 for (auto *PredN : N->preds(DAG)) {
108 ReadyList.remove(N: PredN);
109 PredN->incrUnscheduledSuccs();
110 }
111 }
112}
113
114SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
115 SchedBundle::ContainerTy Nodes;
116 Nodes.reserve(N: Instrs.size());
117 for (auto *I : Instrs)
118 Nodes.push_back(Elt: DAG.getNode(I));
119 auto BndlPtr = std::make_unique<SchedBundle>(args: std::move(Nodes));
120 auto *Bndl = BndlPtr.get();
121 Bndls[Bndl] = std::move(BndlPtr);
122 return Bndl;
123}
124
125void Scheduler::eraseBundle(SchedBundle *SB) { Bndls.erase(Val: SB); }
126
127bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
128 // Create a bundle for Instrs. If it turns out the schedule is infeasible we
129 // will dismantle it.
130 auto *InstrsSB = createBundle(Instrs);
131 // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
132 // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
133 // which are collected in DeferredNodes) are all ready to schedule.
134 SmallVector<DGNode *> Retry;
135 bool KeepScheduling = true;
136 while (KeepScheduling) {
137 enum class TryScheduleRes {
138 Success, ///> We successfully scheduled the bundle.
139 Failure, ///> We failed to schedule the bundle.
140 Finished, ///> We successfully scheduled the bundle and it is the last
141 /// bundle to be scheduled.
142 };
143 /// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that
144 /// ReadyN is in. If it's not in a bundle it will create a singleton bundle
145 /// and will try to schedule it.
146 auto TryScheduleBndl = [this, InstrsSB](DGNode *ReadyN) -> TryScheduleRes {
147 auto *SB = ReadyN->getSchedBundle();
148 if (SB == nullptr) {
149 // If ReadyN does not belong to a bundle, create a singleton bundle
150 // and schedule it.
151 auto *SingletonSB = createBundle(Instrs: {ReadyN->getInstruction()});
152 scheduleAndUpdateReadyList(Bndl&: *SingletonSB);
153 return TryScheduleRes::Success;
154 }
155 if (SB->ready()) {
156 // Remove the rest of the bundle from the ready list.
157 // TODO: Perhaps change the Scheduler + ReadyList to operate on
158 // SchedBundles instead of DGNodes.
159 for (auto *N : *SB) {
160 if (N != ReadyN)
161 ReadyList.remove(N);
162 }
163 // If all nodes in the bundle are ready.
164 scheduleAndUpdateReadyList(Bndl&: *SB);
165 if (SB == InstrsSB)
166 // We just scheduled InstrsSB bundle, so we are done scheduling.
167 return TryScheduleRes::Finished;
168 return TryScheduleRes::Success;
169 }
170 return TryScheduleRes::Failure;
171 };
172 while (!ReadyList.empty()) {
173 auto *ReadyN = ReadyList.pop();
174 auto Res = TryScheduleBndl(ReadyN);
175 switch (Res) {
176 case TryScheduleRes::Success:
177 // We successfully scheduled ReadyN, keep scheduling.
178 continue;
179 case TryScheduleRes::Failure:
180 // We failed to schedule ReadyN, defer it to later and keep scheduling
181 // other ready instructions.
182 Retry.push_back(Elt: ReadyN);
183 continue;
184 case TryScheduleRes::Finished:
185 // We successfully scheduled the instruction bundle, so we are done.
186 return true;
187 }
188 llvm_unreachable("Unhandled TrySchedule() result");
189 }
190 // Try to schedule nodes from the Retry list.
191 KeepScheduling = false;
192 for (auto *N : make_early_inc_range(Range&: Retry)) {
193 auto Res = TryScheduleBndl(N);
194 if (Res == TryScheduleRes::Success) {
195 Retry.erase(CI: find(Range&: Retry, Val: N));
196 KeepScheduling = true;
197 }
198 }
199 }
200
201 eraseBundle(SB: InstrsSB);
202 return false;
203}
204
205Scheduler::BndlSchedState
206Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
207 assert(!Instrs.empty() && "Expected non-empty bundle");
208 auto *N0 = DAG.getNode(I: Instrs[0]);
209 auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
210 bool AllUnscheduled = SB0 == nullptr;
211 bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
212 for (auto *I : drop_begin(RangeOrContainer&: Instrs)) {
213 auto *N = DAG.getNode(I);
214 auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
215 if (SB != nullptr) {
216 // We found a scheduled instr, so there is now way all are unscheduled.
217 AllUnscheduled = false;
218 if (SB->isSingleton()) {
219 // We found an instruction in a temporarily scheduled singleton. There
220 // is no way that all instructions are scheduled in the same bundle.
221 FullyScheduled = false;
222 }
223 }
224
225 if (SB != SB0) {
226 // Either one of SB, SB0 is null, or they are in different bundles, so
227 // Instrs are definitely not in the same vector bundle.
228 FullyScheduled = false;
229 // One of SB, SB0 are in a vector bundle and they differ.
230 if ((SB != nullptr && !SB->isSingleton()) ||
231 (SB0 != nullptr && !SB0->isSingleton()))
232 return BndlSchedState::AlreadyScheduled;
233 }
234 }
235 return AllUnscheduled ? BndlSchedState::NoneScheduled
236 : FullyScheduled ? BndlSchedState::FullyScheduled
237 : BndlSchedState::TemporarilyScheduled;
238}
239
240void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
241 // | Legend: N: DGNode
242 // N <- DAGInterval.top() | B: SchedBundle
243 // N | *: Contains instruction in Instrs
244 // B <- TopI (Top of schedule) +-------------------------------------------
245 // B
246 // B *
247 // B
248 // B * <- LowestI (Lowest in Instrs)
249 // B
250 // N
251 // N
252 // N <- DAGInterval.bottom()
253 //
254 Instruction *TopI = &*ScheduleTopItOpt.value();
255 Instruction *LowestI = VecUtils::getLowest(Instrs);
256 // Destroy the singleton schedule bundles from LowestI all the way to the top.
257 for (auto *I = LowestI, *E = TopI->getPrevNode(); I != E;
258 I = I->getPrevNode()) {
259 auto *N = DAG.getNode(I);
260 if (N == nullptr)
261 continue;
262 auto *SB = N->getSchedBundle();
263 if (SB->isSingleton())
264 eraseBundle(SB);
265 }
266 // The DAG Nodes contain state like the number of UnscheduledSuccs and the
267 // Scheduled flag. We need to reset their state. We need to do this for all
268 // nodes from LowestI to the top of the schedule. DAG Nodes that are above the
269 // top of schedule that depend on nodes that got reset need to have their
270 // UnscheduledSuccs adjusted.
271 Interval<Instruction> ResetIntvl(TopI, LowestI);
272 for (Instruction &I : ResetIntvl) {
273 auto *N = DAG.getNode(I: &I);
274 N->resetScheduleState();
275 // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even for
276 // nodes above the top of schedule.
277 for (auto *PredN : N->preds(DAG))
278 PredN->incrUnscheduledSuccs();
279 }
280 // Refill the ready list by visiting all nodes from the top of DAG to LowestI.
281 ReadyList.clear();
282 Interval<Instruction> RefillIntvl(DAG.getInterval().top(), LowestI);
283 for (Instruction &I : RefillIntvl) {
284 auto *N = DAG.getNode(I: &I);
285 if (N->ready())
286 ReadyList.insert(N);
287 }
288}
289
290bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
291 assert(all_of(drop_begin(Instrs),
292 [Instrs](Instruction *I) {
293 return I->getParent() == (*Instrs.begin())->getParent();
294 }) &&
295 "Instrs not in the same BB, should have been rejected by Legality!");
296 // TODO: For now don't cross BBs.
297 if (!DAG.getInterval().empty()) {
298 auto *BB = DAG.getInterval().top()->getParent();
299 if (any_of(Range&: Instrs, P: [BB](auto *I) { return I->getParent() != BB; }))
300 return false;
301 }
302 if (ScheduledBB == nullptr)
303 ScheduledBB = Instrs[0]->getParent();
304 // We don't support crossing BBs for now.
305 if (any_of(Range&: Instrs,
306 P: [this](Instruction *I) { return I->getParent() != ScheduledBB; }))
307 return false;
308 auto SchedState = getBndlSchedState(Instrs);
309 switch (SchedState) {
310 case BndlSchedState::FullyScheduled:
311 // Nothing to do.
312 return true;
313 case BndlSchedState::AlreadyScheduled:
314 // Instructions are part of a different vector schedule, so we can't
315 // schedule \p Instrs in the same bundle (without destroying the existing
316 // schedule).
317 return false;
318 case BndlSchedState::TemporarilyScheduled:
319 // If one or more instrs are already scheduled we need to destroy the
320 // top-most part of the schedule that includes the instrs in the bundle and
321 // re-schedule.
322 DAG.extend(Instrs);
323 trimSchedule(Instrs);
324 ScheduleTopItOpt = std::next(x: VecUtils::getLowest(Instrs)->getIterator());
325 return tryScheduleUntil(Instrs);
326 case BndlSchedState::NoneScheduled: {
327 // TODO: Set the window of the DAG that we are interested in.
328 if (!ScheduleTopItOpt)
329 // We start scheduling at the bottom instr of Instrs.
330 ScheduleTopItOpt = std::next(x: VecUtils::getLowest(Instrs)->getIterator());
331 // Extend the DAG to include Instrs.
332 Interval<Instruction> Extension = DAG.extend(Instrs);
333 // Add nodes to ready list.
334 for (auto &I : Extension) {
335 auto *N = DAG.getNode(I: &I);
336 if (N->ready())
337 ReadyList.insert(N);
338 }
339 // Try schedule all nodes until we can schedule Instrs back-to-back.
340 return tryScheduleUntil(Instrs);
341 }
342 }
343 llvm_unreachable("Unhandled BndlSchedState enum");
344}
345
346#ifndef NDEBUG
347void Scheduler::dump(raw_ostream &OS) const {
348 OS << "ReadyList:\n";
349 ReadyList.dump(OS);
350 OS << "Top of schedule: ";
351 if (ScheduleTopItOpt)
352 OS << **ScheduleTopItOpt;
353 else
354 OS << "Empty";
355 OS << "\n";
356}
357void Scheduler::dump() const { dump(dbgs()); }
358#endif // NDEBUG
359
360} // namespace llvm::sandboxir
361