361 lines · cpp
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-exception6//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(Nodes)) {18 if (N->getInstruction()->comesBefore(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(Nodes)) {27 if (BotN->getInstruction()->comesBefore(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(*Where.getNodeParent(), Where);39 }40}41 42#ifndef NDEBUG43void 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 // NDEBUG53 54#ifndef NDEBUG55void 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 // NDEBUG68 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 all78 // 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(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 the93 // 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 as97 // "scheduled".98 bool IsScheduled = ScheduleTopItOpt &&99 *ScheduleTopItOpt != I->getParent()->end() &&100 (*ScheduleTopItOpt.value()).comesBefore(I);101 if (IsScheduled)102 N->setScheduled(true);103 // If the new instruction is above the top of schedule we need to remove its104 // dependency predecessors from the ready list and increment their105 // `UnscheduledSuccs` counters.106 if (!IsScheduled) {107 for (auto *PredN : N->preds(DAG)) {108 ReadyList.remove(PredN);109 PredN->incrUnscheduledSuccs();110 }111 }112}113 114SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {115 SchedBundle::ContainerTy Nodes;116 Nodes.reserve(Instrs.size());117 for (auto *I : Instrs)118 Nodes.push_back(DAG.getNode(I));119 auto BndlPtr = std::make_unique<SchedBundle>(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(SB); }126 127bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {128 // Create a bundle for Instrs. If it turns out the schedule is infeasible we129 // 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 of133 // 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 last141 /// bundle to be scheduled.142 };143 /// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that144 /// ReadyN is in. If it's not in a bundle it will create a singleton bundle145 /// 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 bundle150 // and schedule it.151 auto *SingletonSB = createBundle({ReadyN->getInstruction()});152 scheduleAndUpdateReadyList(*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 on158 // 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(*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 scheduling181 // other ready instructions.182 Retry.push_back(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(Retry)) {193 auto Res = TryScheduleBndl(N);194 if (Res == TryScheduleRes::Success) {195 Retry.erase(find(Retry, N));196 KeepScheduling = true;197 }198 }199 }200 201 eraseBundle(InstrsSB);202 return false;203}204 205Scheduler::BndlSchedState206Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {207 assert(!Instrs.empty() && "Expected non-empty bundle");208 auto *N0 = DAG.getNode(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(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. There220 // 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, so227 // 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::NoneScheduled236 : FullyScheduled ? BndlSchedState::FullyScheduled237 : BndlSchedState::TemporarilyScheduled;238}239 240void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {241 // | Legend: N: DGNode242 // N <- DAGInterval.top() | B: SchedBundle243 // N | *: Contains instruction in Instrs244 // B <- TopI (Top of schedule) +-------------------------------------------245 // B246 // B *247 // B248 // B * <- LowestI (Lowest in Instrs)249 // B250 // N251 // N252 // 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 the267 // Scheduled flag. We need to reset their state. We need to do this for all268 // nodes from LowestI to the top of the schedule. DAG Nodes that are above the269 // top of schedule that depend on nodes that got reset need to have their270 // UnscheduledSuccs adjusted.271 Interval<Instruction> ResetIntvl(TopI, LowestI);272 for (Instruction &I : ResetIntvl) {273 auto *N = DAG.getNode(&I);274 N->resetScheduleState();275 // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even for276 // 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);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(Instrs, [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(Instrs,306 [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't315 // schedule \p Instrs in the same bundle (without destroying the existing316 // schedule).317 return false;318 case BndlSchedState::TemporarilyScheduled:319 // If one or more instrs are already scheduled we need to destroy the320 // top-most part of the schedule that includes the instrs in the bundle and321 // re-schedule.322 DAG.extend(Instrs);323 trimSchedule(Instrs);324 ScheduleTopItOpt = std::next(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(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);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 NDEBUG347void 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 else354 OS << "Empty";355 OS << "\n";356}357void Scheduler::dump() const { dump(dbgs()); }358#endif // NDEBUG359 360} // namespace llvm::sandboxir361