[llvm-branch-commits] [llvm] [AMDGPU] Count WMMA scale copies against both operands (PR #234104)
Anshil Gandhi via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Oct 22 10:54:35 PDT 2026
https://github.com/gandhi56 updated https://github.com/llvm/llvm-project/pull/234104
>From 1c0efbaae4a4e1b9b797ad83637f7de3b2eceb4f Mon Sep 17 00:00:00 2001
From: Anshil Gandhi <Anshil.Gandhi at amd.com>
Date: Wed, 21 Oct 2026 03:24:33 -0500
Subject: [PATCH] [AMDGPU] Count WMMA scale copies against both operands
Copying a shared scale while its partner stays an SGPR adds a va_ssrc
issue without removing one. Choose copies from both scale operands.
---
.../Target/AMDGPU/SIMaterializeWMMAScale.cpp | 238 +++++++++++++++---
.../AMDGPU/wmma-scale-sgpr-materialize.mir | 68 ++++-
2 files changed, 261 insertions(+), 45 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/SIMaterializeWMMAScale.cpp b/llvm/lib/Target/AMDGPU/SIMaterializeWMMAScale.cpp
index 0c8ae8a802d1a..7a1857fc4358a 100644
--- a/llvm/lib/Target/AMDGPU/SIMaterializeWMMAScale.cpp
+++ b/llvm/lib/Target/AMDGPU/SIMaterializeWMMAScale.cpp
@@ -14,12 +14,16 @@
/// scalar-source slot and a run of such WMMA ops serializes. One cross-bank
/// copy pays for that slot once; later WMMA ops then read a VGPR.
///
-/// A single WMMA is left on the SGPR: the copy would still occupy va_ssrc and
-/// would add a VALU op. A scale defined outside a loop is copied at its def
-/// even for one in-loop reader, so the loop body does not hit the counter. A
-/// copy inside a loop that modifies exec stays in that loop. The VGPR latches
-/// each lane as it drops out, so a use after the loop keeps the SGPR and reads
-/// the reconverged scalar. A subregister use is copied at that width, so two
+/// A WMMA occupies va_ssrc when either scale is an SGPR. That includes a
+/// physical register and a subregister that cannot be copied. Copies are
+/// chosen together: (S, T0) and (S, T1) stay on the SGPR, because copying
+/// only S leaves both instructions on the counter and the copy adds one.
+/// A set of copies is emitted when the instructions that then read no SGPR
+/// scale outnumber the copies. A tie is kept when every copy sits outside a
+/// loop that would otherwise repeat that many reads. A copy inside a loop
+/// that modifies exec stays in that loop. The VGPR latches each lane as it
+/// drops out, so a use after the loop keeps the SGPR and reads the
+/// reconverged scalar. A subregister use is copied at that width, so two
/// lanes of one SGPR tuple stay independent.
//
//===----------------------------------------------------------------------===//
@@ -30,13 +34,15 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/InitializePasses.h"
+#include "llvm/Support/MathExtras.h"
using namespace llvm;
@@ -48,8 +54,93 @@ STATISTIC(NumWMMAScalesMaterialized,
/// VALU SGPR-source issue slots. depctr_va_ssrc is a 1-bit counter.
static constexpr unsigned VaSsrcDepth = 1;
+/// Exact search over copy sets. Larger components fall back to single-key
+/// decisions, which never copy a scale that still leaves its reader on an SGPR.
+static constexpr unsigned MaxExactCopySet = 16;
+
namespace {
+using ScaleKey = std::pair<Register, unsigned>;
+
+struct ScaleCandidate {
+ SmallVector<unsigned, 2> KeyIdx;
+ bool LoopHoist = false;
+};
+
+/// Pick a subset of \p CompKeys. An instruction is cleared only when every
+/// key it needs is copied, so a shared scale is not copied while its partner
+/// stays on an SGPR. \p Selected is indexed by the global key id.
+void chooseCopies(ArrayRef<ScaleCandidate> Cands, ArrayRef<unsigned> CompKeys,
+ SmallVectorImpl<bool> &Selected) {
+ const unsigned K = CompKeys.size();
+ DenseMap<unsigned, unsigned> Local;
+ for (auto [I, Key] : enumerate(CompKeys))
+ Local[Key] = I;
+
+ auto consider = [&](unsigned Cost, unsigned Benefit, bool AllHoist,
+ function_ref<void()> Accept) {
+ if (Benefit == 0 || Cost == 0)
+ return;
+ const bool Pays = Benefit > Cost * VaSsrcDepth ||
+ (Benefit == Cost * VaSsrcDepth && AllHoist);
+ if (Pays)
+ Accept();
+ };
+
+ if (K > MaxExactCopySet) {
+ for (unsigned L = 0; L < K; ++L) {
+ unsigned Benefit = 0;
+ bool AllHoist = true;
+ for (const ScaleCandidate &C : Cands) {
+ if (C.KeyIdx.size() != 1 || Local[C.KeyIdx[0]] != L)
+ continue;
+ ++Benefit;
+ AllHoist &= C.LoopHoist;
+ }
+ consider(1, Benefit, AllHoist, [&] { Selected[CompKeys[L]] = true; });
+ }
+ return;
+ }
+
+ int BestScore = 0;
+ bool BestHoist = false;
+ unsigned BestCost = 0;
+ uint32_t BestMask = 0;
+ for (uint32_t Mask = 1; Mask < (1u << K); ++Mask) {
+ const unsigned Cost = llvm::popcount(Mask);
+ unsigned Benefit = 0;
+ bool AllHoist = true;
+ for (const ScaleCandidate &C : Cands) {
+ bool Covered = true;
+ for (unsigned Key : C.KeyIdx)
+ Covered &= (Mask & (1u << Local[Key])) != 0;
+ if (!Covered)
+ continue;
+ ++Benefit;
+ AllHoist &= C.LoopHoist;
+ }
+ if (Benefit == 0)
+ continue;
+ const int Score =
+ static_cast<int>(Benefit) - static_cast<int>(Cost * VaSsrcDepth);
+ const bool Hoist = Score == 0 && AllHoist;
+ if (Score < 0 || (Score == 0 && !Hoist))
+ continue;
+ const bool Better = Score != BestScore ? Score > BestScore
+ : Hoist != BestHoist ? Hoist
+ : Cost < BestCost;
+ if (!Better)
+ continue;
+ BestScore = Score;
+ BestHoist = Hoist;
+ BestCost = Cost;
+ BestMask = Mask;
+ }
+ for (unsigned L = 0; L < K; ++L)
+ if (BestMask & (1u << L))
+ Selected[CompKeys[L]] = true;
+}
+
class SIMaterializeWMMAScale {
struct ScaleUse {
MachineInstr *MI;
@@ -145,29 +236,6 @@ bool SIMaterializeWMMAScale::materialize(Register Reg, unsigned SubReg,
if (!Def || Def->isTerminator())
return false;
- const MachineBasicBlock *DefBB = Def->getParent();
- SmallVector<ScaleUse, 4> Rewritable;
- for (const ScaleUse &Use : Uses) {
- if (!crossesUnsafeDivergentLoop(DefBB, Use.MI->getParent()))
- Rewritable.push_back(Use);
- }
- if (Rewritable.empty())
- return false;
-
- SmallPtrSet<MachineInstr *, 4> Readers;
- bool DefOutsideReaderLoop = false;
- for (const ScaleUse &Use : Rewritable) {
- Readers.insert(Use.MI);
- if (const MachineLoop *UseLoop = MLI->getLoopFor(Use.MI->getParent()))
- DefOutsideReaderLoop |= !UseLoop->contains(DefBB);
- }
-
- // One SGPR read occupies the only va_ssrc slot. A copy does too, so it
- // wins only when more than one WMMA would otherwise take that slot, or
- // when a loop would take it on every iteration.
- if (Readers.size() <= VaSsrcDepth && !DefOutsideReaderLoop)
- return false;
-
// Size the VGPR from the subregister, not the surrounding tuple. Copying
// the whole SGPR would move lanes this use does not read.
unsigned BitWidth =
@@ -175,17 +243,20 @@ bool SIMaterializeWMMAScale::materialize(Register Reg, unsigned SubReg,
: TRI->getRegSizeInBits(*MRI->getRegClass(Reg)).getFixedValue();
const TargetRegisterClass *VRC =
TRI->getAlignedLo256VGPRClassForBitWidth(BitWidth);
- if (!VRC)
+ if (!VRC || Uses.empty())
return false;
Register VReg = MRI->createVirtualRegister(VRC);
+ const MachineBasicBlock *DefBB = Def->getParent();
MachineBasicBlock &MBB = *Def->getParent();
MachineBasicBlock::iterator InsertPt =
Def->isPHI() ? MBB.getFirstNonPHI() : std::next(Def->getIterator());
BuildMI(MBB, InsertPt, Def->getDebugLoc(), TII->get(TargetOpcode::COPY), VReg)
.addUse(Reg, {}, SubReg);
- for (const ScaleUse &Use : Rewritable) {
+ for (const ScaleUse &Use : Uses) {
+ if (crossesUnsafeDivergentLoop(DefBB, Use.MI->getParent()))
+ continue;
MachineOperand &MO = Use.MI->getOperand(Use.OpNo);
MO.setReg(VReg);
MO.setSubReg(0);
@@ -210,34 +281,119 @@ bool SIMaterializeWMMAScale::run(MachineFunction &MF,
MLI = &Loops;
ExecLoopCache.clear();
- // Insertion order keeps copy placement deterministic. Subregister lanes of
- // one tuple are separate scales.
- MapVector<std::pair<Register, unsigned>, SmallVector<ScaleUse, 4>> ScaleUses;
+ // Insertion order keeps copy placement deterministic. Each lane is its own
+ // scale. A physical SGPR, or a subregister with no legal VGPR copy, still
+ // occupies va_ssrc and is not a key.
+ MapVector<ScaleKey, SmallVector<ScaleUse, 4>> ScaleUses;
+ DenseMap<ScaleKey, unsigned> KeyIndex;
+ SmallVector<const MachineBasicBlock *, 8> KeyDefBB;
+ SmallVector<ScaleCandidate, 8> Candidates;
+
+ auto canCopy = [&](const MachineOperand &MO, const MachineBasicBlock *UseBB,
+ const MachineBasicBlock *&DefBB) {
+ if (!MO.isReg() || !MO.getReg().isVirtual())
+ return false;
+ MachineInstr *Def = MRI->getVRegDef(MO.getReg());
+ if (!Def || Def->isTerminator())
+ return false;
+ unsigned SubReg = MO.getSubReg();
+ unsigned BitWidth =
+ SubReg ? TRI->getSubRegIdxSize(SubReg)
+ : TRI->getRegSizeInBits(*MRI->getRegClass(MO.getReg()))
+ .getFixedValue();
+ if (!TRI->getAlignedLo256VGPRClassForBitWidth(BitWidth))
+ return false;
+ DefBB = Def->getParent();
+ return !crossesUnsafeDivergentLoop(DefBB, UseBB);
+ };
+
for (MachineBasicBlock &MBB : MF) {
for (MachineInstr &MI : MBB) {
if (!SIInstrInfo::isWMMA(MI) ||
!AMDGPU::hasNamedOperand(MI.getOpcode(), AMDGPU::OpName::scale_src0))
continue;
+ bool FixedSGPR = false;
+ SmallVector<unsigned, 2> Keys;
for (AMDGPU::OpName Name :
{AMDGPU::OpName::scale_src0, AMDGPU::OpName::scale_src1}) {
int Idx = AMDGPU::getNamedOperandIdx(MI.getOpcode(), Name);
if (Idx < 0)
continue;
MachineOperand &MO = MI.getOperand(Idx);
- if (!MO.isReg() || !MO.getReg().isVirtual())
+ if (!MO.isReg() || !MO.getReg() || !TRI->isSGPRReg(*MRI, MO.getReg()))
continue;
- if (!TRI->isSGPRReg(*MRI, MO.getReg()))
+ const MachineBasicBlock *DefBB = nullptr;
+ if (!canCopy(MO, MI.getParent(), DefBB)) {
+ FixedSGPR = true;
continue;
- ScaleUses[{MO.getReg(), MO.getSubReg()}].push_back(
- {&MI, static_cast<unsigned>(Idx)});
+ }
+ ScaleKey Key{MO.getReg(), MO.getSubReg()};
+ auto [It, Inserted] = KeyIndex.try_emplace(Key, KeyIndex.size());
+ if (Inserted) {
+ ScaleUses[Key] = {};
+ KeyDefBB.push_back(DefBB);
+ }
+ ScaleUses[Key].push_back({&MI, static_cast<unsigned>(Idx)});
+ if (!is_contained(Keys, It->second))
+ Keys.push_back(It->second);
+ }
+ if (FixedSGPR || Keys.empty())
+ continue;
+
+ const MachineLoop *UseLoop = MLI->getLoopFor(MI.getParent());
+ bool Hoist = UseLoop != nullptr;
+ if (Hoist) {
+ for (unsigned Key : Keys)
+ Hoist &= !UseLoop->contains(KeyDefBB[Key]);
+ }
+ Candidates.push_back({std::move(Keys), Hoist});
+ }
+ }
+
+ SmallVector<bool, 8> Selected(KeyIndex.size(), false);
+ if (!Candidates.empty()) {
+ SmallVector<unsigned, 8> Parent(KeyIndex.size());
+ for (unsigned I = 0, E = Parent.size(); I != E; ++I)
+ Parent[I] = I;
+ auto findRoot = [&](unsigned X) {
+ while (Parent[X] != X) {
+ Parent[X] = Parent[Parent[X]];
+ X = Parent[X];
}
+ return X;
+ };
+ for (const ScaleCandidate &C : Candidates) {
+ unsigned Root = findRoot(C.KeyIdx.front());
+ for (unsigned Key : C.KeyIdx)
+ Parent[findRoot(Key)] = Root;
+ }
+
+ MapVector<unsigned, SmallVector<unsigned, 8>> ByRoot;
+ for (auto [I, C] : enumerate(Candidates))
+ ByRoot[findRoot(C.KeyIdx.front())].push_back(I);
+
+ for (const auto &[Root, CandIdxs] : ByRoot) {
+ (void)Root;
+ SmallVector<unsigned, 8> CompKeys;
+ SmallVector<ScaleCandidate, 8> CompCands;
+ for (unsigned CI : CandIdxs) {
+ CompCands.push_back(Candidates[CI]);
+ for (unsigned Key : Candidates[CI].KeyIdx)
+ if (!is_contained(CompKeys, Key))
+ CompKeys.push_back(Key);
+ }
+ chooseCopies(CompCands, CompKeys, Selected);
}
}
bool Changed = false;
- for (const auto &[Key, Uses] : ScaleUses)
- Changed |= materialize(Key.first, Key.second, Uses);
+ unsigned Idx = 0;
+ for (const auto &[Key, Uses] : ScaleUses) {
+ if (Selected[Idx])
+ Changed |= materialize(Key.first, Key.second, Uses);
+ ++Idx;
+ }
return Changed;
}
diff --git a/llvm/test/CodeGen/AMDGPU/wmma-scale-sgpr-materialize.mir b/llvm/test/CodeGen/AMDGPU/wmma-scale-sgpr-materialize.mir
index a4ddda70c62cb..caaa67a1d1794 100644
--- a/llvm/test/CodeGen/AMDGPU/wmma-scale-sgpr-materialize.mir
+++ b/llvm/test/CodeGen/AMDGPU/wmma-scale-sgpr-materialize.mir
@@ -2,10 +2,12 @@
# RUN: llc -mtriple=amdgpu12.50 -passes=si-materialize-wmma-scale -o - %s | FileCheck %s
# Two WMMA instructions reading the same SGPR scale must share one VGPR copy.
-# Both scale operands of a single WMMA count as one va_ssrc issue, so that
-# case stays on the SGPR. A scale defined outside a loop is copied there.
-# V_WMMA_SCALE16_* scale operands are 64 bits. A subregister use is copied at
-# that lane's width; the other lane of the tuple stays an SGPR.
+# Both scale operands of one WMMA are one va_ssrc issue. (S, T0) and (S, T1)
+# stay on SGPRs: copying only S removes nothing, and copying all three adds
+# more issues than it removes. A physical SGPR partner cannot be rewritten, so
+# it keeps the shared scale on the counter too. A scale defined outside a loop
+# is copied there. V_WMMA_SCALE16_* scale operands are 64 bits. A subregister
+# use is copied at that lane's width; the other lane of the tuple stays an SGPR.
---
name: two_wmma_shared_sgpr_scale
@@ -406,3 +408,61 @@ body: |
%7:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %1, %2, 0, %6, %0.sub1, %4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
S_ENDPGM 0
...
+
+---
+name: mixed_shared_sgpr_scales
+isSSA: true
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $sgpr0, $sgpr1, $sgpr2
+ ; Two WMMAs read (S, T0) and (S, T1). Copying S leaves both on va_ssrc.
+ ; CHECK-LABEL: name: mixed_shared_sgpr_scales
+ ; CHECK: liveins: $sgpr0, $sgpr1, $sgpr2
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %s:sreg_32 = COPY $sgpr0
+ ; CHECK-NEXT: %t0:sreg_32 = COPY $sgpr1
+ ; CHECK-NEXT: %t1:sreg_32 = COPY $sgpr2
+ ; CHECK-NEXT: %a:vreg_512_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: %b:vreg_512_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: %c:vreg_256_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: early-clobber %w0:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %c, %s, %t0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ ; CHECK-NEXT: early-clobber %w1:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %w0, %s, %t1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ ; CHECK-NEXT: S_ENDPGM 0
+ %s:sreg_32 = COPY $sgpr0
+ %t0:sreg_32 = COPY $sgpr1
+ %t1:sreg_32 = COPY $sgpr2
+ %a:vreg_512_align2 = IMPLICIT_DEF
+ %b:vreg_512_align2 = IMPLICIT_DEF
+ %c:vreg_256_align2 = IMPLICIT_DEF
+ %w0:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %c, %s, %t0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ %w1:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %w0, %s, %t1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ S_ENDPGM 0
+...
+
+---
+name: shared_scale_with_physical_sgpr
+isSSA: true
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $sgpr0, $sgpr1, $sgpr2
+ ; $sgpr1 and $sgpr2 cannot be copied. Both WMMAs still occupy va_ssrc.
+ ; CHECK-LABEL: name: shared_scale_with_physical_sgpr
+ ; CHECK: liveins: $sgpr0, $sgpr1, $sgpr2
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %s:sreg_32 = COPY $sgpr0
+ ; CHECK-NEXT: %a:vreg_512_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: %b:vreg_512_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: %c:vreg_256_align2 = IMPLICIT_DEF
+ ; CHECK-NEXT: early-clobber %w0:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %c, %s, $sgpr1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ ; CHECK-NEXT: early-clobber %w1:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %w0, %s, $sgpr2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ ; CHECK-NEXT: S_ENDPGM 0
+ %s:sreg_32 = COPY $sgpr0
+ %a:vreg_512_align2 = IMPLICIT_DEF
+ %b:vreg_512_align2 = IMPLICIT_DEF
+ %c:vreg_256_align2 = IMPLICIT_DEF
+ %w0:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %c, %s, $sgpr1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ %w1:vreg_256_align2 = V_WMMA_SCALE_F32_16X16X128_F8F6F4_f8_f8_w32_threeaddr %a, %b, 0, %w0, %s, $sgpr2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+ S_ENDPGM 0
+...
More information about the llvm-branch-commits
mailing list