Skip to content

Commit

Permalink
[SandboxVec][DAG] Implement UnscheduledSuccs (#112255)
Browse files Browse the repository at this point in the history
This patch implements the UnscheduledSuccs counter in DGNode. It counts
the number of unscheduled successors and is used by the scheduler to
determine when a node is ready.
  • Loading branch information
vporpo authored Oct 14, 2024
1 parent 3484ed9 commit fc08ad6
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -96,16 +96,25 @@ class DGNode {
// TODO: Use a PointerIntPair for SubclassID and I.
/// For isa/dyn_cast etc.
DGNodeID SubclassID;
/// The number of unscheduled successors.
unsigned UnscheduledSuccs = 0;
/// This is true if this node has been scheduled.
bool Scheduled = false;

DGNode(Instruction *I, DGNodeID ID) : I(I), SubclassID(ID) {}
friend class MemDGNode; // For constructor.
friend class MemDGNode; // For constructor.
friend class DependencyGraph; // For UnscheduledSuccs

public:
DGNode(Instruction *I) : I(I), SubclassID(DGNodeID::DGNode) {
assert(!isMemDepNodeCandidate(I) && "Expected Non-Mem instruction, ");
}
DGNode(const DGNode &Other) = delete;
virtual ~DGNode() = default;
/// \Returns the number of unscheduled successors.
unsigned getNumUnscheduledSuccs() const { return UnscheduledSuccs; }
/// \Returns true if this node has been scheduled.
bool scheduled() const { return Scheduled; }
/// \Returns true if this is before \p Other in program order.
bool comesBefore(const DGNode *Other) { return I->comesBefore(Other->I); }
using iterator = PredIterator;
Expand Down Expand Up @@ -215,8 +224,16 @@ class MemDGNode final : public DGNode {
MemDGNode *getPrevNode() const { return PrevMemN; }
/// \Returns the next Mem DGNode in instruction order.
MemDGNode *getNextNode() const { return NextMemN; }
/// Adds the mem dependency edge PredN->this.
void addMemPred(MemDGNode *PredN) { MemPreds.insert(PredN); }
/// Adds the mem dependency edge PredN->this. This also increments the
/// UnscheduledSuccs counter of the predecessor if this node has not been
/// scheduled.
void addMemPred(MemDGNode *PredN) {
auto Inserted = MemPreds.insert(PredN).second;
assert(Inserted && "PredN already exists!");
if (!Scheduled) {
++PredN->UnscheduledSuccs;
}
}
/// \Returns true if there is a memory dependency N->this.
bool hasMemPred(DGNode *N) const {
if (auto *MN = dyn_cast<MemDGNode>(N))
Expand Down Expand Up @@ -284,6 +301,10 @@ class DependencyGraph {
/// \p DstN.
void scanAndAddDeps(MemDGNode &DstN, const Interval<MemDGNode> &SrcScanRange);

/// Sets the UnscheduledSuccs of all DGNodes in \p NewInterval based on
/// def-use edges.
void setDefUseUnscheduledSuccs(const Interval<Instruction> &NewInterval);

/// Create DAG nodes for instrs in \p NewInterval and update the MemNode
/// chain.
void createNewNodes(const Interval<Instruction> &NewInterval);
Expand Down
70 changes: 60 additions & 10 deletions llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -59,21 +59,17 @@ bool PredIterator::operator==(const PredIterator &Other) const {
}

#ifndef NDEBUG
void DGNode::print(raw_ostream &OS, bool PrintDeps) const { I->dumpOS(OS); }
void DGNode::dump() const {
print(dbgs());
dbgs() << "\n";
void DGNode::print(raw_ostream &OS, bool PrintDeps) const {
OS << *I << " USuccs:" << UnscheduledSuccs << "\n";
}
void DGNode::dump() const { print(dbgs()); }
void MemDGNode::print(raw_ostream &OS, bool PrintDeps) const {
I->dumpOS(OS);
DGNode::print(OS, false);
if (PrintDeps) {
// Print memory preds.
static constexpr const unsigned Indent = 4;
for (auto *Pred : MemPreds) {
OS.indent(Indent) << "<-";
Pred->print(OS, false);
OS << "\n";
}
for (auto *Pred : MemPreds)
OS.indent(Indent) << "<-" << *Pred->getInstruction() << "\n";
}
}
#endif // NDEBUG
Expand Down Expand Up @@ -215,6 +211,58 @@ void DependencyGraph::scanAndAddDeps(MemDGNode &DstN,
}
}

void DependencyGraph::setDefUseUnscheduledSuccs(
const Interval<Instruction> &NewInterval) {
// +---+
// | | Def
// | | |
// | | v
// | | Use
// +---+
// Set the intra-interval counters in NewInterval.
for (Instruction &I : NewInterval) {
for (Value *Op : I.operands()) {
auto *OpI = dyn_cast<Instruction>(Op);
if (OpI == nullptr)
continue;
if (!NewInterval.contains(OpI))
continue;
auto *OpN = getNode(OpI);
if (OpN == nullptr)
continue;
++OpN->UnscheduledSuccs;
}
}

// Now handle the cross-interval edges.
bool NewIsAbove = DAGInterval.empty() || NewInterval.comesBefore(DAGInterval);
const auto &TopInterval = NewIsAbove ? NewInterval : DAGInterval;
const auto &BotInterval = NewIsAbove ? DAGInterval : NewInterval;
// +---+
// |Top|
// | | Def
// +---+ |
// | | v
// |Bot| Use
// | |
// +---+
// Walk over all instructions in "BotInterval" and update the counter
// of operands that are in "TopInterval".
for (Instruction &BotI : BotInterval) {
for (Value *Op : BotI.operands()) {
auto *OpI = dyn_cast<Instruction>(Op);
if (OpI == nullptr)
continue;
if (!TopInterval.contains(OpI))
continue;
auto *OpN = getNode(OpI);
if (OpN == nullptr)
continue;
++OpN->UnscheduledSuccs;
}
}
}

void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
// Create Nodes only for the new sections of the DAG.
DGNode *LastN = getOrCreateNode(NewInterval.top());
Expand Down Expand Up @@ -260,6 +308,8 @@ void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
}
#endif // NDEBUG
}

setDefUseUnscheduledSuccs(NewInterval);
}

Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,11 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
EXPECT_TRUE(N0->memPreds().empty());
EXPECT_THAT(N1->memPreds(), testing::ElementsAre(N0));
EXPECT_TRUE(N2->preds(DAG).empty());

// Check UnscheduledSuccs.
EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u); // N1
EXPECT_EQ(N1->getNumUnscheduledSuccs(), 0u);
EXPECT_EQ(N2->getNumUnscheduledSuccs(), 0u);
}

TEST_F(DependencyGraphTest, Preds) {
Expand Down Expand Up @@ -286,6 +291,14 @@ define i8 @foo(i8 %v0, i8 %v1) {
EXPECT_THAT(StN->preds(DAG),
testing::UnorderedElementsAre(CallN, CallN, AddN2));
EXPECT_THAT(RetN->preds(DAG), testing::ElementsAre(AddN2));

// Check UnscheduledSuccs.
EXPECT_EQ(AddN0->getNumUnscheduledSuccs(), 1u); // AddN2
EXPECT_EQ(AddN1->getNumUnscheduledSuccs(), 2u); // AddN2, CallN
EXPECT_EQ(AddN2->getNumUnscheduledSuccs(), 2u); // StN, RetN
EXPECT_EQ(CallN->getNumUnscheduledSuccs(), 2u); // StN, StN
EXPECT_EQ(StN->getNumUnscheduledSuccs(), 0u);
EXPECT_EQ(RetN->getNumUnscheduledSuccs(), 0u);
}

TEST_F(DependencyGraphTest, MemDGNode_getPrevNode_getNextNode) {
Expand Down Expand Up @@ -711,6 +724,8 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
EXPECT_EQ(DAG.getInterval().top(), S3);
EXPECT_EQ(DAG.getInterval().bottom(), S3);
[[maybe_unused]] auto *S3N = cast<sandboxir::MemDGNode>(DAG.getNode(S3));
// Check UnscheduledSuccs.
EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 0u);
}
{
// Scenario 2: Extend below
Expand All @@ -722,6 +737,10 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
EXPECT_TRUE(S4N->hasMemPred(S3N));
EXPECT_TRUE(S5N->hasMemPred(S4N));
EXPECT_TRUE(S5N->hasMemPred(S3N));
// Check UnscheduledSuccs.
EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 2u); // S4N, S5N
EXPECT_EQ(S4N->getNumUnscheduledSuccs(), 1u); // S5N
EXPECT_EQ(S5N->getNumUnscheduledSuccs(), 0u);
}
{
// Scenario 3: Extend above
Expand All @@ -746,5 +765,12 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
EXPECT_TRUE(S5N->hasMemPred(S3N));
EXPECT_TRUE(S5N->hasMemPred(S2N));
EXPECT_TRUE(S5N->hasMemPred(S1N));

// Check UnscheduledSuccs.
EXPECT_EQ(S1N->getNumUnscheduledSuccs(), 4u); // S2N, S3N, S4N, S5N
EXPECT_EQ(S2N->getNumUnscheduledSuccs(), 3u); // S3N, S4N, S5N
EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 2u); // S4N, S5N
EXPECT_EQ(S4N->getNumUnscheduledSuccs(), 1u); // S5N
EXPECT_EQ(S5N->getNumUnscheduledSuccs(), 0u);
}
}

0 comments on commit fc08ad6

Please sign in to comment.