Replace the for loop that iterates between TopI and LowestI with a range loop over ResetIntvl and move the reset of the scheduling state to this loop. We iterate over the loop in the reverse order than before but the functionality should not change.
426 lines
15 KiB
C++
426 lines
15 KiB
C++
//===- Scheduler.cpp ------------------------------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
|
|
#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"
|
|
|
|
namespace llvm::sandboxir {
|
|
|
|
#ifndef NDEBUG
|
|
StringLiteral schedDirectionToStr(SchedDirection Dir) {
|
|
switch (Dir) {
|
|
case SchedDirection::BottomUp:
|
|
return "BottomUp";
|
|
case SchedDirection::TopDown:
|
|
return "TopDown";
|
|
}
|
|
llvm_unreachable("Unhandled Dir!");
|
|
}
|
|
#endif // NDEBUG
|
|
|
|
// TODO: Check if we can cache top/bottom to reduce compile-time.
|
|
DGNode *SchedBundle::getTop() const {
|
|
DGNode *TopN = Nodes.front();
|
|
for (auto *N : drop_begin(Nodes)) {
|
|
if (N->getInstruction()->comesBefore(TopN->getInstruction()))
|
|
TopN = N;
|
|
}
|
|
return TopN;
|
|
}
|
|
|
|
DGNode *SchedBundle::getBot() const {
|
|
DGNode *BotN = Nodes.front();
|
|
for (auto *N : drop_begin(Nodes)) {
|
|
if (BotN->getInstruction()->comesBefore(N->getInstruction()))
|
|
BotN = N;
|
|
}
|
|
return BotN;
|
|
}
|
|
|
|
void SchedBundle::cluster(BasicBlock::iterator Where) {
|
|
for (auto *N : Nodes) {
|
|
auto *I = N->getInstruction();
|
|
if (I->getIterator() == Where)
|
|
++Where; // Try to maintain bundle order.
|
|
I->moveBefore(*Where.getNodeParent(), Where);
|
|
}
|
|
}
|
|
|
|
#ifndef NDEBUG
|
|
void SchedBundle::dump(raw_ostream &OS) const {
|
|
for (auto *N : Nodes)
|
|
OS << *N;
|
|
}
|
|
|
|
void SchedBundle::dump() const {
|
|
dump(dbgs());
|
|
dbgs() << "\n";
|
|
}
|
|
#endif // NDEBUG
|
|
|
|
#ifndef NDEBUG
|
|
void ReadyListContainer::dump(raw_ostream &OS) const {
|
|
auto ListCopy = List;
|
|
while (!ListCopy.empty()) {
|
|
OS << *ListCopy.top() << "\n";
|
|
ListCopy.pop();
|
|
}
|
|
}
|
|
|
|
void ReadyListContainer::dump() const {
|
|
dump(dbgs());
|
|
dbgs() << "\n";
|
|
}
|
|
#endif // NDEBUG
|
|
|
|
void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
|
|
// Find where we should schedule the instructions.
|
|
assert(ScheduleTopItOpt && "Should have been set by now!");
|
|
auto Where = Dir == SchedDirection::BottomUp ? *ScheduleTopItOpt
|
|
: std::next(*ScheduleTopItOpt);
|
|
// Move all instructions in `Bndl` to `Where`.
|
|
Bndl.cluster(Where);
|
|
// Update the last scheduled bundle.
|
|
ScheduleTopItOpt = Dir == SchedDirection::BottomUp
|
|
? Bndl.getTop()->getInstruction()->getIterator()
|
|
: Bndl.getBot()->getInstruction()->getIterator();
|
|
// Set nodes as "scheduled" and decrement the UnscheduledSuccs/Preds counter
|
|
// of all dependency predecessors/successors.
|
|
for (DGNode *N : Bndl) {
|
|
switch (Dir) {
|
|
case SchedDirection::BottomUp: {
|
|
for (auto *DepN : N->preds(DAG)) {
|
|
DepN->decrUnscheduledSuccs();
|
|
if (DepN->readyBottomUp() && !DepN->scheduled())
|
|
ReadyList.insert(DepN);
|
|
}
|
|
break;
|
|
}
|
|
case SchedDirection::TopDown: {
|
|
for (auto *DepN : N->succs(DAG)) {
|
|
DepN->decrUnscheduledPreds();
|
|
if (DepN->readyTopDown() && !DepN->scheduled())
|
|
ReadyList.insert(DepN);
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
N->setScheduled();
|
|
}
|
|
}
|
|
|
|
void Scheduler::notifyCreateInstr(Instruction *I) {
|
|
// The DAG notifier should have run by now.
|
|
auto *N = DAG.getNode(I);
|
|
// If there is no DAG node for `I` it means that this is out of scope for the
|
|
// DAG and as such out of scope for the scheduler too, so nothing to do.
|
|
if (N == nullptr)
|
|
return;
|
|
// If the instruction is inserted below the top-of-schedule then we mark it as
|
|
// "scheduled".
|
|
bool IsScheduled = ScheduleTopItOpt &&
|
|
*ScheduleTopItOpt != I->getParent()->end() &&
|
|
((Dir == SchedDirection::BottomUp &&
|
|
(*ScheduleTopItOpt.value()).comesBefore(I)) ||
|
|
(Dir == SchedDirection::TopDown &&
|
|
I->comesBefore(&*ScheduleTopItOpt.value())));
|
|
if (IsScheduled)
|
|
N->setScheduled();
|
|
// If the new instruction is above the top of schedule we need to remove its
|
|
// dependency predecessors from the ready list and increment their
|
|
// `UnscheduledSuccs` counters.
|
|
if (!IsScheduled) {
|
|
if (Dir == SchedDirection::BottomUp) {
|
|
for (auto *PredN : N->preds(DAG)) {
|
|
ReadyList.remove(PredN);
|
|
PredN->incrUnscheduledSuccs();
|
|
}
|
|
} else {
|
|
for (auto *SuccN : N->succs(DAG)) {
|
|
ReadyList.remove(SuccN);
|
|
SuccN->incrUnscheduledPreds();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
|
|
SchedBundle::ContainerTy Nodes;
|
|
Nodes.reserve(Instrs.size());
|
|
for (auto *I : Instrs)
|
|
Nodes.push_back(DAG.getNode(I));
|
|
auto BndlPtr = std::make_unique<SchedBundle>(std::move(Nodes));
|
|
auto *Bndl = BndlPtr.get();
|
|
Bndls[Bndl] = std::move(BndlPtr);
|
|
return Bndl;
|
|
}
|
|
|
|
void Scheduler::eraseBundle(SchedBundle *SB) { Bndls.erase(SB); }
|
|
|
|
bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
|
|
// Create a bundle for Instrs. If it turns out the schedule is infeasible we
|
|
// will dismantle it.
|
|
auto *InstrsSB = createBundle(Instrs);
|
|
// Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
|
|
// ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
|
|
// which are collected in DeferredNodes) are all ready to schedule.
|
|
SmallVector<DGNode *> Retry;
|
|
bool KeepScheduling = true;
|
|
while (KeepScheduling) {
|
|
enum class TryScheduleRes {
|
|
Success, ///> We successfully scheduled the bundle.
|
|
Failure, ///> We failed to schedule the bundle.
|
|
Finished, ///> We successfully scheduled the bundle and it is the last
|
|
/// bundle to be scheduled.
|
|
};
|
|
/// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that
|
|
/// ReadyN is in. If it's not in a bundle it will create a singleton bundle
|
|
/// and will try to schedule it.
|
|
auto TryScheduleBndl = [this, InstrsSB](DGNode *ReadyN) -> TryScheduleRes {
|
|
auto *SB = ReadyN->getSchedBundle();
|
|
if (SB == nullptr) {
|
|
// If ReadyN does not belong to a bundle, create a singleton bundle
|
|
// and schedule it.
|
|
auto *SingletonSB = createBundle({ReadyN->getInstruction()});
|
|
scheduleAndUpdateReadyList(*SingletonSB);
|
|
return TryScheduleRes::Success;
|
|
}
|
|
if (SB->ready(Dir)) {
|
|
// Remove the rest of the bundle from the ready list.
|
|
// TODO: Perhaps change the Scheduler + ReadyList to operate on
|
|
// SchedBundles instead of DGNodes.
|
|
for (auto *N : *SB) {
|
|
if (N != ReadyN)
|
|
ReadyList.remove(N);
|
|
}
|
|
// If all nodes in the bundle are ready.
|
|
scheduleAndUpdateReadyList(*SB);
|
|
if (SB == InstrsSB)
|
|
// We just scheduled InstrsSB bundle, so we are done scheduling.
|
|
return TryScheduleRes::Finished;
|
|
return TryScheduleRes::Success;
|
|
}
|
|
return TryScheduleRes::Failure;
|
|
};
|
|
while (!ReadyList.empty()) {
|
|
auto *ReadyN = ReadyList.pop();
|
|
auto Res = TryScheduleBndl(ReadyN);
|
|
switch (Res) {
|
|
case TryScheduleRes::Success:
|
|
// We successfully scheduled ReadyN, keep scheduling.
|
|
continue;
|
|
case TryScheduleRes::Failure:
|
|
// We failed to schedule ReadyN, defer it to later and keep scheduling
|
|
// other ready instructions.
|
|
Retry.push_back(ReadyN);
|
|
continue;
|
|
case TryScheduleRes::Finished:
|
|
// We successfully scheduled the instruction bundle, so we are done.
|
|
return true;
|
|
}
|
|
llvm_unreachable("Unhandled TrySchedule() result");
|
|
}
|
|
// Try to schedule nodes from the Retry list.
|
|
KeepScheduling = false;
|
|
for (auto *N : make_early_inc_range(Retry)) {
|
|
auto Res = TryScheduleBndl(N);
|
|
if (Res == TryScheduleRes::Success) {
|
|
Retry.erase(find(Retry, N));
|
|
KeepScheduling = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
eraseBundle(InstrsSB);
|
|
return false;
|
|
}
|
|
|
|
Scheduler::BndlSchedState
|
|
Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
|
|
assert(!Instrs.empty() && "Expected non-empty bundle");
|
|
auto *N0 = DAG.getNode(Instrs[0]);
|
|
auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
|
|
bool AllUnscheduled = SB0 == nullptr;
|
|
bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
|
|
for (auto *I : drop_begin(Instrs)) {
|
|
auto *N = DAG.getNode(I);
|
|
auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
|
|
if (SB != nullptr) {
|
|
// We found a scheduled instr, so there is now way all are unscheduled.
|
|
AllUnscheduled = false;
|
|
if (SB->isSingleton()) {
|
|
// We found an instruction in a temporarily scheduled singleton. There
|
|
// is no way that all instructions are scheduled in the same bundle.
|
|
FullyScheduled = false;
|
|
}
|
|
}
|
|
|
|
if (SB != SB0) {
|
|
// Either one of SB, SB0 is null, or they are in different bundles, so
|
|
// Instrs are definitely not in the same vector bundle.
|
|
FullyScheduled = false;
|
|
// One of SB, SB0 are in a vector bundle and they differ.
|
|
if ((SB != nullptr && !SB->isSingleton()) ||
|
|
(SB0 != nullptr && !SB0->isSingleton()))
|
|
return BndlSchedState::AlreadyScheduled;
|
|
}
|
|
}
|
|
return AllUnscheduled ? BndlSchedState::NoneScheduled
|
|
: FullyScheduled ? BndlSchedState::FullyScheduled
|
|
: BndlSchedState::TemporarilyScheduled;
|
|
}
|
|
|
|
void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
|
|
// | Legend: N: DGNode
|
|
// N <- DAGInterval.top() | B: SchedBundle
|
|
// N | *: Contains instruction in Instrs
|
|
// B <- TopI (Top of schedule) +-------------------------------------------
|
|
// B
|
|
// B *
|
|
// B
|
|
// B * <- LowestI (Lowest in Instrs)
|
|
// B
|
|
// N
|
|
// N
|
|
// N <- DAGInterval.bottom()
|
|
//
|
|
// Note: this figure assumes bottom-up scheduling. In top-down we have the
|
|
// top-down mirror image.
|
|
Instruction *TopI = Dir == SchedDirection::BottomUp
|
|
? &*ScheduleTopItOpt.value()
|
|
: VecUtils::getHighest(Instrs);
|
|
Instruction *LowestI = Dir == SchedDirection::BottomUp
|
|
? VecUtils::getLowest(Instrs)
|
|
: &*ScheduleTopItOpt.value();
|
|
Interval<Instruction> ResetIntvl(TopI, LowestI);
|
|
// The DAG Nodes contain state like the number of UnscheduledSuccs and the
|
|
// Scheduled flag. We need to reset their state. We need to do this for all
|
|
// nodes in ResetIntvl. Also destroy the singleton schedule bundles from
|
|
// LowestI all the way to the top.
|
|
for (auto &I : ResetIntvl) {
|
|
auto *N = DAG.getNode(&I);
|
|
if (N == nullptr)
|
|
continue;
|
|
auto *SB = N->getSchedBundle();
|
|
if (SB->isSingleton())
|
|
eraseBundle(SB);
|
|
N->resetScheduleState();
|
|
}
|
|
// Nodes that depend on the nodes in ResetIntvl also need to have their
|
|
// UnscheduledSuccs/UnscheduledPreds adjusted.
|
|
for (Instruction &I : ResetIntvl) {
|
|
auto *N = DAG.getNode(&I);
|
|
if (Dir == SchedDirection::BottomUp) {
|
|
// Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even
|
|
// for nodes above the top of schedule.
|
|
for (auto *PredN : N->preds(DAG))
|
|
PredN->incrUnscheduledSuccs();
|
|
} else {
|
|
assert(Dir == SchedDirection::TopDown);
|
|
// Recompute UnscheduledPreds for nodes not only in ResetIntvl but even
|
|
// for nodes below the bottom of schedule.
|
|
for (auto *SuccN : N->succs(DAG))
|
|
SuccN->incrUnscheduledPreds();
|
|
}
|
|
}
|
|
|
|
// Refill the ready list by visiting all nodes from the top of DAG to LowestI.
|
|
ReadyList.clear();
|
|
Interval<Instruction> RefillIntvl(DAG.getInterval().top(), LowestI);
|
|
for (Instruction &I : RefillIntvl) {
|
|
auto *N = DAG.getNode(&I);
|
|
if (N->readyBottomUp())
|
|
ReadyList.insert(N);
|
|
}
|
|
}
|
|
|
|
bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
|
|
assert(all_of(drop_begin(Instrs),
|
|
[Instrs](Instruction *I) {
|
|
return I->getParent() == (*Instrs.begin())->getParent();
|
|
}) &&
|
|
"Instrs not in the same BB, should have been rejected by Legality!");
|
|
// TODO: For now don't cross BBs.
|
|
if (!DAG.getInterval().empty()) {
|
|
auto *BB = DAG.getInterval().top()->getParent();
|
|
if (any_of(Instrs, [BB](auto *I) { return I->getParent() != BB; }))
|
|
return false;
|
|
}
|
|
if (ScheduledBB == nullptr)
|
|
ScheduledBB = Instrs[0]->getParent();
|
|
// We don't support crossing BBs for now.
|
|
if (any_of(Instrs,
|
|
[this](Instruction *I) { return I->getParent() != ScheduledBB; }))
|
|
return false;
|
|
|
|
auto GetSchedPoint = [](SchedDirection Dir, const auto &Instrs) {
|
|
switch (Dir) {
|
|
case SchedDirection::BottomUp:
|
|
return std::next(VecUtils::getLowest(Instrs)->getIterator());
|
|
case SchedDirection::TopDown:
|
|
return VecUtils::getHighest(Instrs)->getIterator();
|
|
}
|
|
llvm_unreachable("Unhandled Dir!");
|
|
};
|
|
auto SchedState = getBndlSchedState(Instrs);
|
|
switch (SchedState) {
|
|
case BndlSchedState::FullyScheduled:
|
|
// Nothing to do.
|
|
return true;
|
|
case BndlSchedState::AlreadyScheduled:
|
|
// Instructions are part of a different vector schedule, so we can't
|
|
// schedule \p Instrs in the same bundle (without destroying the existing
|
|
// schedule).
|
|
return false;
|
|
case BndlSchedState::TemporarilyScheduled:
|
|
// If one or more instrs are already scheduled we need to destroy the
|
|
// top-most part of the schedule that includes the instrs in the bundle and
|
|
// re-schedule.
|
|
DAG.extend(Instrs);
|
|
trimSchedule(Instrs);
|
|
ScheduleTopItOpt = GetSchedPoint(Dir, Instrs);
|
|
return tryScheduleUntil(Instrs);
|
|
case BndlSchedState::NoneScheduled: {
|
|
// TODO: Set the window of the DAG that we are interested in.
|
|
if (!ScheduleTopItOpt)
|
|
// We start scheduling at the bottom instr of Instrs (top in TopDown).
|
|
ScheduleTopItOpt = GetSchedPoint(Dir, Instrs);
|
|
// Extend the DAG to include Instrs.
|
|
Interval<Instruction> Extension = DAG.extend(Instrs);
|
|
// Add nodes from the new interval to ready list if they are ready.
|
|
for (auto &I : Extension) {
|
|
auto *N = DAG.getNode(&I);
|
|
if (Dir == SchedDirection::BottomUp ? N->readyBottomUp()
|
|
: N->readyTopDown())
|
|
ReadyList.insert(N);
|
|
}
|
|
// Try schedule all nodes until we can schedule Instrs back-to-back.
|
|
return tryScheduleUntil(Instrs);
|
|
}
|
|
}
|
|
llvm_unreachable("Unhandled BndlSchedState enum");
|
|
}
|
|
|
|
#ifndef NDEBUG
|
|
void Scheduler::dump(raw_ostream &OS) const {
|
|
OS << "ReadyList:\n";
|
|
ReadyList.dump(OS);
|
|
OS << "Dir=" << schedDirectionToStr(Dir) << " "
|
|
<< (Dir == SchedDirection::BottomUp ? "Top" : "Bottom")
|
|
<< " of schedule: ";
|
|
if (ScheduleTopItOpt)
|
|
OS << **ScheduleTopItOpt;
|
|
else
|
|
OS << "Empty";
|
|
OS << "\n";
|
|
}
|
|
void Scheduler::dump() const { dump(dbgs()); }
|
|
#endif // NDEBUG
|
|
|
|
} // namespace llvm::sandboxir
|