-
Notifications
You must be signed in to change notification settings - Fork 11.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[SandboxVec][Scheduler] Boilerplate and initial implementation.
This patch implements a ready-list-based scheduler that operates on DependencyGraph. It is used by the sandbox vectorizer to test the legality of vectorizing a group of instrs. SchedBundle is a helper container, containing all DGNodes that correspond to the instructions that we are attempting to schedule with trySchedule(Instrs).
- Loading branch information
Showing
8 changed files
with
529 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
124 changes: 124 additions & 0 deletions
124
llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,124 @@ | ||
//===- Scheduler.h ----------------------------------------------*- C++ -*-===// | ||
// | ||
// 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 | ||
// | ||
//===----------------------------------------------------------------------===// | ||
// | ||
// This is the bottom-up list scheduler used by the vectorizer. It is used for | ||
// checking the legality of vectorization and for scheduling instructions in | ||
// such a way that makes vectorization possible, if legal. | ||
// | ||
// The legality check is performed by `trySchedule(Instrs)`, which will try to | ||
// schedule the IR until all instructions in `Instrs` can be scheduled together | ||
// back-to-back. If this fails then it is illegal to vectorize `Instrs`. | ||
// | ||
// Internally the scheduler uses the vectorizer-specific DependencyGraph class. | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H | ||
#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H | ||
|
||
#include "llvm/SandboxIR/Instruction.h" | ||
#include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h" | ||
#include <queue> | ||
|
||
namespace llvm::sandboxir { | ||
|
||
class PriorityCmp { | ||
public: | ||
bool operator()(const DGNode *N1, const DGNode *N2) { | ||
// TODO: This should be a hierarchical comparator. | ||
return N1->getInstruction()->comesBefore(N2->getInstruction()); | ||
} | ||
}; | ||
|
||
/// The list holding nodes that are ready to schedule. Used by the scheduler. | ||
class ReadyListContainer { | ||
PriorityCmp Cmp; | ||
/// Control/Other dependencies are not modeled by the DAG to save memory. | ||
/// These have to be modeled in the ready list for correctness. | ||
/// This means that the list will hold back nodes that need to meet such | ||
/// unmodeled dependencies. | ||
std::priority_queue<DGNode *, std::vector<DGNode *>, PriorityCmp> List; | ||
|
||
public: | ||
ReadyListContainer() : List(Cmp) {} | ||
void insert(DGNode *N) { List.push(N); } | ||
DGNode *pop() { | ||
auto *Back = List.top(); | ||
List.pop(); | ||
return Back; | ||
} | ||
bool empty() const { return List.empty(); } | ||
#ifndef NDEBUG | ||
void dump(raw_ostream &OS) const; | ||
LLVM_DUMP_METHOD void dump() const; | ||
#endif // NDEBUG | ||
}; | ||
|
||
/// The nodes that need to be scheduled back-to-back in a single scheduling | ||
/// cycle form a SchedBundle. | ||
class SchedBundle { | ||
public: | ||
using ContainerTy = SmallVector<DGNode *, 4>; | ||
|
||
private: | ||
ContainerTy Nodes; | ||
|
||
public: | ||
SchedBundle() = default; | ||
SchedBundle(ContainerTy &&Nodes) : Nodes(std::move(Nodes)) {} | ||
using iterator = ContainerTy::iterator; | ||
using const_iterator = ContainerTy::const_iterator; | ||
iterator begin() { return Nodes.begin(); } | ||
iterator end() { return Nodes.end(); } | ||
const_iterator begin() const { return Nodes.begin(); } | ||
const_iterator end() const { return Nodes.end(); } | ||
/// \Returns the bundle node that comes before the others in program order. | ||
DGNode *getTop() const; | ||
/// \Returns the bundle node that comes after the others in program order. | ||
DGNode *getBot() const; | ||
/// Move all bundle instructions to \p Where back-to-back. | ||
void cluster(BasicBlock::iterator Where); | ||
#ifndef NDEBUG | ||
void dump(raw_ostream &OS) const; | ||
LLVM_DUMP_METHOD void dump() const; | ||
#endif | ||
}; | ||
|
||
/// The list scheduler. | ||
class Scheduler { | ||
ReadyListContainer ReadyList; | ||
DependencyGraph DAG; | ||
std::optional<BasicBlock::iterator> ScheduleTopItOpt; | ||
SmallVector<std::unique_ptr<SchedBundle>> Bndls; | ||
|
||
/// \Returns a scheduling bundle containing \p Instrs. | ||
SchedBundle *createBundle(ArrayRef<Instruction *> Instrs); | ||
/// Schedule nodes until we can schedule \p Instrs back-to-back. | ||
bool tryScheduleUntil(ArrayRef<Instruction *> Instrs); | ||
|
||
void scheduleAndUpdateReadyList(SchedBundle &Bndl); | ||
|
||
/// Disable copies. | ||
Scheduler(const Scheduler &) = delete; | ||
Scheduler &operator=(const Scheduler &) = delete; | ||
|
||
public: | ||
Scheduler(AAResults &AA) : DAG(AA) {} | ||
~Scheduler() {} | ||
|
||
bool trySchedule(ArrayRef<Instruction *> Instrs); | ||
|
||
#ifndef NDEBUG | ||
void dump(raw_ostream &OS) const; | ||
LLVM_DUMP_METHOD void dump() const; | ||
#endif | ||
}; | ||
|
||
} // namespace llvm::sandboxir | ||
|
||
#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
160 changes: 160 additions & 0 deletions
160
llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,160 @@ | ||
//===- 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" | ||
|
||
namespace llvm::sandboxir { | ||
|
||
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 = *ScheduleTopItOpt; | ||
// Move all instructions in `Bndl` to `Where`. | ||
Bndl.cluster(Where); | ||
// Update the last scheduled bundle. | ||
ScheduleTopItOpt = Bndl.getTop()->getInstruction()->getIterator(); | ||
// Set nodes as "scheduled" and decrement the UnsceduledSuccs counter of all | ||
// dependency predecessors. | ||
for (DGNode *N : Bndl) { | ||
N->setScheduled(true); | ||
for (auto *DepN : N->preds(DAG)) { | ||
// TODO: preds() should not return nullptr. | ||
if (DepN == nullptr) | ||
continue; | ||
DepN->decrUnscheduledSuccs(); | ||
if (DepN->ready()) | ||
ReadyList.insert(DepN); | ||
} | ||
} | ||
} | ||
|
||
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.push_back(std::move(BndlPtr)); | ||
return Bndl; | ||
} | ||
|
||
bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) { | ||
// Use a set for fast lookups. | ||
DenseSet<Instruction *> InstrsToDefer(Instrs.begin(), Instrs.end()); | ||
SmallVector<DGNode *, 8> DeferredNodes; | ||
|
||
// Keep scheduling ready nodes. | ||
while (!ReadyList.empty()) { | ||
auto *ReadyN = ReadyList.pop(); | ||
// We defer scheduling of instructions in `Instrs` until we can schedule all | ||
// of them at the same time in a single scheduling bundle. | ||
if (InstrsToDefer.contains(ReadyN->getInstruction())) { | ||
DeferredNodes.push_back(ReadyN); | ||
bool ReadyToScheduleDeferred = DeferredNodes.size() == Instrs.size(); | ||
if (ReadyToScheduleDeferred) { | ||
scheduleAndUpdateReadyList(*createBundle(Instrs)); | ||
return true; | ||
} | ||
} else { | ||
scheduleAndUpdateReadyList(*createBundle({ReadyN->getInstruction()})); | ||
} | ||
} | ||
assert(DeferredNodes.size() != Instrs.size() && | ||
"We should have succesfully scheduled and early-returned!"); | ||
return false; | ||
} | ||
|
||
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!"); | ||
// Extend the DAG to include Instrs. | ||
Interval<Instruction> Extension = DAG.extend(Instrs); | ||
// TODO: Set the window of the DAG that we are interested in. | ||
// We start scheduling at the bottom instr of Instrs. | ||
auto getBottomI = [](ArrayRef<Instruction *> Instrs) -> Instruction * { | ||
return *min_element(Instrs, | ||
[](auto *I1, auto *I2) { return I1->comesBefore(I2); }); | ||
}; | ||
ScheduleTopItOpt = std::next(getBottomI(Instrs)->getIterator()); | ||
// Add nodes to ready list. | ||
for (auto &I : Extension) { | ||
auto *N = DAG.getNode(&I); | ||
if (N->ready()) | ||
ReadyList.insert(N); | ||
} | ||
// Try schedule all nodes until we can schedule Instrs back-to-back. | ||
return tryScheduleUntil(Instrs); | ||
} | ||
|
||
#ifndef NDEBUG | ||
void Scheduler::dump(raw_ostream &OS) const { | ||
OS << "ReadyList:\n"; | ||
ReadyList.dump(OS); | ||
} | ||
void Scheduler::dump() const { dump(dbgs()); } | ||
#endif // NDEBUG | ||
|
||
} // namespace llvm::sandboxir |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.