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(ScheduleTopItOpt && "Should have been set by now!");
86 auto Where = Dir == SchedDirection::BottomUp
87 ? ScheduleTopItOpt->getIterator()
88 : ScheduleTopItOpt->getNext().getIterator();
89 // Move all instructions in `Bndl` to `Where`.
90 Bndl.cluster(Where);
91 // Update the last scheduled bundle.
92 ScheduleTopItOpt = 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 = ScheduleTopItOpt &&
130 ScheduleTopItOpt->getIterator() != I->getParent()->end() &&
131 ((Dir == SchedDirection::BottomUp &&
132 (*ScheduleTopItOpt.value()).comesBefore(Other: I)) ||
133 (Dir == SchedDirection::TopDown &&
134 I->comesBefore(Other: &*ScheduleTopItOpt.value())));
135 if (IsScheduled)
136 N->setScheduled();
137 // If the new instruction is above the top of schedule we need to remove its
138 // dependency predecessors from the ready list and increment their
139 // `UnscheduledSuccs` counters.
140 if (!IsScheduled) {
141 if (Dir == SchedDirection::BottomUp) {
142 for (auto *PredN : N->preds(DAG)) {
143 ReadyList.remove(N: PredN);
144 PredN->incrUnscheduledDeps();
145 }
146 } else {
147 for (auto *SuccN : N->succs(DAG)) {
148 ReadyList.remove(N: SuccN);
149 SuccN->incrUnscheduledDeps();
150 }
151 }
152 }
153}
154
155SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
156 SchedBundle::ContainerTy Nodes;
157 Nodes.reserve(N: Instrs.size());
158 for (auto *I : Instrs)
159 Nodes.push_back(Elt: DAG.getNode(I));
160 auto BndlPtr = std::make_unique<SchedBundle>(args: std::move(Nodes));
161 auto *Bndl = BndlPtr.get();
162 Bndls[Bndl] = std::move(BndlPtr);
163 return Bndl;
164}
165
166void Scheduler::eraseBundle(SchedBundle *SB) { Bndls.erase(Val: SB); }
167
168bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
169 // Create a bundle for Instrs. If it turns out the schedule is infeasible we
170 // will dismantle it.
171 auto *InstrsSB = createBundle(Instrs);
172 // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
173 // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
174 // which are collected in DeferredNodes) are all ready to schedule.
175 SmallVector<DGNode *> Retry;
176 bool KeepScheduling = true;
177 while (KeepScheduling) {
178 enum class TryScheduleRes {
179 Success, ///> We successfully scheduled the bundle.
180 Failure, ///> We failed to schedule the bundle.
181 Finished, ///> We successfully scheduled the bundle and it is the last
182 /// bundle to be scheduled.
183 };
184 /// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that
185 /// ReadyN is in. If it's not in a bundle it will create a singleton bundle
186 /// and will try to schedule it.
187 auto TryScheduleBndl = [this, InstrsSB](DGNode *ReadyN) -> TryScheduleRes {
188 auto *SB = ReadyN->getSchedBundle();
189 if (SB == nullptr) {
190 // If ReadyN does not belong to a bundle, create a singleton bundle
191 // and schedule it.
192 auto *SingletonSB = createBundle(Instrs: {ReadyN->getInstruction()});
193 scheduleAndUpdateReadyList(Bndl&: *SingletonSB);
194 return TryScheduleRes::Success;
195 }
196 if (SB->ready(Dir)) {
197 // Remove the rest of the bundle from the ready list.
198 // TODO: Perhaps change the Scheduler + ReadyList to operate on
199 // SchedBundles instead of DGNodes.
200 for (auto *N : *SB) {
201 if (N != ReadyN)
202 ReadyList.remove(N);
203 }
204 // If all nodes in the bundle are ready.
205 scheduleAndUpdateReadyList(Bndl&: *SB);
206 if (SB == InstrsSB)
207 // We just scheduled InstrsSB bundle, so we are done scheduling.
208 return TryScheduleRes::Finished;
209 return TryScheduleRes::Success;
210 }
211 return TryScheduleRes::Failure;
212 };
213 while (!ReadyList.empty()) {
214 auto *ReadyN = ReadyList.pop();
215 auto Res = TryScheduleBndl(ReadyN);
216 switch (Res) {
217 case TryScheduleRes::Success:
218 // We successfully scheduled ReadyN, keep scheduling.
219 continue;
220 case TryScheduleRes::Failure:
221 // We failed to schedule ReadyN, defer it to later and keep scheduling
222 // other ready instructions.
223 Retry.push_back(Elt: ReadyN);
224 continue;
225 case TryScheduleRes::Finished:
226 // We successfully scheduled the instruction bundle, so we are done.
227 return true;
228 }
229 llvm_unreachable("Unhandled TrySchedule() result");
230 }
231 // Try to schedule nodes from the Retry list.
232 KeepScheduling = false;
233 for (auto *N : make_early_inc_range(Range&: Retry)) {
234 auto Res = TryScheduleBndl(N);
235 if (Res == TryScheduleRes::Success) {
236 Retry.erase(CI: find(Range&: Retry, Val: N));
237 KeepScheduling = true;
238 }
239 }
240 }
241
242 eraseBundle(SB: InstrsSB);
243 return false;
244}
245
246Scheduler::BndlSchedState
247Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
248 assert(!Instrs.empty() && "Expected non-empty bundle");
249 auto *N0 = DAG.getNode(I: Instrs[0]);
250 auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
251 bool AllUnscheduled = SB0 == nullptr;
252 bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
253 for (auto *I : drop_begin(RangeOrContainer&: Instrs)) {
254 auto *N = DAG.getNode(I);
255 auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
256 if (SB != nullptr) {
257 // We found a scheduled instr, so there is now way all are unscheduled.
258 AllUnscheduled = false;
259 if (SB->isSingleton()) {
260 // We found an instruction in a temporarily scheduled singleton. There
261 // is no way that all instructions are scheduled in the same bundle.
262 FullyScheduled = false;
263 }
264 }
265
266 if (SB != SB0) {
267 // Either one of SB, SB0 is null, or they are in different bundles, so
268 // Instrs are definitely not in the same vector bundle.
269 FullyScheduled = false;
270 // One of SB, SB0 are in a vector bundle and they differ.
271 if ((SB != nullptr && !SB->isSingleton()) ||
272 (SB0 != nullptr && !SB0->isSingleton()))
273 return BndlSchedState::AlreadyScheduled;
274 }
275 }
276 return AllUnscheduled ? BndlSchedState::NoneScheduled
277 : FullyScheduled ? BndlSchedState::FullyScheduled
278 : BndlSchedState::TemporarilyScheduled;
279}
280
281void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
282 // | Legend: N: DGNode
283 // N <- DAGInterval.top() | B: SchedBundle
284 // N | *: Contains instruction in Instrs
285 // B <- TopI (Top of schedule) +-------------------------------------------
286 // B
287 // B *
288 // B
289 // B * <- LowestI (Lowest in Instrs)
290 // B
291 // N
292 // N
293 // N <- DAGInterval.bottom()
294 //
295 // Note: this figure assumes bottom-up scheduling. In top-down we have the
296 // top-down mirror image.
297 Instruction *TopI = Dir == SchedDirection::BottomUp
298 ? &*ScheduleTopItOpt.value()
299 : VecUtils::getHighest(Instrs);
300 Instruction *LowestI = Dir == SchedDirection::BottomUp
301 ? VecUtils::getLowest(Instrs)
302 : &*ScheduleTopItOpt.value();
303 Interval<Instruction> ResetIntvl(TopI, LowestI);
304 // The DAG Nodes contain state like the number of UnscheduledSuccs and the
305 // Scheduled flag. We need to reset their state. We need to do this for all
306 // nodes in ResetIntvl. Also destroy the singleton schedule bundles from
307 // LowestI all the way to the top.
308 for (auto &I : ResetIntvl) {
309 auto *N = DAG.getNode(I: &I);
310 if (N == nullptr)
311 continue;
312 auto *SB = N->getSchedBundle();
313 if (SB->isSingleton())
314 eraseBundle(SB);
315 N->resetScheduleState();
316 }
317 // Nodes that depend on the nodes in ResetIntvl also need to have their
318 // UnscheduledSuccs/UnscheduledPreds adjusted.
319 for (Instruction &I : ResetIntvl) {
320 auto *N = DAG.getNode(I: &I);
321 if (Dir == SchedDirection::BottomUp) {
322 // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even
323 // for nodes above the top of schedule.
324 for (auto *PredN : N->preds(DAG))
325 PredN->incrUnscheduledDeps();
326 } else {
327 assert(Dir == SchedDirection::TopDown);
328 // Recompute UnscheduledPreds for nodes not only in ResetIntvl but even
329 // for nodes below the bottom of schedule.
330 for (auto *SuccN : N->succs(DAG))
331 SuccN->incrUnscheduledDeps();
332 }
333 }
334
335 // Refill the ready list by visiting all the nodes in the unscheduled part of
336 // the DAG. In bottom-up that is from the top of the DAG down to LowestI; in
337 // top-down it is the mirror image, from TopI down to the bottom of the DAG.
338 ReadyList.clear();
339 Interval<Instruction> RefillIntvl =
340 Dir == SchedDirection::BottomUp
341 ? Interval<Instruction>(DAG.getInterval().top(), LowestI)
342 : Interval<Instruction>(TopI, DAG.getInterval().bottom());
343 for (Instruction &I : RefillIntvl) {
344 auto *N = DAG.getNode(I: &I);
345 if (N->ready())
346 ReadyList.insert(N);
347 }
348}
349
350bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
351 assert(all_of(drop_begin(Instrs),
352 [Instrs](Instruction *I) {
353 return I->getParent() == (*Instrs.begin())->getParent();
354 }) &&
355 "Instrs not in the same BB, should have been rejected by Legality!");
356 // TODO: For now don't cross BBs.
357 if (!DAG.getInterval().empty()) {
358 auto *BB = DAG.getInterval().top()->getParent();
359 if (any_of(Range&: Instrs, P: [BB](auto *I) { return I->getParent() != BB; }))
360 return false;
361 }
362 if (ScheduledBB == nullptr)
363 ScheduledBB = Instrs[0]->getParent();
364 // We don't support crossing BBs for now.
365 if (any_of(Range&: Instrs,
366 P: [this](Instruction *I) { return I->getParent() != ScheduledBB; }))
367 return false;
368
369 auto GetSchedPoint = [](SchedDirection Dir,
370 const auto &Instrs) -> SchedulingPoint {
371 switch (Dir) {
372 case SchedDirection::BottomUp:
373 return SchedulingPoint(VecUtils::getLowest(Instrs)->getIterator())
374 .getNext();
375 case SchedDirection::TopDown:
376 return SchedulingPoint(VecUtils::getHighest(Instrs)->getIterator())
377 .getPrev();
378 }
379 llvm_unreachable("Unhandled Dir!");
380 };
381 auto SchedState = getBndlSchedState(Instrs);
382 switch (SchedState) {
383 case BndlSchedState::FullyScheduled:
384 // Nothing to do.
385 return true;
386 case BndlSchedState::AlreadyScheduled:
387 // Instructions are part of a different vector schedule, so we can't
388 // schedule \p Instrs in the same bundle (without destroying the existing
389 // schedule).
390 return false;
391 case BndlSchedState::TemporarilyScheduled:
392 // If one or more instrs are already scheduled we need to destroy the
393 // top-most part of the schedule that includes the instrs in the bundle and
394 // re-schedule.
395 DAG.extend(Instrs);
396 trimSchedule(Instrs);
397 ScheduleTopItOpt = GetSchedPoint(Dir, Instrs);
398 return tryScheduleUntil(Instrs);
399 case BndlSchedState::NoneScheduled: {
400 // TODO: Set the window of the DAG that we are interested in.
401 if (!ScheduleTopItOpt)
402 // We start scheduling at the bottom instr of Instrs (top in TopDown).
403 ScheduleTopItOpt = GetSchedPoint(Dir, Instrs);
404 // Extend the DAG to include Instrs.
405 Interval<Instruction> Extension = DAG.extend(Instrs);
406 // Add nodes from the new interval to ready list if they are ready.
407 Interval<Instruction> InstrsInterval(Instrs);
408 Interval<Instruction> ScanForReady =
409 InstrsInterval.getUnionInterval(Other: Extension);
410 for (auto &I : ScanForReady) {
411 auto *N = DAG.getNode(I: &I);
412 if (N->scheduled())
413 continue;
414 if (N->ready() && !ReadyList.contains(N))
415 ReadyList.insert(N);
416 }
417 // Try schedule all nodes until we can schedule Instrs back-to-back.
418 return tryScheduleUntil(Instrs);
419 }
420 }
421 llvm_unreachable("Unhandled BndlSchedState enum");
422}
423
424#ifndef NDEBUG
425void Scheduler::dump(raw_ostream &OS) const {
426 OS << "ReadyList:\n";
427 ReadyList.dump(OS);
428 OS << "Dir=" << schedDirectionToStr(Dir) << " "
429 << (Dir == SchedDirection::BottomUp ? "Top" : "Bottom")
430 << " of schedule: ";
431 if (ScheduleTopItOpt)
432 OS << **ScheduleTopItOpt;
433 else
434 OS << "Empty";
435 OS << "\n";
436}
437void Scheduler::dump() const { dump(dbgs()); }
438#endif // NDEBUG
439
440} // namespace llvm::sandboxir
441