-
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. (#112449
) 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
540 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
126 changes: 126 additions & 0 deletions
126
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,126 @@ | ||
//===- 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); | ||
/// Schedules all nodes in \p Bndl, marks them as scheduled, updates the | ||
/// UnscheduledSuccs counter of all dependency predecessors, and adds any of | ||
/// them that become ready to the ready list. | ||
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
169 changes: 169 additions & 0 deletions
169
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,169 @@ | ||
//===- 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 { | ||
|
||
// 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 = *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 of instructions, instead of `Instrs` for fast lookups. | ||
DenseSet<Instruction *> InstrsToDefer(Instrs.begin(), Instrs.end()); | ||
// This collects the nodes that correspond to instructions found in `Instrs` | ||
// that have just become ready. These nodes won't be scheduled right away. | ||
SmallVector<DGNode *, 8> DeferredNodes; | ||
|
||
// 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. | ||
while (!ReadyList.empty()) { | ||
auto *ReadyN = ReadyList.pop(); | ||
if (InstrsToDefer.contains(ReadyN->getInstruction())) { | ||
// If the ready instruction is one of those in `Instrs`, then we don't | ||
// schedule it right away. Instead we defer it until we can schedule it | ||
// along with the rest of the instructions in `Instrs`, at the same | ||
// time in a single scheduling bundle. | ||
DeferredNodes.push_back(ReadyN); | ||
bool ReadyToScheduleDeferred = DeferredNodes.size() == Instrs.size(); | ||
if (ReadyToScheduleDeferred) { | ||
scheduleAndUpdateReadyList(*createBundle(Instrs)); | ||
return true; | ||
} | ||
} else { | ||
// If the ready instruction is not found in `Instrs`, then we wrap it in a | ||
// scheduling bundle and schedule it right away. | ||
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.