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