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
68void SchedulingPoint::print(raw_ostream &OS) const {
69 if (BasicBlock *BB = atBeforeBeginOrNull())
70 OS << "Before begin of BB " << BB->getName();
71 else if (BasicBlock *BB = atEndOrNull())
72 OS << "At end of BB " << BB->getName();
73 else
74 OS << "At instr: " << *atInstrOrNull();
75}
76
77void SchedulingPoint::dump() const {
78 print(dbgs());
79 dbgs() << "\n";
80}
81#endif // NDEBUG
82
83void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
84 // Find where we should schedule the instructions.
85 assert(ScheduleFrontierOpt && "Should have been set by now!");
86 auto Where = Dir == SchedDirection::BottomUp
87 ? ScheduleFrontierOpt->getIterator()
88 : ScheduleFrontierOpt->getNext().getIterator();
89 // Move all instructions in `Bndl` to `Where`.
90 Bndl.cluster(Where);
91 // Update the last scheduled bundle.
92 ScheduleFrontierOpt = Dir == SchedDirection::BottomUp
93 ? Bndl.getTop()->getInstruction()->getIterator()
94 : Bndl.getBot()->getInstruction()->getIterator();
95 // Set nodes as "scheduled" and decrement the UnscheduledSuccs/Preds counter
96 // of all dependency predecessors/successors.
97 for (DGNode *N : Bndl) {
98 switch (Dir) {
99 case SchedDirection::BottomUp: {
100 for (auto *DepN : N->preds(DAG)) {
101 DepN->decrUnscheduledDeps();
102 if (DepN->ready() && !DepN->scheduled())
103 ReadyList.insert(N: DepN);
104 }
105 break;
106 }
107 case SchedDirection::TopDown: {
108 for (auto *DepN : N->succs(DAG)) {
109 DepN->decrUnscheduledDeps();
110 if (DepN->ready() && !DepN->scheduled())
111 ReadyList.insert(N: DepN);
112 }
113 break;
114 }
115 }
116 N->setScheduled();
117 }
118}
119
120void Scheduler::notifyCreateInstr(Instruction *I) {
121 // The DAG notifier should have run by now.
122 auto *N = DAG.getNode(I);
123 // If there is no DAG node for `I` it means that this is out of scope for the
124 // DAG and as such out of scope for the scheduler too, so nothing to do.
125 if (N == nullptr)
126 return;
127 // If the instruction is inserted below the top-of-schedule then we mark it as
128 // "scheduled".
129 bool IsScheduled =
130 ScheduleFrontierOpt &&
131 ScheduleFrontierOpt->getIterator() != I->getParent()->end() &&
132 ((Dir == SchedDirection::BottomUp &&
133 (*ScheduleFrontierOpt.value()).comesBefore(Other: I)) ||
134 (Dir == SchedDirection::TopDown &&
135 I->comesBefore(Other: &*ScheduleFrontierOpt.value())));
136 if (IsScheduled)
137 N->setScheduled();
138 // If the new instruction is above the top of schedule we need to remove its
139 // dependency predecessors from the ready list and increment their
140 // `UnscheduledSuccs` counters.
141 if (!IsScheduled) {
142 if (Dir == SchedDirection::BottomUp) {
143 for (auto *PredN : N->preds(DAG)) {
144 ReadyList.remove(N: PredN);
145 PredN->incrUnscheduledDeps();
146 }
147 } else {
148 for (auto *SuccN : N->succs(DAG)) {
149 ReadyList.remove(N: SuccN);
150 SuccN->incrUnscheduledDeps();
151 }
152 }
153 }
154}
155
156void Scheduler::notifyEraseInstr(Instruction *I) {
157 // We don't maintain the state while reverting.
158 if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
159 return;
160 auto *N = DAG.getNode(I);
161 if (N == nullptr)
162 return;
163 ReadyList.remove(N);
164 // Also decrement the unscheduledDep counter for the dependents and add them
165 // to the ready list if they become ready.
166 auto UpdateNodeAndTryAddToReadyList = [this, N](DGNode *DepN) {
167 if (DepN->scheduled())
168 return;
169 if (!N->scheduled() && !DepN->ready())
170 DepN->decrUnscheduledDeps();
171 if (DepN->ready() && !ReadyList.contains(N: DepN))
172 ReadyList.insert(N: DepN);
173 };
174 if (Dir == SchedDirection::BottomUp) {
175 for (auto *DepN : N->preds(DAG))
176 UpdateNodeAndTryAddToReadyList(DepN);
177 } else if (Dir == SchedDirection::TopDown) {
178 for (auto *DepN : N->succs(DAG))
179 UpdateNodeAndTryAddToReadyList(DepN);
180 }
181}
182
183void Scheduler::notifyMoveInstr(Instruction *I, const BBIterator &To) {
184 // We don't maintain the state while reverting.
185 if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
186 return;
187 // We assume that the dependencies have not changed because the user will
188 // only attempt instruction moves that don't modify the dependencies, because
189 // if they did they would not be legal.
190 //
191 // If this assumption does not hold, we would need to empty the ready list and
192 // re-fill it.
193}
194void Scheduler::notifySetUse(const Use &U, Value *NewSrc) {
195 // We don't maintain the state while reverting.
196 if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
197 return;
198 Instruction *DstI = cast<Instruction>(Val: U.getUser());
199 DGNode *DstN = DAG.getNode(I: DstI);
200 Value *OldSrc = U.get();
201 DGNode *OldSrcN = isa<Instruction>(Val: OldSrc)
202 ? DAG.getNode(I: cast<Instruction>(Val: OldSrc))
203 : nullptr;
204 DGNode *NewSrcN = isa<Instruction>(Val: NewSrc)
205 ? DAG.getNode(I: cast<Instruction>(Val: NewSrc))
206 : nullptr;
207 switch (Dir) {
208 case SchedDirection::BottomUp: {
209 // Check if OldSrc is now ready and add it to the ready list.
210 if (OldSrcN && OldSrcN->ready() && !OldSrcN->scheduled() &&
211 !ReadyList.contains(N: OldSrcN))
212 ReadyList.insert(N: OldSrcN);
213 // Check if NewSrcN needs to be removed from the ready list.
214 if (NewSrcN && (!DstN || !DstN->scheduled()) && !NewSrcN->ready())
215 ReadyList.remove(N: NewSrcN);
216 break;
217 }
218 case SchedDirection::TopDown: {
219 // Check if we need to add DstN to the ready list.
220 if (DstN && DstN->ready() && !NewSrcN->scheduled() &&
221 !ReadyList.contains(N: NewSrcN))
222 ReadyList.insert(N: NewSrcN);
223 // Check if we need to remove DstN from the ready list.
224 if (DstN && !DstN->ready())
225 ReadyList.remove(N: NewSrcN);
226 break;
227 }
228 }
229}
230
231SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
232 SchedBundle::ContainerTy Nodes;
233 Nodes.reserve(N: Instrs.size());
234 for (auto *I : Instrs)
235 Nodes.push_back(Elt: DAG.getNode(I));
236 auto BndlPtr = std::make_unique<SchedBundle>(args: std::move(Nodes));
237 auto *Bndl = BndlPtr.get();
238 Bndls[Bndl] = std::move(BndlPtr);
239 return Bndl;
240}
241
242void Scheduler::eraseBundle(SchedBundle *SB) { Bndls.erase(Val: SB); }
243
244bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
245 // Create a bundle for Instrs. If it turns out the schedule is infeasible we
246 // will dismantle it.
247 auto *InstrsSB = createBundle(Instrs);
248 // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
249 // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
250 // which are collected in DeferredNodes) are all ready to schedule.
251 SmallVector<DGNode *> Retry;
252 bool KeepScheduling = true;
253 while (KeepScheduling) {
254 enum class TryScheduleRes {
255 Success, ///> We successfully scheduled the bundle.
256 Failure, ///> We failed to schedule the bundle.
257 Finished, ///> We successfully scheduled the bundle and it is the last
258 /// bundle to be scheduled.
259 };
260 /// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that
261 /// ReadyN is in. If it's not in a bundle it will create a singleton bundle
262 /// and will try to schedule it.
263 auto TryScheduleBndl = [this, InstrsSB](DGNode *ReadyN) -> TryScheduleRes {
264 auto *SB = ReadyN->getSchedBundle();
265 if (SB == nullptr) {
266 // If ReadyN does not belong to a bundle, create a singleton bundle
267 // and schedule it.
268 auto *SingletonSB = createBundle(Instrs: {ReadyN->getInstruction()});
269 scheduleAndUpdateReadyList(Bndl&: *SingletonSB);
270 return TryScheduleRes::Success;
271 }
272 if (SB->ready(Dir)) {
273 // Remove the rest of the bundle from the ready list.
274 // TODO: Perhaps change the Scheduler + ReadyList to operate on
275 // SchedBundles instead of DGNodes.
276 for (auto *N : *SB) {
277 if (N != ReadyN)
278 ReadyList.remove(N);
279 }
280 // If all nodes in the bundle are ready.
281 scheduleAndUpdateReadyList(Bndl&: *SB);
282 if (SB == InstrsSB)
283 // We just scheduled InstrsSB bundle, so we are done scheduling.
284 return TryScheduleRes::Finished;
285 return TryScheduleRes::Success;
286 }
287 return TryScheduleRes::Failure;
288 };
289 while (!ReadyList.empty()) {
290 auto *ReadyN = ReadyList.pop();
291 auto Res = TryScheduleBndl(ReadyN);
292 switch (Res) {
293 case TryScheduleRes::Success:
294 // We successfully scheduled ReadyN's bundle, keep scheduling.
295 continue;
296 case TryScheduleRes::Failure:
297 // We failed to schedule ReadyN's bundle, defer it to later and keep
298 // scheduling other ready instructions.
299 Retry.push_back(Elt: ReadyN);
300 continue;
301 case TryScheduleRes::Finished:
302 // We successfully scheduled the instruction bundle, so we are done.
303 return true;
304 }
305 llvm_unreachable("Unhandled TrySchedule() result");
306 }
307 // Try to schedule nodes from the Retry list.
308 KeepScheduling = false;
309 for (auto *N : make_early_inc_range(Range&: Retry)) {
310 auto Res = TryScheduleBndl(N);
311 if (Res == TryScheduleRes::Success) {
312 Retry.erase(CI: find(Range&: Retry, Val: N));
313 KeepScheduling = true;
314 }
315 }
316 }
317
318 // The Retry vector contains the ready nodes that were removed from the ready
319 // list but we could not schedule them (along with their parent bundle).
320 // Insert them back in.
321 for (auto *RetryN : Retry)
322 ReadyList.insert(N: RetryN);
323
324 eraseBundle(SB: InstrsSB);
325 return false;
326}
327
328Scheduler::BndlSchedState
329Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
330 assert(!Instrs.empty() && "Expected non-empty bundle");
331 auto *N0 = DAG.getNode(I: Instrs[0]);
332 auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
333 bool AllUnscheduled = SB0 == nullptr;
334 bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
335 for (auto *I : drop_begin(RangeOrContainer&: Instrs)) {
336 auto *N = DAG.getNode(I);
337 auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
338 if (SB != nullptr) {
339 // We found a scheduled instr, so there is now way all are unscheduled.
340 AllUnscheduled = false;
341 if (SB->isSingleton()) {
342 // We found an instruction in a temporarily scheduled singleton. There
343 // is no way that all instructions are scheduled in the same bundle.
344 FullyScheduled = false;
345 }
346 }
347
348 if (SB != SB0) {
349 // Either one of SB, SB0 is null, or they are in different bundles, so
350 // Instrs are definitely not in the same vector bundle.
351 FullyScheduled = false;
352 // One of SB, SB0 are in a vector bundle and they differ.
353 if ((SB != nullptr && !SB->isSingleton()) ||
354 (SB0 != nullptr && !SB0->isSingleton()))
355 return BndlSchedState::AlreadyScheduled;
356 }
357 }
358 return AllUnscheduled ? BndlSchedState::NoneScheduled
359 : FullyScheduled ? BndlSchedState::FullyScheduled
360 : BndlSchedState::TemporarilyScheduled;
361}
362
363void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
364 // | Legend: N: DGNode
365 // N <- DAGInterval.top() | B: SchedBundle
366 // N | *: Contains instruction in Instrs
367 // B <- TopI (Top of schedule) +-------------------------------------------
368 // B
369 // B *
370 // B
371 // B * <- LowestI (Lowest in Instrs)
372 // B
373 // N
374 // N
375 // N <- DAGInterval.bottom()
376 //
377 // Note: this figure assumes bottom-up scheduling. In top-down we have the
378 // top-down mirror image.
379 Instruction *TopI = Dir == SchedDirection::BottomUp
380 ? &*ScheduleFrontierOpt.value()
381 : VecUtils::getHighest(Instrs);
382 Instruction *LowestI = Dir == SchedDirection::BottomUp
383 ? VecUtils::getLowest(Instrs)
384 : &*ScheduleFrontierOpt.value();
385 Interval<Instruction> ResetIntvl(TopI, LowestI);
386 // The DAG Nodes contain state like the number of UnscheduledSuccs and the
387 // Scheduled flag. We need to reset their state. We need to do this for all
388 // nodes in ResetIntvl. Also destroy the singleton schedule bundles from
389 // LowestI all the way to the top.
390 for (auto &I : ResetIntvl) {
391 auto *N = DAG.getNode(I: &I);
392 if (N == nullptr)
393 continue;
394 auto *SB = N->getSchedBundle();
395 if (SB->isSingleton())
396 eraseBundle(SB);
397 N->resetScheduleState();
398 }
399 // Nodes that depend on the nodes in ResetIntvl also need to have their
400 // UnscheduledSuccs/UnscheduledPreds adjusted.
401 for (Instruction &I : ResetIntvl) {
402 auto *N = DAG.getNode(I: &I);
403 if (Dir == SchedDirection::BottomUp) {
404 // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even
405 // for nodes above the top of schedule.
406 for (auto *PredN : N->preds(DAG))
407 PredN->incrUnscheduledDeps();
408 } else {
409 assert(Dir == SchedDirection::TopDown);
410 // Recompute UnscheduledPreds for nodes not only in ResetIntvl but even
411 // for nodes below the bottom of schedule.
412 for (auto *SuccN : N->succs(DAG))
413 SuccN->incrUnscheduledDeps();
414 }
415 }
416
417 // Refill the ready list by visiting all the nodes in the unscheduled part of
418 // the DAG. In bottom-up that is from the top of the DAG down to LowestI; in
419 // top-down it is the mirror image, from TopI down to the bottom of the DAG.
420 ReadyList.clear();
421 Interval<Instruction> RefillIntvl =
422 Dir == SchedDirection::BottomUp
423 ? Interval<Instruction>(DAG.getInterval().top(), LowestI)
424 : Interval<Instruction>(TopI, DAG.getInterval().bottom());
425 for (Instruction &I : RefillIntvl) {
426 auto *N = DAG.getNode(I: &I);
427 if (N->ready())
428 ReadyList.insert(N);
429 }
430}
431
432#ifndef NDEBUG
433void Scheduler::assertSameDirection(ArrayRef<Instruction *> Instrs) const {
434 // Check that we are not switching scheduling direction.
435 switch (Dir) {
436 case SchedDirection::BottomUp:
437 assert(none_of(Instrs,
438 [this](Instruction *I) {
439 return ScheduleFrontierOpt->comesBefore(*I);
440 }) &&
441 "Wrong scheduling direction!");
442 break;
443 case SchedDirection::TopDown:
444 assert(all_of(Instrs,
445 [this](Instruction *I) {
446 return ScheduleFrontierOpt->comesBefore(*I);
447 }) &&
448 "Wrong scheduling direction!");
449 break;
450 }
451}
452#endif // NDEBUG
453
454bool Scheduler::trySchedule(BndlRef<Instruction *> Instrs) {
455 assert(all_of(drop_begin(Instrs),
456 [Instrs](Instruction *I) {
457 return I->getParent() == (*Instrs.begin())->getParent();
458 }) &&
459 "Instrs not in the same BB, should have been rejected by Legality!");
460 // TODO: For now don't cross BBs.
461 if (!DAG.getInterval().empty()) {
462 auto *BB = DAG.getInterval().top()->getParent();
463 if (any_of(Range&: Instrs, P: [BB](auto *I) { return I->getParent() != BB; }))
464 return false;
465 }
466 if (ScheduledBB == nullptr)
467 ScheduledBB = Instrs[0]->getParent();
468 // We don't support crossing BBs for now.
469 if (any_of(Range&: Instrs,
470 P: [this](Instruction *I) { return I->getParent() != ScheduledBB; }))
471 return false;
472
473 auto GetSchedPoint = [](SchedDirection Dir,
474 const auto &Instrs) -> SchedulingPoint {
475 switch (Dir) {
476 case SchedDirection::BottomUp:
477 return SchedulingPoint(VecUtils::getLowest(Instrs)->getIterator())
478 .getNext();
479 case SchedDirection::TopDown:
480 return SchedulingPoint(VecUtils::getHighest(Instrs)->getIterator())
481 .getPrev();
482 }
483 llvm_unreachable("Unhandled Dir!");
484 };
485 auto SchedState = getBndlSchedState(Instrs);
486 switch (SchedState) {
487 case BndlSchedState::FullyScheduled:
488 // Nothing to do.
489 return true;
490 case BndlSchedState::AlreadyScheduled:
491 // Instructions are part of a different vector schedule, so we can't
492 // schedule \p Instrs in the same bundle (without destroying the existing
493 // schedule).
494 return false;
495 case BndlSchedState::TemporarilyScheduled:
496 // If one or more instrs are already scheduled we need to destroy the
497 // top-most part of the schedule that includes the instrs in the bundle and
498 // re-schedule.
499 DAG.extend(Instrs);
500 trimSchedule(Instrs);
501 ScheduleFrontierOpt = GetSchedPoint(Dir, Instrs);
502 return tryScheduleUntil(Instrs);
503 case BndlSchedState::NoneScheduled: {
504 // TODO: Set the window of the DAG that we are interested in.
505 if (!ScheduleFrontierOpt) {
506 // We start scheduling at the bottom instr of Instrs (top in TopDown).
507 ScheduleFrontierOpt = GetSchedPoint(Dir, Instrs);
508 } else {
509#ifndef NDEBUG
510 assertSameDirection(Instrs);
511#endif
512 }
513 // Extend the DAG to include Instrs.
514 Interval<Instruction> Extension = DAG.extend(Instrs);
515 // Add nodes from the new interval to ready list if they are ready.
516 for (auto &I : Extension) {
517 auto *N = DAG.getNode(I: &I);
518 if (N->scheduled())
519 continue;
520 if (N->ready() && !ReadyList.contains(N))
521 ReadyList.insert(N);
522 }
523 // Try schedule all nodes until we can schedule Instrs back-to-back.
524 return tryScheduleUntil(Instrs);
525 }
526 }
527 llvm_unreachable("Unhandled BndlSchedState enum");
528}
529
530#ifndef NDEBUG
531void Scheduler::dump(raw_ostream &OS) const {
532 OS << "ReadyList:\n";
533 ReadyList.dump(OS);
534 OS << "Dir=" << schedDirectionToStr(Dir) << " "
535 << (Dir == SchedDirection::BottomUp ? "Top" : "Bottom")
536 << " of schedule: ";
537 if (ScheduleFrontierOpt)
538 OS << **ScheduleFrontierOpt;
539 else
540 OS << "Empty";
541 OS << "\n";
542}
543void Scheduler::dump() const { dump(dbgs()); }
544#endif // NDEBUG
545
546} // namespace llvm::sandboxir
547