[llvm] [AMDGPU][CoExecSched] Add free-proximity scheduling heuristic (PR #225334)
Romanov Vlad via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 22 02:07:03 PDT 2026
https://github.com/romanovvlad updated https://github.com/llvm/llvm-project/pull/225334
>From 4599d0ccaf692ef8d709911c289cc5917d61153c Mon Sep 17 00:00:00 2001
From: Vlad <Vladislav.Romanov at amd.com>
Date: Tue, 22 Sep 2026 02:52:13 -0500
Subject: [PATCH] [AMDGPU][CoExecSched] Add free-proximity scheduling heuristic
Add a new scheduling heuristic that sorts PrioritySUs preferring
candidates closer to killing registers, reducing live register pressure.
Two-level comparison:
1. Prefer cand with more registers whose producers all have exactly 1
unscheduled successor (scheduling can free those registers immediately).
2. Prefer cand with lower min unscheduled successors among remaining
registers (closer to killing in future).
Unscheduled successors are computed manually as NumSuccsLeft is not
decremented in top-down scheduling.
When active, PrioritySUs are rebuilt every pick cycle to reflect the
current scheduling state.
By default (auto mode) the heuristic is enabled when VGPR pressure
approaches the excess limit, which is recomputed per each pick.
Option -amdgpu-coexec-reg-free-proximity={off,auto,always} controls the
behavior.
---
.../AMDGPU/AMDGPUCoExecSchedStrategy.cpp | 211 +++++++++++++++---
.../Target/AMDGPU/AMDGPUCoExecSchedStrategy.h | 32 ++-
.../AMDGPU/coexec-reg-free-proximity.mir | 56 +++++
3 files changed, 261 insertions(+), 38 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/coexec-reg-free-proximity.mir
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.cpp
index 50d444068e01d..23e4117157119 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.cpp
@@ -38,6 +38,19 @@ static cl::opt<CarriedLatency> BlockCarriedLatency(
CarriedLatency::All, "all",
"Pad latency for any SU with an incoming ds_load dependency.")));
+enum class RegFreeProximityMode { Off, Auto, Always };
+
+static cl::opt<RegFreeProximityMode> CoexecRegFreeProximity(
+ "amdgpu-coexec-reg-free-proximity", cl::Hidden,
+ cl::init(RegFreeProximityMode::Auto),
+ cl::desc("Prioritize instructions which are expected to free a register "
+ "sooner (lower min NumSuccsLeft)."),
+ cl::values(clEnumValN(RegFreeProximityMode::Off, "off", "Disabled."),
+ clEnumValN(RegFreeProximityMode::Auto, "auto",
+ "Enabled when HighPressure is set."),
+ clEnumValN(RegFreeProximityMode::Always, "always",
+ "Always enabled.")));
+
namespace {
// Used to disable post-RA scheduling with function level granularity.
@@ -463,32 +476,163 @@ SUnit *HardwareUnitInfo::getNextTargetSU(bool LookDeep) const {
return TargetSU;
}
-void HardwareUnitInfo::insert(SUnit *SU, unsigned BlockingCycles) {
- if (!AllSUs.insert(SU))
- llvm_unreachable("HardwareUnit already contains SU!");
+int HardwareUnitInfo::compareDepth(SUnit *Candidate, SUnit *Existing) const {
+ const unsigned CurDepth = Existing->getDepth();
+ const unsigned CandDepth = Candidate->getDepth();
- TotalCycles += BlockingCycles;
+ if (CandDepth > CurDepth)
+ return -1;
+ if (CandDepth == CurDepth)
+ return 0;
+ return 1;
+}
+
+int HardwareUnitInfo::compareRegFreeProximity(SUnit *Candidate,
+ SUnit *Existing) const {
+ const MachineInstr *CandMI = Candidate->getInstr();
+ const bool IsMemOp = SIInstrInfo::isDS(*CandMI) ||
+ SIInstrInfo::isFLAT(*CandMI) ||
+ SIInstrInfo::isVMEM(*CandMI);
+ if (IsMemOp)
+ return compareDepth(Candidate, Existing);
+ auto getRegStats = [](const SUnit *SU) {
+ SmallDenseMap<Register, unsigned, 8> RegMaxUnsched;
+
+ for (const SDep &Pred : SU->Preds) {
+ if (Pred.getKind() != SDep::Data)
+ continue;
+ Register Reg = Pred.getReg();
+ if (!Reg)
+ continue;
+ const SUnit *PredSU = Pred.getSUnit();
+ unsigned Unscheduled = 0;
+ for (const SDep &Succ : PredSU->Succs)
+ if (Succ.getKind() == SDep::Data && !Succ.getSUnit()->isScheduled)
+ ++Unscheduled;
+
+ LLVM_DEBUG({
+ dbgs() << " pred SU(" << PredSU->NodeNum
+ << ") reg=" << printReg(Reg) << " unschedSuccs=" << Unscheduled
+ << " ";
+ if (PredSU->getInstr())
+ PredSU->getInstr()->print(dbgs(), /*IsStandalone=*/true,
+ /*SkipOpers=*/false, /*SkipDebugLoc=*/true);
+ else
+ dbgs() << "<no instr>";
+ dbgs() << "\n";
+ });
+ RegMaxUnsched[Reg] = std::max(RegMaxUnsched[Reg], Unscheduled);
+ }
+
+ unsigned Frees = 0;
+ unsigned MinOther =
+ RegMaxUnsched.empty() ? 0 : RegMaxUnsched.begin()->second;
+ for (auto &[Reg, MaxUnsched] : RegMaxUnsched) {
+ if (MaxUnsched < 2)
+ ++Frees;
+ else
+ MinOther = std::min(MinOther, MaxUnsched);
+ }
+ LLVM_DEBUG(dbgs() << " => frees=" << Frees
+ << " minOther=" << MinOther
+ << " (regs=" << RegMaxUnsched.size() << ")\n");
+ return std::pair(Frees, MinOther);
+ };
+
+ LLVM_DEBUG(dbgs() << " RegFreeProximity: Existing SU("
+ << Existing->NodeNum << ") vs Candidate SU("
+ << Candidate->NodeNum << ")\n");
+ auto [CurFrees, CurMinOther] = getRegStats(Existing);
+ auto [CandFrees, CandMinOther] = getRegStats(Candidate);
+
+ LLVM_DEBUG(dbgs() << " Existing: frees=" << CurFrees << " minOther="
+ << CurMinOther << " Candidate: frees=" << CandFrees
+ << " minOther=" << CandMinOther << "\n");
+
+ if (CandFrees > CurFrees) {
+ LLVM_DEBUG(dbgs() << " -> Candidate wins (more frees)\n");
+ return 1;
+ }
+ if (CandFrees < CurFrees) {
+ LLVM_DEBUG(dbgs() << " -> Existing wins (more frees)\n");
+ return -1;
+ }
+
+ if (CandMinOther < CurMinOther) {
+ LLVM_DEBUG(dbgs() << " -> Candidate wins (closer to freeing)\n");
+ return 1;
+ }
+ if (CandMinOther > CurMinOther) {
+ LLVM_DEBUG(dbgs() << " -> Existing wins (closer to freeing)\n");
+ return -1;
+ }
+ LLVM_DEBUG(dbgs() << " -> Tie\n");
+ return 0;
+}
+
+void HardwareUnitInfo::updatePrioritySUsWith(SUnit *Cand,
+ bool IsCloseToRegPressureLimit) {
if (PrioritySUs.empty()) {
- PrioritySUs.insert(SU);
+ PrioritySUs.insert(Cand);
return;
}
- unsigned SUDepth = SU->getDepth();
- unsigned CurrDepth = (*PrioritySUs.begin())->getDepth();
- if (SUDepth > CurrDepth)
+
+ int Decision = 0;
+
+ SUnit *Existing = *PrioritySUs.begin();
+ bool UseRegFree = CoexecRegFreeProximity == RegFreeProximityMode::Always ||
+ (CoexecRegFreeProximity == RegFreeProximityMode::Auto &&
+ IsCloseToRegPressureLimit);
+
+ LLVM_DEBUG(dbgs() << " updatePrioritySUs: SU(" << Cand->NodeNum
+ << ") vs existing SU(" << Existing->NodeNum << ")"
+ << " mode=" << (UseRegFree ? "RegFreeProximity" : "Depth")
+ << "\n");
+
+ if (!UseRegFree)
+ Decision = compareDepth(Cand, Existing);
+ else
+ Decision = compareRegFreeProximity(Cand, Existing);
+
+ LLVM_DEBUG(dbgs() << " decision=" << Decision
+ << (Decision > 0 ? " (candidate better)"
+ : Decision < 0 ? " (existing better)"
+ : " (tie)")
+ << "\n");
+
+ if (Decision < 0)
return;
- if (SUDepth == CurrDepth) {
- PrioritySUs.insert(SU);
+ if (Decision == 0) {
+ PrioritySUs.insert(Cand);
return;
}
- // SU is lower depth and should be prioritized.
PrioritySUs.clear();
- PrioritySUs.insert(SU);
+ PrioritySUs.insert(Cand);
+}
+
+void HardwareUnitInfo::rebuildPrioritySUs(bool IsCloseToRegPressureLimit) {
+ if (AllSUs.empty())
+ return;
+ PrioritySUs.clear();
+ for (auto *SU : AllSUs)
+ updatePrioritySUsWith(SU, IsCloseToRegPressureLimit);
}
-void HardwareUnitInfo::markScheduled(SUnit *SU, unsigned BlockingCycles) {
+void HardwareUnitInfo::insert(SUnit *SU, unsigned BlockingCycles,
+ bool IsCloseToRegPressureLimit) {
+ if (!AllSUs.insert(SU))
+ llvm_unreachable("HardwareUnit already contains SU!");
+
+ TotalCycles += BlockingCycles;
+
+ updatePrioritySUsWith(SU, IsCloseToRegPressureLimit);
+}
+
+void HardwareUnitInfo::markScheduled(SUnit *SU, unsigned BlockingCycles,
+ bool IsCloseToRegPressureLimit) {
// We may want to ignore some HWUIs (e.g. InstructionFlavor::Other). To do so,
// we just clear the HWUI. However, we still have instructions which map to
// this HWUI. Don't bother managing the state for these HWUI.
@@ -506,27 +650,8 @@ void HardwareUnitInfo::markScheduled(SUnit *SU, unsigned BlockingCycles) {
if (AllSUs.empty())
return;
- if (PrioritySUs.empty()) {
- for (auto SU : AllSUs) {
- if (PrioritySUs.empty()) {
- PrioritySUs.insert(SU);
- continue;
- }
- unsigned SUDepth = SU->getDepth();
- unsigned CurrDepth = (*PrioritySUs.begin())->getDepth();
- if (SUDepth > CurrDepth)
- continue;
-
- if (SUDepth == CurrDepth) {
- PrioritySUs.insert(SU);
- continue;
- }
-
- // SU is lower depth and should be prioritized.
- PrioritySUs.clear();
- PrioritySUs.insert(SU);
- }
- }
+ if (PrioritySUs.empty())
+ rebuildPrioritySUs(IsCloseToRegPressureLimit);
}
void HardwareUnitInfo::finalizeCycles() {
@@ -599,7 +724,7 @@ void CandidateHeuristics::updateForScheduling(SUnit *SU) {
HardwareUnitInfo *HWUI =
getHWUIFromFlavor(classifyFlavor(*SU->getInstr(), *SII));
assert(HWUI);
- HWUI->markScheduled(SU, getHWUICyclesForSU(SU));
+ HWUI->markScheduled(SU, getHWUICyclesForSU(SU), IsCloseToRegPressureLimit);
}
void CandidateHeuristics::initialize(ScheduleDAGMI *SchedDAG,
@@ -718,7 +843,8 @@ void CandidateHeuristics::collectRegionSummary() {
for (auto &SU : DAG->SUnits) {
MachineInstr *MI = SU.getInstr();
const InstructionFlavor Flavor = classifyFlavor(*MI, *SII);
- HWUInfo[static_cast<int>(Flavor)].insert(&SU, getHWUICyclesForSU(&SU));
+ HWUInfo[static_cast<int>(Flavor)].insert(&SU, getHWUICyclesForSU(&SU),
+ IsCloseToRegPressureLimit);
unsigned CarriedLatency = getCarriedLatency(&SU);
if (CarriedLatency)
CarriedLatencies[MI] = CarriedLatency;
@@ -1190,6 +1316,19 @@ void AMDGPUCoExecSchedStrategy::pickNodeFromQueue(
}
}
+ constexpr unsigned MaxVGPRPressureInc = 16;
+ constexpr unsigned MaxVGPRPressureIncFactor = 2;
+ const bool IsCloseToRegPressureLimit =
+ DAG->isTrackingPressure() &&
+ VGPRPressure + MaxVGPRPressureIncFactor * MaxVGPRPressureInc >=
+ VGPRExcessLimit;
+ LLVM_DEBUG(dbgs() << "IsCloseToRegPressureLimit=" << IsCloseToRegPressureLimit
+ << " (VGPR=" << VGPRPressure << " limit=" << VGPRExcessLimit
+ << ")\n");
+ Heurs.setIsCloseToRegPressureLimit(IsCloseToRegPressureLimit);
+ if (IsCloseToRegPressureLimit)
+ Heurs.rebuildAllPrioritySUs();
+
auto EvaluateQueue = [&](ReadyQueue &Q, bool FromPending) {
for (SUnit *SU : Q) {
SchedCandidate TryCand(ZonePolicy);
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.h b/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.h
index 6ce5aaeda8c82..38dbe9a3ee115 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCoExecSchedStrategy.h
@@ -107,7 +107,20 @@ class HardwareUnitInfo {
/// behavior which is not modelled in the compiler.
unsigned BufferCycles = 0;
+ /// Compares two SUnits by depth (lower depth = higher priority for top-down).
+ /// \returns -1 if Candidate is worse, 0 if equal, 1 if Candidate is better.
+ int compareDepth(SUnit *Candidate, SUnit *Existing) const;
+
+ /// Compares two SUnits by proximity to freeing a register.
+ /// \returns -1 if Candidate is worse, 0 if equal, 1 if Candidate is better.
+ int compareRegFreeProximity(SUnit *Candidate, SUnit *Existing) const;
+
+ /// Try to update PrioritySUs with a new \p SU.
+ void updatePrioritySUsWith(SUnit *SU, bool IsCloseToRegPressureLimit = false);
+
public:
+ /// Rebuild PrioritySUs from AllSUs using the given pressure flag.
+ void rebuildPrioritySUs(bool IsCloseToRegPressureLimit);
HardwareUnitInfo() {}
unsigned size() { return AllSUs.size(); }
@@ -192,11 +205,13 @@ class HardwareUnitInfo {
SUnit *getNextTargetSU(bool LookDeep = false) const;
/// Insert the \p SU into AllSUs and account its \p BlockingCycles into
/// the TotalCycles. This maintains the list of PrioritySUs.
- void insert(SUnit *SU, unsigned BlockingCycles);
+ void insert(SUnit *SU, unsigned BlockingCycles,
+ bool IsCloseToRegPressureLimit);
/// Update the state for \p SU being scheduled by removing it from the AllSUs
/// and reducing its \p BlockingCycles from the TotalCycles. This maintains
/// the list of PrioritySUs.
- void markScheduled(SUnit *SU, unsigned BlockingCycles);
+ void markScheduled(SUnit *SU, unsigned BlockingCycles,
+ bool IsCloseToRegPressureLimit);
/// After we've collected all the region pressure for this HWUI, correct for
/// any specifics of the behavior of this resource. For example, if the
/// HardwareUnit can hold N instructions simultaneously, then there is no
@@ -255,6 +270,10 @@ class CandidateHeuristics {
StallCosts getStallCosts(SUnit *SU, SchedBoundary &Zone);
+ /// Controls whether or not the KillProximity heuristic is used when
+ /// selecting the next candidate SU for scheduling.
+ bool IsCloseToRegPressureLimit = false;
+
public:
CandidateHeuristics() = default;
@@ -310,6 +329,15 @@ class CandidateHeuristics {
SchedBoundary *Zone) const;
void dumpRegionSummary();
+
+ void setIsCloseToRegPressureLimit(bool Value) {
+ IsCloseToRegPressureLimit = Value;
+ }
+
+ void rebuildAllPrioritySUs() {
+ for (auto &HWUI : HWUInfo)
+ HWUI.rebuildPrioritySUs(IsCloseToRegPressureLimit);
+ }
};
class AMDGPUCoExecSchedStrategy final : public GCNSchedStrategy {
diff --git a/llvm/test/CodeGen/AMDGPU/coexec-reg-free-proximity.mir b/llvm/test/CodeGen/AMDGPU/coexec-reg-free-proximity.mir
new file mode 100644
index 0000000000000..37f89fbb92f31
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/coexec-reg-free-proximity.mir
@@ -0,0 +1,56 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -stress-regalloc=8 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx950 -run-pass=machine-scheduler -amdgpu-sched-strategy=coexec %s -o - | FileCheck %s
+
+# Test the register free proximity heuristic in the coexec scheduler.
+#
+# %s0 and %s1 are used by add_a1, add_a2, add_a3 (3 users each).
+# %s2 and %s3 are used by add_b1, add_b2, add_b3 (3 users each).
+# All six adds are independent and ready simultaneously.
+#
+# Without the heuiristic: scheduled in program order (a1, b1, a2, b2, a3, b3).
+# With the heuristic: it groups adds by operand so that registers are freed
+# sooner.
+
+--- |
+ define void @free_proximity() #0 { ret void }
+ attributes #0 = { "amdgpu-waves-per-eu"="1,1" "amdgpu-flat-work-group-size"="64,64" }
+...
+
+---
+
+name: free_proximity
+tracksRegLiveness: true
+machineFunctionInfo:
+ isEntryFunction: true
+body: |
+ bb.0:
+ liveins: $vgpr0, $vgpr1, $vgpr2, $vgpr3, $vgpr4
+
+ ; CHECK-LABEL: name: free_proximity
+ ; CHECK: liveins: $vgpr0, $vgpr1, $vgpr2, $vgpr3, $vgpr4
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %s0:vgpr_32 = COPY $vgpr0
+ ; CHECK-NEXT: %s1:vgpr_32 = COPY $vgpr1
+ ; CHECK-NEXT: %s2:vgpr_32 = COPY $vgpr2
+ ; CHECK-NEXT: %s3:vgpr_32 = COPY $vgpr3
+ ; CHECK-NEXT: %add_a1:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ ; CHECK-NEXT: %add_a2:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ ; CHECK-NEXT: %add_a3:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ ; CHECK-NEXT: %add_b1:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+ ; CHECK-NEXT: %add_b2:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+ ; CHECK-NEXT: %add_b3:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+ ; CHECK-NEXT: S_ENDPGM 0, implicit %add_a1, implicit %add_b1, implicit %add_a2, implicit %add_b2, implicit %add_a3, implicit %add_b3
+ %s0:vgpr_32 = COPY $vgpr0
+ %s1:vgpr_32 = COPY $vgpr1
+ %s2:vgpr_32 = COPY $vgpr2
+ %s3:vgpr_32 = COPY $vgpr3
+
+ %add_a1:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ %add_b1:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+ %add_a2:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ %add_b2:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+ %add_a3:vgpr_32 = V_ADD_U32_e32 %s0, %s1, implicit $exec
+ %add_b3:vgpr_32 = V_ADD_U32_e32 %s2, %s3, implicit $exec
+
+ S_ENDPGM 0, implicit %add_a1, implicit %add_b1, implicit %add_a2, implicit %add_b2, implicit %add_a3, implicit %add_b3
+...
More information about the llvm-commits
mailing list