[AMDGPU] Propagate defining src reg for AGPR to AGPR Copys

[vangthao]

Summary

On targets that do not support AGPR to AGPR copying directly, try to find the
defining accvgpr_write and propagate its source vgpr register to the copies
before register allocation so the source vgpr register does not get clobbered.

The postrapseudos pass also attempt to propagate the defining accvgpr_write but
if the register to propagate is clobbered, it will give up and create new
temporary vgpr registers instead.

[vangthao]

This attempts to address the issue seen here https://reviews.llvm.org/D105062

[arsenm]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

160

Can skip this whole pass if this isn't gfx908

161–162

Since you're visiting all the defs anyway, could you unify this with the previous loop when an agpr register is found?

201

I'm worried this might not work correctly with overlapping subregisters

211–214

Should not need to totally trash the liveness and start over

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

161–162

Second to that. We already scanning all registers in the processReg() and all their defs, so it shall not be hard to find a def of an AGPR which is a COPY.

186

It is better to use isAGPRClass() instead. AV classes may become allocatable one day [soon].

[vangthao]

Merged loop. Changed uses of TRI->hasAGPRs() to TRI->isAGPRClass(). Added some subregs test.

[vangthao]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

201

I'm worried this might not work correctly with overlapping subregisters

I added some new tests. Do they address this? If not can you give an example?

211–214

Should not need to totally trash the liveness and start over

Is there a way to update the liveness without trashing it?

[vangthao]

Fixed clang-tidy variable case style warning and rebased.

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

102

I do not think you need this case. This is a single 32 bit register, it will have a single def.

102–103

... and here.

107

You know you are dealing with an AGPR or SGPR right at the beginning of the function. You can quickly bail from here based on that.

147

You can just return false here.

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

107

DstReg always the same as Reg or at least should be unless it is not an implicit-def. At maximum I would expect assert(Reg == I.getOperand(0).getReg()) here.

110

Reg is always virtual, it is a result of index2VirtReg() call.

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

125

I am no sure why do you need to match SRC and DST subregs. Why not just copy a source from Def here whatever it is?

[vangthao]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

102

This may not necessarily be a single def. This may contain multiple defs if subregs are used. For example in the case that I am looking at:

undef %1.sub0:areg_128 = V_ACCVGPR_WRITE_B32_e64 %0.sub0:vreg_128, implicit $exec
%1.sub1:areg_128 = COPY %1.sub0:areg_128
%1.sub2:areg_128 = COPY %1.sub0:areg_128
%1.sub3:areg_128 = COPY %1.sub0:areg_128

These subregs are considered to be under %1 since def_instruction() does not consider different subregs as far as I can tell. I think index2VirtReg() does not consider subregs either.

125

Same as other comment. There may be subregs that def_instructions does not account for. We need to check if we are copying the correct subreg.

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

102–103

Same thing, it should continue outer loop.

125

I see.

141

Actually you can defer LIS update until after the loop. If you are dealing with many subregs of the same register you would only need to recreate interval once.

147

In fact not, you still want to continue outer loop.

[vangthao]

Check if Reg is AGPR before doing anything in Copy switch case.
Remove isVirtual() check for Reg since index2VirtReg() guarantees virtual.
Defer LIS update until end.

[rampitec]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

89

A largest tuple is 32, so maybe SmallSet<Register, 32>?

108

Can you hoist it out the loop? Reg is loop invariant.

113

Also replace it with hoisted IsAGPRDst.

[vangthao]

Hoisted IsAGPRDst out of loop. Change SmallSet size to 32.

[rampitec]

LGTM

[piotr]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

109

I'm seeing issues with this commit on targets without agpr (e.g. gfx900).
I think there is a change in behaviour here - previously on COPY there was an early return. Now, there is only a switch break, so the loop continues. Is this intended?

[foad]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

109

@vangthao ping! This is causing codegen differences on gfx900 which has no AGPRs. Can you please comment on whether that is expected, and if so why? Thanks!

[vangthao]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

109

Sorry I missed this comment.

This behavior is not intended. The issue was found and should be fixed in https://reviews.llvm.org/D113005

[foad]

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

109

Thanks. We are testing that fix now.

[piotr]

For the record, D113005 indeed fixes the CTS issue I was seeing with D108830. Thanks.

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

FunctionPass *llvm::createGCNPreRAOptimizationsPass() {
   return new GCNPreRAOptimizations();
 }
 
 bool GCNPreRAOptimizations::processReg(Register Reg) {
   MachineInstr *Def0 = nullptr;
   MachineInstr *Def1 = nullptr;
   uint64_t Init = 0;
 
   for (MachineInstr &I : MRI->def_instructions(Reg)) {
     if (I.getOpcode() != AMDGPU::S_MOV_B32 || I.getOperand(0).getReg() != Reg ||
         !I.getOperand(1).isImm() || I.getNumOperands() != 2)
       return false;
 
     switch (I.getOperand(0).getSubReg()) {
     default:
       return false;
     case AMDGPU::sub0:
       if (Def0)
         return false;
       Def0 = &I;
       Init |= I.getOperand(1).getImm() & 0xffffffff;
       break;
     case AMDGPU::sub1:
       if (Def1)
         return false;
       Def1 = &I;
       Init |= static_cast<uint64_t>(I.getOperand(1).getImm()) << 32;
       break;
     }
   }
 
   if (!Def0 || !Def1 || Def0->getParent() != Def1->getParent())
     return false;
 
   LLVM_DEBUG(dbgs() << "Combining:\n  " << *Def0 << "  " << *Def1
                     << "    =>\n");
 
   if (SlotIndex::isEarlierInstr(LIS->getInstructionIndex(*Def1),
                                 LIS->getInstructionIndex(*Def0)))
     std::swap(Def0, Def1);
 
   LIS->RemoveMachineInstrFromMaps(*Def0);
   LIS->RemoveMachineInstrFromMaps(*Def1);
   auto NewI = BuildMI(*Def0->getParent(), *Def0, Def0->getDebugLoc(),
                       TII->get(AMDGPU::S_MOV_B64_IMM_PSEUDO), Reg)
                   .addImm(Init);
 
   Def0->eraseFromParent();
   Def1->eraseFromParent();
   LIS->InsertMachineInstrInMaps(*NewI);
   LIS->removeInterval(Reg);
   LIS->createAndComputeVirtRegInterval(Reg);
 
   LLVM_DEBUG(dbgs() << "  " << *NewI);
 
   return true;
 }
 
 bool GCNPreRAOptimizations::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   TII = ST.getInstrInfo();
   MRI = &MF.getRegInfo();
   LIS = &getAnalysis<LiveIntervals>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
 
   bool Changed = false;
 
   for (unsigned I = 0, E = MRI->getNumVirtRegs(); I != E; ++I) {
     Register Reg = Register::index2VirtReg(I);
     if (!LIS->hasInterval(Reg))
       continue;
     const TargetRegisterClass *RC = MRI->getRegClass(Reg);
     if (RC->MC->getSizeInBits() != 64 || !TRI->isSGPRClass(RC))
       continue;
     Changed |= processReg(Reg);
   }
 
+  for (MachineFunction::iterator BI = MF.begin(), BE = MF.end(); BI != BE;
+       ++BI) {
+    MachineBasicBlock *MBB = &*BI;
+    for (MachineBasicBlock::iterator I = MBB->begin(), E = MBB->end(); I != E;
+         ++I) {
+      MachineInstr &MI = *I;
+
+      switch (MI.getOpcode()) {
+      default:
+        continue;
+      case AMDGPU::COPY: {
+        // Some subtargets cannot do an AGPR to AGPR copy directly, and need an
+        // intermdiate temporary VGPR register. Try to find the defining
+        // accvgpr_write to avoid temporary registers.
+
+        // GFX90A can directly copy AGPR to AGPR so nothing to do.
+        if (ST.hasGFX90AInsts())
+          continue;
+
+        Register DstReg = MI.getOperand(0).getReg();
+        Register SrcReg = MI.getOperand(1).getReg();
+        if (!DstReg.isVirtual() || !SrcReg.isVirtual())
+          continue;
+
+        // Check if copy is AGPR to AGPR
+        bool IsAGPRDst = TRI->hasAGPRs(MRI->getRegClass(DstReg));
+        bool IsAGPRSrc = TRI->hasAGPRs(MRI->getRegClass(SrcReg));
+        if (!IsAGPRDst || !IsAGPRSrc)
+          continue;
+
+        // def_instructions() does not look at subregs so it may give us a
+        // different instruction that defines the same vreg but diferent subreg
+        // so we have to manually check subreg.
+        Register SrcSubReg = MI.getOperand(1).getSubReg();
+
+        for (auto &Def : MRI->def_instructions(SrcReg)) {
+          if (SrcSubReg != Def.getOperand(0).getSubReg())
+            continue;
+
+          if (Def.getOpcode() == AMDGPU::V_ACCVGPR_WRITE_B32_e64) {
+            MachineOperand DefSrcMO = Def.getOperand(1);
+
+            // Immediates are not an issue and can be propagated in
+            // postrapseudos pass. Only handle cases where defining
+            // accvgpr_write source is a vreg.
+            if (DefSrcMO.isReg() && DefSrcMO.getReg().isVirtual()) {
+              // Propagate source reg of accvgpr write to this copy instruction
+              MI.getOperand(1).setReg(DefSrcMO.getReg());
+              MI.getOperand(1).setSubReg(DefSrcMO.getSubReg());
+
+              LIS->removeInterval(DefSrcMO.getReg());
+              LIS->removeInterval(SrcReg);
+              LIS->createAndComputeVirtRegInterval(DefSrcMO.getReg());
+              LIS->createAndComputeVirtRegInterval(SrcReg);
+              Changed = true;
+            }
+
+            break;
+          }
+        }
+
+        break;
+      }
+      }
+    }
+  }
+
   return Changed;
 }

llvm/test/CodeGen/AMDGPU/agpr-to-agpr-copy.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx908 -run-pass=liveintervals,amdgpu-pre-ra-optimizations -verify-machineinstrs %s -o - | FileCheck -check-prefix=GFX908 %s
+
+---
+name: test_mfma_f32_4x4x1f32_propagate_vgpr
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1
+    ; GFX908-LABEL: name: test_mfma_f32_4x4x1f32_propagate_vgpr
+    ; GFX908: liveins: $sgpr0_sgpr1
+    ; GFX908: [[COPY:%[0-9]+]]:sgpr_64(p4) = COPY $sgpr0_sgpr1
+    ; GFX908: [[S_LOAD_DWORDX2_IMM:%[0-9]+]]:sreg_64_xexec = S_LOAD_DWORDX2_IMM [[COPY]](p4), 36, 0 :: (dereferenceable invariant load (s64), align 4, addrspace 4)
+    ; GFX908: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    ; GFX908: [[V_MOV_B32_e32_1:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 1123418112, implicit $exec
+    ; GFX908: undef %4.sub0:areg_128 = V_ACCVGPR_WRITE_B32_e64 [[V_MOV_B32_e32_1]], implicit $exec
+    ; GFX908: %4.sub1:areg_128 = COPY [[V_MOV_B32_e32_1]]
+    ; GFX908: %4.sub2:areg_128 = COPY [[V_MOV_B32_e32_1]]
+    ; GFX908: %4.sub3:areg_128 = COPY [[V_MOV_B32_e32_1]]
+    ; GFX908: [[V_MOV_B32_e32_2:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 1073741824, implicit $exec
+    ; GFX908: [[V_MOV_B32_e32_3:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 1065353216, implicit $exec
+    ; GFX908: [[V_MFMA_F32_4X4X1F32_e64_:%[0-9]+]]:areg_128 = V_MFMA_F32_4X4X1F32_e64 [[V_MOV_B32_e32_3]], [[V_MOV_B32_e32_2]], %4, 0, 0, 0, implicit $mode, implicit $exec
+    ; GFX908: [[COPY1:%[0-9]+]]:vreg_128 = COPY [[V_MFMA_F32_4X4X1F32_e64_]]
+    ; GFX908: GLOBAL_STORE_DWORDX4_SADDR [[V_MOV_B32_e32_]], [[COPY1]], [[S_LOAD_DWORDX2_IMM]], 0, 0, implicit $exec :: (store (s128), addrspace 1)
+    ; GFX908: S_ENDPGM 0
+    %1:sgpr_64(p4) = COPY $sgpr0_sgpr1
+    %4:sreg_64_xexec = S_LOAD_DWORDX2_IMM %1:sgpr_64(p4), 36, 0 :: (dereferenceable invariant load (s64), align 4, addrspace 4)
+    %5:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    %13:vgpr_32 = V_MOV_B32_e32 1123418112, implicit $exec
+    undef %11.sub0:areg_128 = V_ACCVGPR_WRITE_B32_e64 %13:vgpr_32, implicit $exec
+    %11.sub1:areg_128 = COPY %11.sub0:areg_128
+    %11.sub2:areg_128 = COPY %11.sub0:areg_128
+    %11.sub3:areg_128 = COPY %11.sub0:areg_128
+    %8:vgpr_32 = V_MOV_B32_e32 1073741824, implicit $exec
+    %9:vgpr_32 = V_MOV_B32_e32 1065353216, implicit $exec
+    %10:areg_128 = V_MFMA_F32_4X4X1F32_e64 %9:vgpr_32, %8:vgpr_32, %11:areg_128, 0, 0, 0, implicit $mode, implicit $exec
+    %12:vreg_128 = COPY %10:areg_128
+    GLOBAL_STORE_DWORDX4_SADDR %5:vgpr_32, %12:vreg_128, %4:sreg_64_xexec, 0, 0, implicit $exec :: (store (s128), addrspace 1)
+    S_ENDPGM 0
+...
+---
+name: test_mfma_f32_4x4x1f32_no_propagate_imm
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1
+    ; GFX908-LABEL: name: test_mfma_f32_4x4x1f32_no_propagate_imm
+    ; GFX908: liveins: $sgpr0_sgpr1
+    ; GFX908: [[COPY:%[0-9]+]]:sgpr_64(p4) = COPY $sgpr0_sgpr1
+    ; GFX908: [[S_LOAD_DWORDX2_IMM:%[0-9]+]]:sreg_64_xexec = S_LOAD_DWORDX2_IMM [[COPY]](p4), 36, 0 :: (dereferenceable invariant load (s64), align 4, addrspace 4)
+    ; GFX908: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    ; GFX908: undef %3.sub0:areg_128 = V_ACCVGPR_WRITE_B32_e64 1073741824, implicit $exec
+    ; GFX908: %3.sub1:areg_128 = COPY %3.sub0
+    ; GFX908: %3.sub2:areg_128 = COPY %3.sub0
+    ; GFX908: %3.sub3:areg_128 = COPY %3.sub0
+    ; GFX908: [[V_MOV_B32_e32_1:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 1073741824, implicit $exec
+    ; GFX908: [[V_MOV_B32_e32_2:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 1065353216, implicit $exec
+    ; GFX908: [[V_MFMA_F32_4X4X1F32_e64_:%[0-9]+]]:areg_128 = V_MFMA_F32_4X4X1F32_e64 [[V_MOV_B32_e32_2]], [[V_MOV_B32_e32_1]], %3, 0, 0, 0, implicit $mode, implicit $exec
+    ; GFX908: [[COPY1:%[0-9]+]]:vreg_128 = COPY [[V_MFMA_F32_4X4X1F32_e64_]]
+    ; GFX908: GLOBAL_STORE_DWORDX4_SADDR [[V_MOV_B32_e32_]], [[COPY1]], [[S_LOAD_DWORDX2_IMM]], 0, 0, implicit $exec :: (store (s128), addrspace 1)
+    ; GFX908: S_ENDPGM 0
+    %1:sgpr_64(p4) = COPY $sgpr0_sgpr1
+    %4:sreg_64_xexec = S_LOAD_DWORDX2_IMM %1:sgpr_64(p4), 36, 0 :: (dereferenceable invariant load (s64), align 4, addrspace 4)
+    %5:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    undef %11.sub0:areg_128 = V_ACCVGPR_WRITE_B32_e64 1073741824, implicit $exec
+    %11.sub1:areg_128 = COPY %11.sub0:areg_128
+    %11.sub2:areg_128 = COPY %11.sub0:areg_128
+    %11.sub3:areg_128 = COPY %11.sub0:areg_128
+    %8:vgpr_32 = V_MOV_B32_e32 1073741824, implicit $exec
+    %9:vgpr_32 = V_MOV_B32_e32 1065353216, implicit $exec
+    %10:areg_128 = V_MFMA_F32_4X4X1F32_e64 %9:vgpr_32, %8:vgpr_32, %11:areg_128, 0, 0, 0, implicit $mode, implicit $exec
+    %12:vreg_128 = COPY %10:areg_128
+    GLOBAL_STORE_DWORDX4_SADDR %5:vgpr_32, %12:vreg_128, %4:sreg_64_xexec, 0, 0, implicit $exec :: (store (s128), addrspace 1)
+    S_ENDPGM 0
+...

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.mfma.ll

bb:
   %tid = call i32 @llvm.amdgcn.workitem.id.x()
   %gep = getelementptr inbounds <4 x float>, <4 x float> addrspace(1)* %arg, i32 %tid
   %mai.1 = tail call <4 x float> @llvm.amdgcn.mfma.f32.4x4x1f32(float 1.0, float 2.0, <4 x float> <float 123.0, float 123.0, float 123.0, float 123.0>, i32 0, i32 0, i32 0)
   ;store <4 x float> %mai.1, <4 x float> addrspace(1)* %arg
   store <4 x float> %mai.1, <4 x float> addrspace(1)* %gep
   ret void
 }
 
-; FIXME: Resulting code for splat is pretty bad. A v_mov_b32 is moved
-; in the middle of the expanded agpr reg_sequence. The broadcast of
-; the individual AGPR->AGPR components should avoid the intermediate AGPR case.
 ; GCN-LABEL: {{^}}test_mfma_f32_4x4x1f32_lit_splat_bad_code:
 ; GFX908_A: v_mov_b32_e32 [[TMP0:v[0-9]+]], 0x42f60000
 ; GCN:      v_accvgpr_write_b32 [[AGPR:a[0-9]+]], [[TMP0]]
-; GFX908:   s_nop 0
-; GFX908:   v_accvgpr_read_b32 [[TMP1:v[0-9]+]], [[AGPR]]
-; GFX908:   v_accvgpr_read_b32 [[TMP2:v[0-9]+]], [[AGPR]]
-; GFX908:   v_accvgpr_read_b32 [[TMP3:v[0-9]+]], [[AGPR]]
-; GFX908:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP1]]
-; GFX908:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP2]]
-; GFX908:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP3]]
+; GFX908-NEXT:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP0]]
+; GFX908-NEXT:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP0]]
+; GFX908-NEXT:   v_accvgpr_write_b32 a{{[0-9]+}}, [[TMP0]]
 ; GFX90A-COUNT-3: v_accvgpr_mov_b32 a{{[0-9]+}}, [[AGPR]]
 ; GCN: s_nop 0
 ; GFX908_A:  v_mfma_f32_4x4x1f32 a[{{[0-9]+:[0-9]+}}], {{v[0-9]+}}, {{v[0-9]+}}, a[{{[0-9]+:[0-9]+}}]
 ; GFX908-COUNT-4: v_accvgpr_read_b32
 ; GFX908:    global_store_dwordx4 v{{[0-9]+}}, v[{{[0-9:]+}}], s[{{[0-9:]+}}]
 ; GFX90A:    global_store_dwordx4 v{{[0-9]+}}, a[{{[0-9:]+}}], s[{{[0-9:]+}}]
 define amdgpu_kernel void @test_mfma_f32_4x4x1f32_lit_splat_bad_code(<4 x float> addrspace(1)* %arg) {
 bb: