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