[llvm] [AMDGPU] Weight spills over copies in RewriteMFMAFormStage cost model (PR #225372)
Romanov Vlad via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 22 04:22:01 PDT 2026
https://github.com/romanovvlad created https://github.com/llvm/llvm-project/pull/225372
getRewriteCost compared a spill/restore pair and a register-to-register
bridge copy at the same price, so a rewrite that removed a handful of
spills was rejected whenever it needed a comparable number of copies.
Scale the spill term by 10 to reflect that a spill goes to scratch.
Also credit each AGPR the rewrite puts to use. We want the AGPR bank
reasonably used and as many ArchVGPRs as possible freed up, which gives
the scheduler more room than the spill term alone can express. The
credit stops at 7/8 of the AGPR budget so that the AGPR bank is not
filled to the point where the spilling merely moves from one bank to
the other.
>From c0515d079fc884cbc9ae7120ffe37242ded020dc Mon Sep 17 00:00:00 2001
From: Vlad <Vladislav.Romanov at amd.com>
Date: Tue, 22 Sep 2026 05:49:18 -0500
Subject: [PATCH] [AMDGPU] Weight spills over copies in RewriteMFMAFormStage
cost model
getRewriteCost compared a spill/restore pair and a register-to-register
bridge copy at the same price, so a rewrite that removed a handful of
spills was rejected whenever it needed a comparable number of copies.
Scale the spill term by 10 to reflect that a spill goes to scratch.
Also credit each AGPR the rewrite puts to use. We want the AGPR bank
reasonably used and as many ArchVGPRs as possible freed up, which gives
the scheduler more room than the spill term alone can express. The credit
stops at 7/8 of the AGPR budget so that the AGPR bank is not filled to the
point where the spilling merely moves from one bank to the other.
---
llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 32 +++-
.../CodeGen/AMDGPU/rewrite-mfma-form-cost.mir | 176 ++++++++++++++++++
2 files changed, 206 insertions(+), 2 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/rewrite-mfma-form-cost.mir
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 6a1250f672b50..55d34f407d266 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -103,6 +103,13 @@ static cl::opt<bool> DisableRewriteMFMAFormSchedStage(
"amdgpu-disable-rewrite-mfma-form-sched-stage", cl::Hidden,
cl::desc("Disable rewrite mfma rewrite scheduling stage"), cl::init(true));
+/// Cost of a spill/restore pair relative to TargetRegisterClass::getCopyCost.
+/// A spill goes to scratch memory, a bridge copy does not.
+static constexpr int64_t SpillVsCopyCostRatio = 10;
+
+/// Credit per AGPR the rewrite puts to use, on the scale of one bridge copy.
+static constexpr int64_t AGPRUtilizationBonus = 4;
+
namespace {
struct VGPRThresholdParser : public cl::parser<unsigned> {
@@ -2544,6 +2551,9 @@ int64_t RewriteMFMAFormStage::getRewriteCost(
unsigned AGPRThreshold = MaxVectorRegs.second;
unsigned CombinedThreshold = ST.getMaxNumVGPRs(MF);
+ // Peak AGPR pressure over the regions we are trying to fix.
+ unsigned MaxAGPRAfter = 0;
+
for (unsigned Region = 0; Region < DAG.Regions.size(); Region++) {
if (!RegionsWithExcessArchVGPR[Region])
continue;
@@ -2560,6 +2570,8 @@ int64_t RewriteMFMAFormStage::getRewriteCost(
unsigned SpillCostAfter = PressureAfter.getVGPRSpills(
MF, ArchVGPRThreshold, AGPRThreshold, CombinedThreshold);
+ MaxAGPRAfter = std::max(MaxAGPRAfter, PressureAfter.getAGPRNum());
+
uint64_t BlockFreq =
MBFI->getBlockFreq(DAG.Regions[Region].first->getParent())
.getFrequency();
@@ -2572,7 +2584,8 @@ int64_t RewriteMFMAFormStage::getRewriteCost(
// This assumes perfect spilling / splitting -- using one spill / copy
// instruction and one restoreFrom / copy for each excess register,
- int64_t SpillCost = ((int)SpillCostAfter - (int)SpillCostBefore) * 2;
+ int64_t SpillCost =
+ ((int)SpillCostAfter - (int)SpillCostBefore) * 2 * SpillVsCopyCostRatio;
// Also account for the block frequency.
if (RelativeFreqIsDenom)
@@ -2621,13 +2634,28 @@ int64_t RewriteMFMAFormStage::getRewriteCost(
}
}
+ // Reward AGPR utilization. We want the AGPR bank reasonably used and as many
+ // ArchVGPRs as possible freed up, which gives the scheduler more room than
+ // the spill term alone can express. The credit stops at a soft limit so that
+ // the AGPR bank is not filled to the point where the spilling merely moves
+ // from one bank to the other.
+ int64_t AGPRBonus = 0;
+ unsigned AGPRSoftLimit = AGPRThreshold * 7 / 8;
+ if (MaxAGPRAfter && MaxAGPRAfter <= AGPRSoftLimit)
+ AGPRBonus = -static_cast<int64_t>(MaxAGPRAfter) * AGPRUtilizationBonus;
+
+ LLVM_DEBUG(dbgs() << "RewriteMFMA cost: spill=" << Cost
+ << " copy=" << CopyCost << " agpr=" << AGPRBonus
+ << " (peak AGPR " << MaxAGPRAfter << "/" << AGPRSoftLimit
+ << ") total=" << (Cost + CopyCost + AGPRBonus) << '\n');
+
// Reset the classes that were changed to AGPR for better register bank
// analysis. We must do rewriting after copy-insertion, as some defs of the
// register may require VGPR. Additionally, if we bail out and don't perform
// the rewrite then these need to be restored anyway.
resetRewriteCandsToVGPR(RewriteCands);
- return Cost + CopyCost;
+ return Cost + CopyCost + AGPRBonus;
}
bool RewriteMFMAFormStage::rewrite(
diff --git a/llvm/test/CodeGen/AMDGPU/rewrite-mfma-form-cost.mir b/llvm/test/CodeGen/AMDGPU/rewrite-mfma-form-cost.mir
new file mode 100644
index 0000000000000..62237bd9ac62b
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/rewrite-mfma-form-cost.mir
@@ -0,0 +1,176 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=amdgpu9.50-amd-amdhsa -run-pass=machine-scheduler -amdgpu-disable-rewrite-mfma-form-sched-stage=false -o - %s | FileCheck %s
+
+# Cost model checks for RewriteMFMAFormStage. Both functions were rejected by
+# the unweighted cost model because the bridge copies outweighed the modelled
+# spill reduction.
+
+--- |
+ define void @spill_saving_outweighs_copies() #0 {
+ entry:
+ unreachable
+ }
+
+ define void @agpr_utilization_bonus() #0 {
+ entry:
+ unreachable
+ }
+
+ attributes #0 = { "amdgpu-waves-per-eu"="1,1" "amdgpu-flat-work-group-size"="64,64"}
+...
+
+# Spill saving of 8 registers (cost -160) beats the bridge copies (cost 18).
+
+---
+name: spill_saving_outweighs_copies
+tracksRegLiveness: true
+machineFunctionInfo:
+ isEntryFunction: true
+body: |
+ ; CHECK-LABEL: name: spill_saving_outweighs_copies
+ ; CHECK: bb.0:
+ ; CHECK-NEXT: successors: %bb.1(0x80000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[DEF:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF1:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF2:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF3:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF4:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF5:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: S_NOP 0, implicit-def %8
+ ; CHECK-NEXT: S_NOP 0, implicit-def %9
+ ; CHECK-NEXT: [[DEF6:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: SCHED_BARRIER 0
+ ; CHECK-NEXT: [[DEF7:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF8:%[0-9]+]]:vreg_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.1:
+ ; CHECK-NEXT: successors: %bb.2(0x80000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[DEF9:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF10:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF11:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_:%[0-9]+]]:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 [[DEF9]], [[DEF10]], [[DEF8]], 4, 4, [[DEF11]].sub0, [[DEF7]], 0, 0, implicit $mode, implicit $exec
+ ; CHECK-NEXT: [[COPY:%[0-9]+]]:areg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]]
+ ; CHECK-NEXT: [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_:%[0-9]+]]:areg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64 [[DEF9]], [[DEF10]], [[COPY]], 4, 4, [[DEF11]].sub0, [[DEF7]], 0, 0, implicit $mode, implicit $exec
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.2:
+ ; CHECK-NEXT: undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_128_align2 = V_ADD_U32_e32 [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]].sub0, [[DEF7]], implicit $exec
+ ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vreg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_]]
+ ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vreg_128_align2 = COPY [[COPY1]]
+ ; CHECK-NEXT: SCHED_BARRIER 0
+ ; CHECK-NEXT: [[DEF12:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: KILL [[DEF12]], [[DEF]], [[DEF1]], [[DEF2]], [[DEF3]], [[DEF4]], [[DEF5]], [[DEF6]], [[DEF8]], [[V_ADD_U32_e32_]], [[COPY2]]
+ ; CHECK-NEXT: S_NOP 0, implicit %8, implicit %9
+ ; CHECK-NEXT: S_ENDPGM 0
+ bb.0:
+ %0:vreg_1024 = IMPLICIT_DEF
+ %1:vreg_1024 = IMPLICIT_DEF
+ %2:vreg_1024 = IMPLICIT_DEF
+ %3:vreg_1024 = IMPLICIT_DEF
+ %4:vreg_1024 = IMPLICIT_DEF
+ %5:vreg_1024 = IMPLICIT_DEF
+ %6:vreg_1024 = IMPLICIT_DEF
+ %7:vreg_1024 = IMPLICIT_DEF
+ S_NOP 0, implicit-def %8:av_512
+ S_NOP 0, implicit-def %9:av_512
+ SCHED_BARRIER 0
+ %10:av_128_align2 = IMPLICIT_DEF
+ %11:av_128_align2 = IMPLICIT_DEF
+ %12:vreg_64_align2 = IMPLICIT_DEF
+ %13:vgpr_32 = IMPLICIT_DEF
+ %14:vreg_128_align2 = IMPLICIT_DEF
+
+ bb.1:
+ %15:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %10, %11, %14, 4, 4, %12.sub0, %13, 0, 0, implicit $mode, implicit $exec
+ %16:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %10, %11, %15, 4, 4, %12.sub0, %13, 0, 0, implicit $mode, implicit $exec
+
+ bb.2:
+ undef %17.sub0:vreg_128_align2 = V_ADD_U32_e32 %15.sub0, %13, implicit $exec
+ %18:vreg_128_align2 = COPY %16
+ SCHED_BARRIER 0
+ KILL %0, %1, %2, %3, %4, %5, %6, %7, %14, %17, %18
+ S_NOP 0, implicit %8, implicit %9
+ S_ENDPGM 0
+...
+
+---
+name: agpr_utilization_bonus
+tracksRegLiveness: true
+machineFunctionInfo:
+ isEntryFunction: true
+body: |
+ ; CHECK-LABEL: name: agpr_utilization_bonus
+ ; CHECK: bb.0:
+ ; CHECK-NEXT: successors: %bb.1(0x08000000), %bb.2(0x78000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[DEF:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF1:%[0-9]+]]:av_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF2:%[0-9]+]]:vreg_64_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF3:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF4:%[0-9]+]]:vreg_128_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_:%[0-9]+]]:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 [[DEF]], [[DEF1]], [[DEF4]], 4, 4, [[DEF2]].sub0, [[DEF3]], 0, 0, implicit $mode, implicit $exec
+ ; CHECK-NEXT: [[COPY:%[0-9]+]]:areg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]]
+ ; CHECK-NEXT: $scc = IMPLICIT_DEF
+ ; CHECK-NEXT: [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_:%[0-9]+]]:areg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64 [[DEF]], [[DEF1]], [[COPY]], 4, 4, [[DEF2]].sub0, [[DEF3]], 0, 0, implicit $mode, implicit $exec
+ ; CHECK-NEXT: S_CBRANCH_SCC1 %bb.1, implicit killed $scc
+ ; CHECK-NEXT: S_BRANCH %bb.2
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.1:
+ ; CHECK-NEXT: successors: %bb.2(0x80000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[DEF5:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF6:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF7:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF8:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF9:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF10:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF11:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF12:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: SCHED_BARRIER 0
+ ; CHECK-NEXT: [[DEF13:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+ ; CHECK-NEXT: KILL [[DEF13]], [[DEF5]], [[DEF6]], [[DEF7]], [[DEF8]], [[DEF9]], [[DEF10]], [[DEF11]], [[DEF12]]
+ ; CHECK-NEXT: S_BRANCH %bb.2
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.2:
+ ; CHECK-NEXT: undef [[V_ADD_U32_e32_:%[0-9]+]].sub0:vreg_128_align2 = V_ADD_U32_e32 [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64_]].sub0, [[DEF3]], implicit $exec
+ ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vreg_128_align2 = COPY [[V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_e64_]]
+ ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vreg_128_align2 = COPY [[COPY1]]
+ ; CHECK-NEXT: SCHED_BARRIER 0
+ ; CHECK-NEXT: KILL [[DEF4]], [[V_ADD_U32_e32_]], [[COPY2]]
+ ; CHECK-NEXT: S_ENDPGM 0
+ bb.0:
+ successors: %bb.1(0x08000000), %bb.2(0x78000000)
+ %0:av_128_align2 = IMPLICIT_DEF
+ %1:av_128_align2 = IMPLICIT_DEF
+ %2:vreg_64_align2 = IMPLICIT_DEF
+ %3:vgpr_32 = IMPLICIT_DEF
+ %4:vreg_128_align2 = IMPLICIT_DEF
+ %5:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %0, %1, %4, 4, 4, %2.sub0, %3, 0, 0, implicit $mode, implicit $exec
+ %6:vreg_128_align2 = contract nofpexcept V_MFMA_SCALE_F32_16X16X128_F8F6F4_f4_f4_vgprcd_e64 %0, %1, %5, 4, 4, %2.sub0, %3, 0, 0, implicit $mode, implicit $exec
+ $scc = IMPLICIT_DEF
+ S_CBRANCH_SCC1 %bb.1, implicit killed $scc
+ S_BRANCH %bb.2
+
+ bb.1:
+ successors: %bb.2(0x80000000)
+ %7:vreg_1024 = IMPLICIT_DEF
+ %8:vreg_1024 = IMPLICIT_DEF
+ %9:vreg_1024 = IMPLICIT_DEF
+ %10:vreg_1024 = IMPLICIT_DEF
+ %11:vreg_1024 = IMPLICIT_DEF
+ %12:vreg_1024 = IMPLICIT_DEF
+ %13:vreg_1024 = IMPLICIT_DEF
+ %14:vreg_1024 = IMPLICIT_DEF
+ %15:vreg_1024 = IMPLICIT_DEF
+ SCHED_BARRIER 0
+ KILL %7, %8, %9, %10, %11, %12, %13, %14, %15
+ S_BRANCH %bb.2
+
+ bb.2:
+ undef %16.sub0:vreg_128_align2 = V_ADD_U32_e32 %5.sub0, %3, implicit $exec
+ %17:vreg_128_align2 = COPY %6
+ SCHED_BARRIER 0
+ KILL %4, %16, %17
+ S_ENDPGM 0
+...
More information about the llvm-commits
mailing list