Skip to content

Commit 603f8df

Browse files
arsenmarsenm
committed
AMDGPU: Handle rewriting VGPR MFMA to AGPR with subregister copies
This should address the case where the result isn't fully used, resulting in partial copy bundles from the MFMA result.
1 parent 3776f12 commit 603f8df

6 files changed

+226
-278
lines changed

‎llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp‎

Lines changed: 45 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@
1414
/// MFMA opcode.
1515
///
1616
/// TODO:
17-
/// - Handle SplitKit partial copy bundles, and not just full copy instructions
18-
///
1917
/// - Update LiveIntervals incrementally instead of recomputing from scratch
2018
///
2119
//===----------------------------------------------------------------------===//
@@ -37,6 +35,7 @@ using namespace llvm;
3735
namespace {
3836

3937
class AMDGPURewriteAGPRCopyMFMAImpl {
38+
MachineFunction &MF;
4039
const GCNSubtarget &ST;
4140
const SIInstrInfo &TII;
4241
const SIRegisterInfo &TRI;
@@ -53,7 +52,7 @@ class AMDGPURewriteAGPRCopyMFMAImpl {
5352
AMDGPURewriteAGPRCopyMFMAImpl(MachineFunction &MF, VirtRegMap &VRM,
5453
LiveRegMatrix &LRM, LiveIntervals &LIS,
5554
const RegisterClassInfo &RegClassInfo)
56-
: ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()),
55+
: MF(MF), ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()),
5756
TRI(*ST.getRegisterInfo()), MRI(MF.getRegInfo()), VRM(VRM), LRM(LRM),
5857
LIS(LIS), RegClassInfo(RegClassInfo) {}
5958

@@ -71,26 +70,26 @@ class AMDGPURewriteAGPRCopyMFMAImpl {
7170
///
7271
/// \p RewriteRegs will accumulate the set of register used by those MFMAs
7372
/// that need to have the register classes adjusted.
74-
const TargetRegisterClass *recomputeRegClassExceptRewritable(
75-
Register Reg, const TargetRegisterClass *OldRC,
76-
const TargetRegisterClass *NewRC,
77-
SmallVectorImpl<MachineInstr *> &RewriteCandidates,
73+
bool recomputeRegClassExceptRewritable(
74+
Register Reg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
7875
SmallSetVector<Register, 4> &RewriteRegs) const;
7976

8077
bool run(MachineFunction &MF) const;
8178
};
8279

83-
const TargetRegisterClass *
84-
AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
85-
Register StartReg, const TargetRegisterClass *OldRC,
86-
const TargetRegisterClass *NewRC,
87-
SmallVectorImpl<MachineInstr *> &RewriteCandidates,
80+
bool AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
81+
Register StartReg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
8882
SmallSetVector<Register, 4> &RewriteRegs) const {
8983
SmallVector<Register, 8> Worklist = {StartReg};
9084

9185
// Recursively visit all transitive MFMA users
9286
while (!Worklist.empty()) {
9387
Register Reg = Worklist.pop_back_val();
88+
const TargetRegisterClass *OldRC = MRI.getRegClass(Reg);
89+
90+
// Inflate to the equivalent AV_* class.
91+
const TargetRegisterClass *NewRC = TRI.getLargestLegalSuperClass(OldRC, MF);
92+
9493
// Accumulate constraints from all uses.
9594
for (MachineOperand &MO : MRI.reg_nodbg_operands(Reg)) {
9695
// Apply the effect of the given operand to NewRC.
@@ -101,23 +100,40 @@ AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
101100
// either AGPR or VGPR in src0/src1, so don't bother checking the
102101
// constraint effects of the individual operands.
103102
if (isRewriteCandidate(*MI)) {
104-
for (AMDGPU::OpName OpName :
105-
{AMDGPU::OpName::vdst, AMDGPU::OpName::src2}) {
106-
const MachineOperand *Op = TII.getNamedOperand(*MI, OpName);
103+
const MachineOperand *VDst =
104+
TII.getNamedOperand(*MI, AMDGPU::OpName::vdst);
105+
const MachineOperand *Src2 =
106+
TII.getNamedOperand(*MI, AMDGPU::OpName::src2);
107+
for (const MachineOperand *Op : {VDst, Src2}) {
107108
if (!Op->isReg())
108109
continue;
109110

110111
Register OtherReg = Op->getReg();
111-
if (OtherReg != Reg) {
112-
if (RewriteRegs.insert(OtherReg))
113-
Worklist.push_back(OtherReg);
114-
}
112+
if (OtherReg.isPhysical())
113+
return false;
114+
115+
if (OtherReg != Reg && RewriteRegs.insert(OtherReg))
116+
Worklist.push_back(OtherReg);
115117
}
116118

117-
LLVM_DEBUG(dbgs() << "Ignoring effects of " << *MI);
119+
if (!is_contained(RewriteCandidates, MI)) {
120+
LLVM_DEBUG({
121+
Register VDstPhysReg = VRM.getPhys(VDst->getReg());
122+
dbgs() << "Attempting to replace VGPR MFMA with AGPR version:"
123+
<< " Dst=[" << printReg(VDst->getReg()) << " => "
124+
<< printReg(VDstPhysReg, &TRI);
125+
126+
if (Src2->isReg()) {
127+
Register Src2PhysReg = VRM.getPhys(Src2->getReg());
128+
dbgs() << ", Src2=[" << printReg(Src2->getReg(), &TRI) << " => "
129+
<< printReg(Src2PhysReg, &TRI);
130+
}
131+
132+
dbgs() << "]: " << MI;
133+
});
118134

119-
if (!is_contained(RewriteCandidates, MI))
120135
RewriteCandidates.push_back(MI);
136+
}
121137

122138
continue;
123139
}
@@ -126,13 +142,14 @@ AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
126142
NewRC = MI->getRegClassConstraintEffect(OpNo, NewRC, &TII, &TRI);
127143
if (!NewRC || NewRC == OldRC) {
128144
LLVM_DEBUG(dbgs() << "User of " << printReg(Reg, &TRI)
129-
<< " cannot be reassigned to AGPR: " << *MI);
130-
return nullptr;
145+
<< " cannot be reassigned to "
146+
<< TRI.getRegClassName(NewRC) << ": " << *MI);
147+
return false;
131148
}
132149
}
133150
}
134151

135-
return NewRC;
152+
return true;
136153
}
137154

138155
/// Attempt to reassign the registers in \p InterferingRegs to be AGPRs, with a
@@ -225,10 +242,7 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
225242

226243
for (VNInfo *VNI : LI.vnis()) {
227244
MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
228-
229-
// TODO: Handle SplitKit produced copy bundles for partially defined
230-
// registers.
231-
if (!DefMI || !DefMI->isFullCopy())
245+
if (!DefMI || !DefMI->isCopy())
232246
continue;
233247

234248
Register MFMADstReg = DefMI->getOperand(1).getReg();
@@ -241,34 +255,6 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
241255
if (!MFMA || !isRewriteCandidate(*MFMA))
242256
continue;
243257

244-
MachineOperand *Src2 = TII.getNamedOperand(*MFMA, AMDGPU::OpName::src2);
245-
Register Src2Reg;
246-
if (Src2->isReg()) {
247-
Src2Reg = Src2->getReg();
248-
if (!Src2Reg.isVirtual())
249-
continue;
250-
}
251-
252-
// FIXME: getMinimalPhysRegClass returns a nonsense AV_* subclass instead
253-
// of an AGPR or VGPR subclass, so we can't simply use the result on the
254-
// assignment.
255-
256-
LLVM_DEBUG({
257-
dbgs() << "Attempting to replace VGPR MFMA with AGPR version:"
258-
<< " Dst=[" << printReg(VReg) << " => "
259-
<< printReg(PhysReg, &TRI);
260-
261-
if (Src2Reg) {
262-
Register Src2PhysReg = VRM.getPhys(Src2Reg);
263-
dbgs() << ", Src2=[" << printReg(Src2Reg, &TRI) << " => "
264-
<< printReg(Src2PhysReg, &TRI);
265-
}
266-
267-
dbgs() << "]: " << *MFMA;
268-
});
269-
270-
const TargetRegisterClass *DstVirtRegRC = MRI.getRegClass(MFMADstReg);
271-
272258
// src2 and dst have the same physical class constraint; try to preserve
273259
// the original src2 subclass if one were to exist.
274260
SmallVector<MachineInstr *, 4> RewriteCandidates = {MFMA};
@@ -287,11 +273,9 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
287273
//
288274
// Note recomputeRegClassExceptRewritable will consider the constraints of
289275
// this MFMA's src2 as well as the src2/dst of any transitive MFMA users.
290-
const TargetRegisterClass *DstExceptRC =
291-
recomputeRegClassExceptRewritable(MFMADstReg, DstVirtRegRC, VirtRegRC,
292-
RewriteCandidates, RewriteRegs);
293-
if (!DstExceptRC) {
294-
LLVM_DEBUG(dbgs() << "Could not recompute the regclass of "
276+
if (!recomputeRegClassExceptRewritable(MFMADstReg, RewriteCandidates,
277+
RewriteRegs)) {
278+
LLVM_DEBUG(dbgs() << "Could not recompute the regclass of dst reg "
295279
<< printReg(MFMADstReg, &TRI) << '\n');
296280
continue;
297281
}

‎llvm/test/CodeGen/AMDGPU/inflate-reg-class-vgpr-mfma-to-agpr-negative-tests.mir‎

Lines changed: 0 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,6 @@
2020
ret void
2121
}
2222

23-
define amdgpu_kernel void @inflate_result_to_agpr__V_MFMA_F32_32X32X8F16_vgprcd_e64_src2_different_subreg() #0 {
24-
ret void
25-
}
26-
2723
define amdgpu_kernel void @inflate_result_to_agpr__V_MFMA_F32_32X32X8F16_vgprcd_e64_chain_no_agprs_first() #1 {
2824
ret void
2925
}
@@ -420,93 +416,6 @@ body: |
420416
421417
...
422418

423-
# Non-mac variant, src2 is the same VGPR, but a different subregister.
424-
---
425-
name: inflate_result_to_agpr__V_MFMA_F32_32X32X8F16_vgprcd_e64_src2_different_subreg
426-
tracksRegLiveness: true
427-
machineFunctionInfo:
428-
isEntryFunction: true
429-
stackPtrOffsetReg: '$sgpr32'
430-
occupancy: 10
431-
sgprForEXECCopy: '$sgpr100_sgpr101'
432-
body: |
433-
; CHECK-LABEL: name: inflate_result_to_agpr__V_MFMA_F32_32X32X8F16_vgprcd_e64_src2_different_subreg
434-
; CHECK: bb.0:
435-
; CHECK-NEXT: successors: %bb.1(0x80000000)
436-
; CHECK-NEXT: {{ $}}
437-
; CHECK-NEXT: S_NOP 0, implicit-def $agpr0
438-
; CHECK-NEXT: renamable $sgpr0 = S_MOV_B32 0
439-
; CHECK-NEXT: renamable $vgpr8 = V_MOV_B32_e32 0, implicit $exec
440-
; CHECK-NEXT: renamable $sgpr1 = COPY renamable $sgpr0
441-
; CHECK-NEXT: renamable $vgpr18_vgpr19 = COPY killed renamable $sgpr0_sgpr1
442-
; CHECK-NEXT: renamable $vcc = S_AND_B64 $exec, -1, implicit-def dead $scc
443-
; CHECK-NEXT: dead renamable $vgpr9 = COPY renamable $vgpr8
444-
; CHECK-NEXT: {{ $}}
445-
; CHECK-NEXT: bb.1:
446-
; CHECK-NEXT: successors: %bb.1(0x40000000), %bb.2(0x40000000)
447-
; CHECK-NEXT: liveins: $vcc, $vgpr18_vgpr19
448-
; CHECK-NEXT: {{ $}}
449-
; CHECK-NEXT: renamable $vgpr0_vgpr1 = GLOBAL_LOAD_DWORDX2 undef renamable $vgpr0_vgpr1, 0, 0, implicit $exec :: (load (s64), addrspace 1)
450-
; CHECK-NEXT: renamable $vgpr16_vgpr17 = GLOBAL_LOAD_DWORDX2 undef renamable $vgpr0_vgpr1, 0, 0, implicit $exec :: (load (s64), addrspace 1)
451-
; CHECK-NEXT: early-clobber renamable $vgpr0_vgpr1_vgpr2_vgpr3_vgpr4_vgpr5_vgpr6_vgpr7_vgpr8_vgpr9_vgpr10_vgpr11_vgpr12_vgpr13_vgpr14_vgpr15 = V_MFMA_F32_32X32X8F16_vgprcd_e64 $vgpr18_vgpr19, $vgpr18_vgpr19, $vgpr16_vgpr17_vgpr18_vgpr19_vgpr20_vgpr21_vgpr22_vgpr23_vgpr24_vgpr25_vgpr26_vgpr27_vgpr28_vgpr29_vgpr30_vgpr31, 0, 0, 0, implicit $mode, implicit $exec
452-
; CHECK-NEXT: renamable $agpr0_agpr1_agpr2_agpr3_agpr4_agpr5_agpr6_agpr7_agpr8_agpr9_agpr10_agpr11_agpr12_agpr13_agpr14_agpr15 = COPY killed renamable $vgpr0_vgpr1_vgpr2_vgpr3_vgpr4_vgpr5_vgpr6_vgpr7_vgpr8_vgpr9_vgpr10_vgpr11_vgpr12_vgpr13_vgpr14_vgpr15
453-
; CHECK-NEXT: S_CBRANCH_VCCNZ %bb.1, implicit $vcc
454-
; CHECK-NEXT: S_BRANCH %bb.2
455-
; CHECK-NEXT: {{ $}}
456-
; CHECK-NEXT: bb.2:
457-
; CHECK-NEXT: liveins: $agpr0_agpr1_agpr2_agpr3_agpr4_agpr5_agpr6_agpr7_agpr8_agpr9_agpr10_agpr11_agpr12_agpr13_agpr14_agpr15_agpr16_agpr17_agpr18_agpr19_agpr20_agpr21_agpr22_agpr23_agpr24_agpr25_agpr26_agpr27_agpr28_agpr29_agpr30_agpr31:0x00000000FFFFFFFF
458-
; CHECK-NEXT: {{ $}}
459-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr0_vgpr1_vgpr2_vgpr3_vgpr4_vgpr5_vgpr6_vgpr7
460-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr8_vgpr9_vgpr10_vgpr11_vgpr12_vgpr13_vgpr14_vgpr15
461-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr16_vgpr17_vgpr18_vgpr19_vgpr20_vgpr21_vgpr22_vgpr23
462-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr24_vgpr25_vgpr26_vgpr27_vgpr28_vgpr29_vgpr30_vgpr31
463-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr32_vgpr33_vgpr34_vgpr35_vgpr36_vgpr37_vgpr38_vgpr39
464-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr40_vgpr41_vgpr42_vgpr43_vgpr44_vgpr45_vgpr46_vgpr47
465-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr48_vgpr49_vgpr50_vgpr51_vgpr52_vgpr53_vgpr54_vgpr55
466-
; CHECK-NEXT: S_NOP 0, implicit-def $vgpr56_vgpr57_vgpr58_vgpr59_vgpr60_vgpr61_vgpr62_vgpr63
467-
; CHECK-NEXT: renamable $vgpr0 = V_MOV_B32_e32 0, implicit $exec
468-
; CHECK-NEXT: GLOBAL_STORE_DWORDX4_SADDR renamable $vgpr0, renamable $agpr8_agpr9_agpr10_agpr11, undef $sgpr0_sgpr1, 32, 0, implicit $exec :: (store (s128), align 32, addrspace 1)
469-
; CHECK-NEXT: GLOBAL_STORE_DWORDX4_SADDR renamable $vgpr0, renamable $agpr12_agpr13_agpr14_agpr15, undef $sgpr0_sgpr1, 48, 0, implicit $exec :: (store (s128), addrspace 1)
470-
; CHECK-NEXT: GLOBAL_STORE_DWORDX4_SADDR renamable $vgpr0, renamable $agpr0_agpr1_agpr2_agpr3, undef $sgpr0_sgpr1, 0, 0, implicit $exec :: (store (s128), align 128, addrspace 1)
471-
; CHECK-NEXT: GLOBAL_STORE_DWORDX4_SADDR killed renamable $vgpr0, killed renamable $agpr4_agpr5_agpr6_agpr7, killed undef $sgpr0_sgpr1, 16, 0, implicit $exec :: (store (s128), addrspace 1)
472-
; CHECK-NEXT: S_ENDPGM 0
473-
bb.0:
474-
S_NOP 0, implicit-def $agpr0
475-
renamable $sgpr0 = S_MOV_B32 0
476-
undef %0.sub8:vreg_1024_align2 = V_MOV_B32_e32 0, implicit $exec
477-
renamable $sgpr1 = COPY renamable $sgpr0
478-
%1:vreg_64_align2 = COPY killed renamable $sgpr0_sgpr1
479-
renamable $vcc = S_AND_B64 $exec, -1, implicit-def dead $scc
480-
%0.sub9:vreg_1024_align2 = COPY %0.sub8
481-
482-
bb.1:
483-
liveins: $vcc
484-
485-
undef %0.sub0_sub1:vreg_1024_align2 = GLOBAL_LOAD_DWORDX2 undef %3:vreg_64_align2, 0, 0, implicit $exec :: (load (s64), addrspace 1)
486-
%0.sub16_sub17:vreg_1024_align2 = GLOBAL_LOAD_DWORDX2 undef %3:vreg_64_align2, 0, 0, implicit $exec :: (load (s64), addrspace 1)
487-
%0.sub0_sub1_sub2_sub3_sub4_sub5_sub6_sub7_sub8_sub9_sub10_sub11_sub12_sub13_sub14_sub15:vreg_1024_align2 = V_MFMA_F32_32X32X8F16_vgprcd_e64 %1, %1, %0.sub16_sub17_sub18_sub19_sub20_sub21_sub22_sub23_sub24_sub25_sub26_sub27_sub28_sub29_sub30_sub31, 0, 0, 0, implicit $mode, implicit $exec
488-
S_CBRANCH_VCCNZ %bb.1, implicit $vcc
489-
S_BRANCH %bb.2
490-
491-
bb.2:
492-
; No VGPRs available for %0
493-
S_NOP 0, implicit-def $vgpr0_vgpr1_vgpr2_vgpr3_vgpr4_vgpr5_vgpr6_vgpr7
494-
S_NOP 0, implicit-def $vgpr8_vgpr9_vgpr10_vgpr11_vgpr12_vgpr13_vgpr14_vgpr15
495-
S_NOP 0, implicit-def $vgpr16_vgpr17_vgpr18_vgpr19_vgpr20_vgpr21_vgpr22_vgpr23
496-
S_NOP 0, implicit-def $vgpr24_vgpr25_vgpr26_vgpr27_vgpr28_vgpr29_vgpr30_vgpr31
497-
S_NOP 0, implicit-def $vgpr32_vgpr33_vgpr34_vgpr35_vgpr36_vgpr37_vgpr38_vgpr39
498-
S_NOP 0, implicit-def $vgpr40_vgpr41_vgpr42_vgpr43_vgpr44_vgpr45_vgpr46_vgpr47
499-
S_NOP 0, implicit-def $vgpr48_vgpr49_vgpr50_vgpr51_vgpr52_vgpr53_vgpr54_vgpr55
500-
S_NOP 0, implicit-def $vgpr56_vgpr57_vgpr58_vgpr59_vgpr60_vgpr61_vgpr62_vgpr63
501-
%2:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
502-
GLOBAL_STORE_DWORDX4_SADDR %2, %0.sub8_sub9_sub10_sub11, undef $sgpr0_sgpr1, 32, 0, implicit $exec :: (store (s128), align 32, addrspace 1)
503-
GLOBAL_STORE_DWORDX4_SADDR %2, %0.sub12_sub13_sub14_sub15, undef $sgpr0_sgpr1, 48, 0, implicit $exec :: (store (s128), addrspace 1)
504-
GLOBAL_STORE_DWORDX4_SADDR %2, %0.sub0_sub1_sub2_sub3, undef $sgpr0_sgpr1, 0, 0, implicit $exec :: (store (s128), align 128, addrspace 1)
505-
GLOBAL_STORE_DWORDX4_SADDR %2, %0.sub4_sub5_sub6_sub7, killed undef $sgpr0_sgpr1, 16, 0, implicit $exec :: (store (s128), addrspace 1)
506-
S_ENDPGM 0
507-
508-
...
509-
510419
# There isn't an assignable AGPR around the first MFMA.
511420
---
512421
name: inflate_result_to_agpr__V_MFMA_F32_32X32X8F16_vgprcd_e64_chain_no_agprs_first

0 commit comments

Comments
 (0)