[llvm] [SandboxVec][Scheduler] Update full state of scheduler when IR changes (PR #222180)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Sep 10 17:52:26 PDT 2026
https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/222180
>From ed7cec8f7cf16dd2959a277622610262112fd18d Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vasileios.porpodas at amd.com>
Date: Fri, 7 Aug 2026 16:30:17 -0700
Subject: [PATCH] [SandboxVec][Scheduler] Update full state of scheduler when
IR changes
Up until now the scheduler state has not been maintained fully upon IR changes.
This is not really needed for correct operation of the current vectorizer
because scheduling happens before emitting the vectorized IR, so there is no
real need for updating the scheduler with the current implementation.
Nevertheless, the implementation may change so the right thing to do is to
make sure the scheduler's state is always correct.
So this patch implements the missing callbacks for erase, move and set use and
updates the scheduler state accordingly.
---
.../SandboxVectorizer/DependencyGraph.h | 3 +
.../Vectorize/SandboxVectorizer/Scheduler.h | 30 +++-
.../SandboxVectorizer/DependencyGraph.cpp | 12 +-
.../Vectorize/SandboxVectorizer/Scheduler.cpp | 75 ++++++++++
.../SandboxVectorizer/DependencyGraphTest.cpp | 43 ------
.../SandboxVectorizer/SchedulerTest.cpp | 129 ++++++++++++++++++
6 files changed, 237 insertions(+), 55 deletions(-)
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index fc80ae68599c7..8f39c4258bb61 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -568,6 +568,9 @@ class DependencyGraph {
InstrToNodeMap.clear();
DAGInterval = {};
}
+ std::optional<Context::CallbackID> getEraseInstrCB() const {
+ return EraseInstrCB;
+ }
#ifndef NDEBUG
/// \Returns true if the DAG's state is clear. Used in assertions.
bool empty() const {
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
index 57b77cdb7e678..9cbf5aa7899e7 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h
@@ -307,7 +307,7 @@ class Scheduler {
/// The dependency graph is used by the scheduler to determine the legal
/// ordering of instructions.
DependencyGraph DAG;
- friend class SchedulerInternalsAttorney; // For DAG.
+ friend class SchedulerInternalsAttorney; // For DAG and ReadyList.
Context &Ctx;
/// This is the top of the schedule during bottom-up scheduling and the bottom
/// of the schedule during top-down. It points to the position of the last
@@ -319,12 +319,23 @@ class Scheduler {
DenseMap<SchedBundle *, std::unique_ptr<SchedBundle>> Bndls;
/// The BB that we are currently scheduling.
BasicBlock *ScheduledBB = nullptr;
- /// The ID of the callback we register with Sandbox IR.
+ /// The IDs of the callbacks we register with Sandbox IR.
std::optional<Context::CallbackID> CreateInstrCB;
+ std::optional<Context::CallbackID> EraseInstrCB;
+ std::optional<Context::CallbackID> MoveInstrCB;
+ std::optional<Context::CallbackID> SetUseCB;
/// Called by Sandbox IR's callback system, after \p I has been created.
/// NOTE: This should run after DAG's callback has run.
// TODO: Perhaps call DAG's notify function from within this one?
LLVM_ABI void notifyCreateInstr(Instruction *I);
+ /// Called by the callbacks when instruction \p I is about to get
+ /// deleted.
+ LLVM_ABI void notifyEraseInstr(Instruction *I);
+ /// Called by the callbacks when instruction \p I is about to be moved to
+ /// \p To.
+ LLVM_ABI void notifyMoveInstr(Instruction *I, const BBIterator &To);
+ /// Called by the callbacks when \p U's source is about to be set to \p NewSrc
+ LLVM_ABI void notifySetUse(const Use &U, Value *NewSrc);
/// \Returns a scheduling bundle containing \p Instrs.
SchedBundle *createBundle(ArrayRef<Instruction *> Instrs);
@@ -368,12 +379,27 @@ class Scheduler {
: DAG(Dir, AA, Ctx), Ctx(Ctx), Dir(Dir) {
// NOTE: The scheduler's callback depends on the DAG's callback running
// before it and updating the DAG accordingly.
+ EraseInstrCB = Ctx.registerEraseInstrCallback(
+ [this](Instruction *I) { notifyEraseInstr(I); },
+ /*BeforeCB=*/DAG.getEraseInstrCB());
CreateInstrCB = Ctx.registerCreateInstrCallback(
[this](Instruction *I) { notifyCreateInstr(I); });
+ MoveInstrCB = Ctx.registerMoveInstrCallback(
+ [this](Instruction *I, const BBIterator &To) {
+ notifyMoveInstr(I, To);
+ });
+ SetUseCB = Ctx.registerSetUseCallback(
+ [this](const Use &U, Value *NewSrc) { notifySetUse(U, NewSrc); });
}
~Scheduler() {
if (CreateInstrCB)
Ctx.unregisterCreateInstrCallback(*CreateInstrCB);
+ if (EraseInstrCB)
+ Ctx.unregisterEraseInstrCallback(*EraseInstrCB);
+ if (MoveInstrCB)
+ Ctx.unregisterMoveInstrCallback(*MoveInstrCB);
+ if (SetUseCB)
+ Ctx.unregisterSetUseCallback(*SetUseCB);
}
/// Tries to build a schedule that includes all of \p Instrs scheduled at the
/// same scheduling cycle. This essentially checks that there are no
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index 752c470975779..1ac84df828075 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -599,22 +599,14 @@ void DependencyGraph::notifyEraseInstr(Instruction *I) {
SuccN->removeMemPred(MemN, Dir);
}
// NOTE: The unscheduled succs for MemNodes get updated be setMemPred().
- } else {
- // If this is a non-mem node we only need to update UnscheduledSuccs.
- if (!N->scheduled()) {
- for (auto *PredN : N->preds(*this))
- if (!PredN->scheduled())
- PredN->decrUnscheduledDeps();
- for (auto *SuccN : N->succs(*this))
- /// TODO: Does the successor also need to be guarded?
- SuccN->decrUnscheduledDeps();
- }
}
// Finally erase the Node.
InstrToNodeMap.erase(I);
}
void DependencyGraph::notifySetUse(const Use &U, Value *NewSrc) {
+ // TODO: We should eventually move the UnschedDep logic to the scheduler.
+
// If U.User is not in the DAG, then we should not attempt to decrement
// CurrSrcN's unscheduled successors.
// ------- ------- -
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
index 8f9388632edb5..ff79bb69f7505 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp
@@ -153,6 +153,81 @@ void Scheduler::notifyCreateInstr(Instruction *I) {
}
}
+void Scheduler::notifyEraseInstr(Instruction *I) {
+ // We don't maintain the state while reverting.
+ if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
+ return;
+ auto *N = DAG.getNode(I);
+ if (N == nullptr)
+ return;
+ ReadyList.remove(N);
+ // Also decrement the unscheduledDep counter for the dependents and add them
+ // to the ready list if they become ready.
+ auto UpdateNodeAndTryAddToReadyList = [this, N](DGNode *DepN) {
+ if (DepN->scheduled())
+ return;
+ if (!N->scheduled() && !DepN->ready())
+ DepN->decrUnscheduledDeps();
+ if (DepN->ready() && !ReadyList.contains(DepN))
+ ReadyList.insert(DepN);
+ };
+ if (Dir == SchedDirection::BottomUp) {
+ for (auto *DepN : N->preds(DAG))
+ UpdateNodeAndTryAddToReadyList(DepN);
+ } else if (Dir == SchedDirection::TopDown) {
+ for (auto *DepN : N->succs(DAG))
+ UpdateNodeAndTryAddToReadyList(DepN);
+ }
+}
+
+void Scheduler::notifyMoveInstr(Instruction *I, const BBIterator &To) {
+ // We don't maintain the state while reverting.
+ if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
+ return;
+ // We assume that the dependencies have not changed because the user will
+ // only attempt instruction moves that don't modify the dependencies, because
+ // if they did they would not be legal.
+ //
+ // If this assumption does not hold, we would need to empty the ready list and
+ // re-fill it.
+}
+void Scheduler::notifySetUse(const Use &U, Value *NewSrc) {
+ // We don't maintain the state while reverting.
+ if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
+ return;
+ Instruction *DstI = cast<Instruction>(U.getUser());
+ DGNode *DstN = DAG.getNode(DstI);
+ Value *OldSrc = U.get();
+ DGNode *OldSrcN = isa<Instruction>(OldSrc)
+ ? DAG.getNode(cast<Instruction>(OldSrc))
+ : nullptr;
+ DGNode *NewSrcN = isa<Instruction>(NewSrc)
+ ? DAG.getNode(cast<Instruction>(NewSrc))
+ : nullptr;
+ switch (Dir) {
+ case SchedDirection::BottomUp: {
+ // Check if OldSrc is now ready and add it to the ready list.
+ if (OldSrcN && OldSrcN->ready() && !OldSrcN->scheduled() &&
+ !ReadyList.contains(OldSrcN))
+ ReadyList.insert(OldSrcN);
+ // Check if NewSrcN needs to be removed from the ready list.
+ if (NewSrcN && (!DstN || !DstN->scheduled()) && !NewSrcN->ready())
+ ReadyList.remove(NewSrcN);
+ break;
+ }
+ case SchedDirection::TopDown: {
+ // Check if we need to add DstN to the ready list.
+ if (DstN && DstN->ready() && !NewSrcN->scheduled() &&
+ !ReadyList.contains(NewSrcN))
+ ReadyList.insert(NewSrcN);
+ // Check if we need to remove DstN from the ready list.
+ if (DstN && !DstN->ready())
+ ReadyList.remove(NewSrcN);
+ break;
+ }
+ }
+}
+
SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
SchedBundle::ContainerTy Nodes;
Nodes.reserve(Instrs.size());
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index 6d49e8c381081..299fa4265f30c 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -1656,46 +1656,3 @@ define void @foo(i8 %v0) {
Add0->setOperand(0, Sched);
EXPECT_EQ(Add0N->getNumUnscheduledDeps(), 0u);
}
-
-// When erasing a non-mem instruction we must not touch the UnscheduledSuccs
-// of an already-scheduled predecessor, since that counter is set to
-// std::nullopt once a node is scheduled.
-TEST_F(DependencyGraphTest, EraseInstrCallbackNonMemWithScheduledPred) {
- parseIR(C, R"IR(
-define void @foo(i8 %v0) {
- %predSched = add i8 %v0, 0
- %predUnsched = add i8 %v0, 1
- %n = add i8 %predSched, %predUnsched
- 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 *PredSched = cast<sandboxir::BinaryOperator>(&*It++);
- auto *PredUnsched = cast<sandboxir::BinaryOperator>(&*It++);
- auto *N = cast<sandboxir::BinaryOperator>(&*It++);
-
- sandboxir::DependencyGraph DAG(BottomUp, getAA(*LLVMF), Ctx);
- DAG.extend({PredSched, N});
- auto *PredSchedN = DAG.getNode(PredSched);
- auto *PredUnschedN = DAG.getNode(PredUnsched);
- EXPECT_EQ(PredSchedN->getNumUnscheduledDeps(), 1u);
- EXPECT_EQ(PredUnschedN->getNumUnscheduledDeps(), 1u);
-
- // Mark one of N's predecessors as scheduled. Its UnscheduledSuccs becomes
- // std::nullopt.
- PredSchedN->setScheduled();
-
- // Erase N, which is *not* scheduled. This must not attempt to decrement
- // the (now invalid) UnscheduledSuccs of PredSchedN, but should still
- // update the counter of the unscheduled predecessor.
- N->eraseFromParent();
- EXPECT_EQ(DAG.getNode(N), nullptr);
- EXPECT_EQ(PredUnschedN->getNumUnscheduledDeps(), 0u);
-#ifndef NDEBUG
- EXPECT_FALSE(PredSchedN->validUnscheduledDeps());
-#endif
-}
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
index 8e9295ab02660..00509eaa88d7c 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp
@@ -924,6 +924,135 @@ define void @foo(ptr noalias %ptr, ptr noalias %ptr1, ptr noalias %ptr2) {
EXPECT_TRUE(ReadyList.empty());
}
+TEST_F(SchedulerTest, NotifyEraseInst) {
+ parseIR(C, R"IR(
+define void @foo(i8 %v0) {
+ %add0 = add i8 %v0, 0
+ %add1 = add i8 %add0, 1
+ %add2 = add i8 %add0, 2
+ 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 *Add0 = &*It++;
+ auto *Add1 = &*It++;
+ auto *Add2 = &*It++;
+
+ sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx,
+ sandboxir::SchedDirection::BottomUp);
+ auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);
+
+ // 1. Check that erasing instrs can automatically insert ready dependents into
+ // the ready list.
+ EXPECT_TRUE(Sched.trySchedule(Add2));
+ // A dummy trySchedule() to make sure the DAG contains all instrs until Add0.
+ EXPECT_FALSE(Sched.trySchedule({Add0, Add1}));
+ // At this point Add0 would have been ready if it weren't for Add1.
+ auto &ReadyList = sandboxir::SchedulerInternalsAttorney::getReadyList(Sched);
+ EXPECT_FALSE(ReadyList.contains(DAG.getNode(Add0)));
+ EXPECT_TRUE(ReadyList.contains(DAG.getNode(Add1)));
+ // Erasing Add1 should automatically get Add0 into the ready list.
+ Add1->eraseFromParent();
+ EXPECT_TRUE(ReadyList.contains(DAG.getNode(Add0)));
+
+ // 2. Check that erasing a ready instr (Add0), automatically removes it from
+ // the ready list. But first remove its def-use edge that connects it to the
+ // other instrs.
+ Add2->eraseFromParent();
+ Add0->eraseFromParent();
+ EXPECT_FALSE(ReadyList.contains(DAG.getNode(Add0)));
+}
+
+// When erasing a non-mem instruction we must not touch the UnscheduledSuccs
+// of an already-scheduled predecessor, since that counter is set to
+// std::nullopt once a node is scheduled.
+TEST_F(SchedulerTest, NotifyEraseInst_NonMemWithScheduledPred) {
+ parseIR(C, R"IR(
+define void @foo(i8 %v0) {
+ %predSched = add i8 %v0, 0
+ %predUnsched = add i8 %v0, 1
+ %n = add i8 %predSched, %predUnsched
+ 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 *PredSched = cast<sandboxir::BinaryOperator>(&*It++);
+ auto *PredUnsched = cast<sandboxir::BinaryOperator>(&*It++);
+ auto *N = cast<sandboxir::BinaryOperator>(&*It++);
+
+ sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx,
+ sandboxir::SchedDirection::BottomUp);
+ auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);
+ DAG.extend({PredSched, N});
+ auto *PredSchedN = DAG.getNode(PredSched);
+ auto *PredUnschedN = DAG.getNode(PredUnsched);
+ EXPECT_EQ(PredSchedN->getNumUnscheduledDeps(), 1u);
+ EXPECT_EQ(PredUnschedN->getNumUnscheduledDeps(), 1u);
+
+ // Mark one of N's predecessors as scheduled. Its UnscheduledSuccs becomes
+ // std::nullopt.
+ PredSchedN->setScheduled();
+
+ // Erase N, which is *not* scheduled. This must not attempt to decrement
+ // the (now invalid) UnscheduledSuccs of PredSchedN, but should still
+ // update the counter of the unscheduled predecessor.
+ N->eraseFromParent();
+ EXPECT_EQ(DAG.getNode(N), nullptr);
+ EXPECT_EQ(PredUnschedN->getNumUnscheduledDeps(), 0u);
+#ifndef NDEBUG
+ EXPECT_FALSE(PredSchedN->validUnscheduledDeps());
+#endif
+}
+
+TEST_F(SchedulerTest, NotifySetUse) {
+ parseIR(C, R"IR(
+define void @foo(i8 %v0, i8 %v1) {
+ %add0 = add i8 %v0, 0
+ %add1 = add i8 %add0, 1
+ %add2 = add i8 %add0, %add1
+ 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 *Add0 = &*It++;
+ auto *Add1 = &*It++;
+ auto *Add2 = &*It++;
+ auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+ auto *V1 = F->getArg(1);
+
+ sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx,
+ sandboxir::SchedDirection::BottomUp);
+ auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);
+ // Dummy trySchedule() to make sure we have built the whole DAG.
+ EXPECT_FALSE(Sched.trySchedule({Ret, Add0, Add1, Add2}));
+ auto &ReadyList = sandboxir::SchedulerInternalsAttorney::getReadyList(Sched);
+
+ // 1. Check that removing the dependency Add0->Add1 will make Add0 ready.
+ Sched.trySchedule({Add2});
+ EXPECT_FALSE(ReadyList.contains(DAG.getNode(Add0)));
+ EXPECT_TRUE(ReadyList.contains(DAG.getNode(Add1)));
+ Add1->setOperand(0, V1);
+ EXPECT_TRUE(ReadyList.contains(DAG.getNode(Add0)));
+
+ // 2. Check that re-introducing the dependency Add0->Add1 will remove Add0
+ // from the ready list.
+ Add1->setOperand(0, Add0);
+ EXPECT_FALSE(ReadyList.contains(DAG.getNode(Add0)));
+ EXPECT_TRUE(ReadyList.contains(DAG.getNode(Add1)));
+}
+
TEST_F(SchedulerTest, ReadyList) {
parseIR(C, R"IR(
define void @foo(ptr %ptr) {
More information about the llvm-commits
mailing list