[llvm] [SandboxVec][Scheduler] Don't insert scheduled instrs into the ready list (PR #127688)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Feb 18 11:27:43 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: vporpo (vporpo)
<details>
<summary>Changes</summary>
In a particular scenario (see test) we used to insert scheduled instructions into the ready list. This patch fixes this by fixing the trimSchedule() function.
---
Full diff: https://github.com/llvm/llvm-project/pull/127688.diff
4 Files Affected:
- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h (+4)
- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h (+12-1)
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp (+44-18)
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp (+71)
``````````diff
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index b25f96571741e..ede51c28fc94d 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -123,6 +123,10 @@ class DGNode {
--UnscheduledSuccs;
}
void incrUnscheduledSuccs() { ++UnscheduledSuccs; }
+ void resetScheduleState() {
+ UnscheduledSuccs = 0;
+ Scheduled = false;
+ }
/// \Returns true if all dependent successors have been scheduled.
bool ready() const { return UnscheduledSuccs == 0; }
/// \Returns true if this node has been scheduled.
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
index 6c2315af0e797..6b56f348f328c 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -61,7 +61,18 @@ class ReadyListContainer {
public:
ReadyListContainer() : List(Cmp) {}
- void insert(DGNode *N) { List.push(N); }
+ void insert(DGNode *N) {
+#ifndef NDEBUG
+ assert(!N->scheduled() && "Don't insert a scheduled node!");
+ auto ListCopy = List;
+ while (!ListCopy.empty()) {
+ DGNode *Top = ListCopy.top();
+ ListCopy.pop();
+ assert(Top != N && "Node already exists in ready list!");
+ }
+#endif
+ List.push(N);
+ }
DGNode *pop() {
auto *Back = List.top();
List.pop();
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
index 2f7d7087ca880..3e37e07aabc5c 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
@@ -77,12 +77,12 @@ void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
// 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)) {
DepN->decrUnscheduledSuccs();
- if (DepN->ready())
+ if (DepN->ready() && !DepN->scheduled())
ReadyList.insert(DepN);
}
+ N->setScheduled(true);
}
}
@@ -188,6 +188,19 @@ Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
}
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()
+ //
Instruction *TopI = &*ScheduleTopItOpt.value();
Instruction *LowestI = VecUtils::getLowest(Instrs);
// Destroy the schedule bundles from LowestI all the way to the top.
@@ -199,13 +212,28 @@ void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
if (auto *SB = N->getSchedBundle())
eraseBundle(SB);
}
- // TODO: For now we clear the DAG. Trim view once it gets implemented.
- Bndls.clear();
- DAG.clear();
-
- // Since we are scheduling NewRegion from scratch, we clear the ready lists.
- // The nodes currently in the list may not be ready after clearing the View.
+ // 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 from LowestI to the top of the schedule. DAG Nodes that are above the
+ // top of schedule that depend on nodes that got reset need to have their
+ // UnscheduledSuccs adjusted.
+ Interval<Instruction> ResetIntvl(TopI, LowestI);
+ for (Instruction &I : ResetIntvl) {
+ auto *N = DAG.getNode(&I);
+ N->resetScheduleState();
+ // 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();
+ }
+ // 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->ready())
+ ReadyList.insert(N);
+ }
}
bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
@@ -214,6 +242,12 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
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.
@@ -230,21 +264,13 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
// top-most part of the schedule that includes the instrs in the bundle and
// re-schedule.
trimSchedule(Instrs);
- ScheduleTopItOpt = std::nullopt;
- [[fallthrough]];
+ ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
+ 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.
ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
-
- // 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;
- }
-
// Extend the DAG to include Instrs.
Interval<Instruction> Extension = DAG.extend(Instrs);
// Add nodes to ready list.
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
index 0d5d86acaee89..7d88337e77faa 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
@@ -246,6 +246,77 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) {
EXPECT_TRUE(Sched.trySchedule({L0, L1}));
}
+// Check scheduling in the following order: {A0,A1},{B0,B1},{C0,C1},{D0,D1}
+// assuming program order: B0,B1,C0,C1,D0,D1,E0,D1.
+// This will effectively schedule nodes below already scheduled nodes, which
+// can expose issues in the code that adds nodes to the ready list.
+// For example, we schedule {D0,D1} while {C0,C1} are scheduled and there is
+// a dependency D0->C0 and D1->C1.
+//
+// {A0,A1} {B0,B1} {C0,C1} {D0,D1}
+// B0,B1 | S
+// |\ |
+// | C0,C1 | | S | S
+// | | \ | |
+// | | D0,D1 | | S
+// | / |
+// A0,A1 | S | S
+// +------------------------+
+// | Legend |: DAG |
+// | S: Scheduled |
+TEST_F(SchedulerTest, ScheduledPredecessors) {
+ parseIR(C, R"IR(
+define void @foo(ptr noalias %ptrA0, ptr noalias %ptrA1,
+ ptr noalias %ptrB0, ptr noalias %ptrB1,
+ ptr noalias %ptrD0, ptr noalias %ptrD1) {
+ %B0 = load i8, ptr %ptrB0
+ %B1 = load i8, ptr %ptrB1
+ %C0 = add i8 %B0, 0
+ %C1 = add i8 %B1, 1
+ store i8 %C0, ptr %ptrD0
+ store i8 %C1, ptr %ptrD1
+ store i8 %B0, ptr %ptrA0
+ store i8 %B1, ptr %ptrA1
+ ret void
+}
+)IR");
+ llvm::Function *LLVMF = &*M->getFunction("foo");
+ sandboxir::Context Ctx(C);
+ auto *F = Ctx.createFunction(LLVMF);
+ auto *BB = &*F->begin();
+ auto It = BB->begin();
+ auto *B1 = cast<sandboxir::LoadInst>(&*It++);
+ auto *B0 = cast<sandboxir::LoadInst>(&*It++);
+ auto *C1 = cast<sandboxir::BinaryOperator>(&*It++);
+ auto *C0 = cast<sandboxir::BinaryOperator>(&*It++);
+ auto *D1 = cast<sandboxir::StoreInst>(&*It++);
+ auto *D0 = cast<sandboxir::StoreInst>(&*It++);
+ auto *A1 = cast<sandboxir::StoreInst>(&*It++);
+ auto *A0 = cast<sandboxir::StoreInst>(&*It++);
+ auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+
+ sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);
+ EXPECT_TRUE(Sched.trySchedule({A0, A1}));
+ // NOTE: We schedule the intermediate nodes between {A0,A1} and {B0,B1} by
+ // hand one by one to make sure they are scheduled in that order because
+ // the scheduler may reorder them a bit if we let it do it.
+ EXPECT_TRUE(Sched.trySchedule(D0));
+ EXPECT_TRUE(Sched.trySchedule(D1));
+ EXPECT_TRUE(Sched.trySchedule(C0));
+ EXPECT_TRUE(Sched.trySchedule(C1));
+ EXPECT_TRUE(Sched.trySchedule({B0, B1}));
+ // At this point all nodes must have been scheduled from B0,B1 to A0,A1.
+ // The ones in between are scheduled as single-instruction nodes.
+ // So when we attempt to schedule {C0,C1} we will need to reschedule.
+ // At this point we will trim the schedule from {C0,C1} upwards.
+ EXPECT_TRUE(Sched.trySchedule({C0, C1}));
+ // Now the schedule should only contain {C0,C1} which should be marked as
+ // "scheduled".
+ // {D0,D1} are below {C0,C1}, so we grow the DAG downwards, while
+ // {C0,C1} are marked as "scheduled" above them.
+ EXPECT_TRUE(Sched.trySchedule({D0, D1}));
+}
+
TEST_F(SchedulerTest, DontCrossBBs) {
parseIR(C, R"IR(
define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %v0, i8 %v1) {
``````````
</details>
https://github.com/llvm/llvm-project/pull/127688
More information about the llvm-commits
mailing list