AMDGPU: Don't fold AGPR copy pairs that need a temp VGPR

[arsenm]

Summary

Do not fold SGPR->VGPR->AGPR copies, or AGPR->VGPR->AGPR copies. We
need to introduce a temporary VGPR. It's better if this this temp
register is explicitly visible to the register allocator, rather than
relying on the register scavenger later. A future change will
introduce the temporary VGPRs, and this combine would undo that.

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

175

It is not a part of this change.

[kerbowa]

I'm working on eliminating some AGPR->VGPR->AGPR copies which commonly cause problems in loops where the VGPR is considered live across iterations by the register allocator even if the VGPR copy should be sunk. Why do we need this temp VGPR?

[arsenm]

In D93715#2477411, @kerbowa wrote:

I'm working on eliminating some AGPR->VGPR->AGPR copies which commonly cause problems in loops where the VGPR is considered live across iterations by the register allocator even if the VGPR copy should be sunk. Why do we need this temp VGPR?

Because you can't directly copy between AGPRs (or directly from SGPR to AGPR)

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

while (ImpOpI != ImpOpE) {
       ImpOpI++;
       UseMI->RemoveOperand(UseMI->getOperandNo(Tmp));
     }
     CopiesToReplace.push_back(UseMI);
   } else {
     if (UseMI->isCopy() && OpToFold.isReg() &&
         UseMI->getOperand(0).getReg().isVirtual() &&
         !UseMI->getOperand(1).getSubReg()) {
-      LLVM_DEBUG(dbgs() << "Folding " << OpToFold
-                        << "\n into " << *UseMI << '\n');
+
       unsigned Size = TII->getOpSize(*UseMI, 1);
       Register UseReg = OpToFold.getReg();
+      Register DstReg = UseMI->getOperand(0).getReg();
+
+      const bool DstIsAGPR = TRI->isAGPR(*MRI, DstReg);
+
+      SmallVector<std::pair<MachineOperand *, unsigned>, 32> Defs;
+      bool IsAGPRRegSeqInit =
+          Size > 4 && DstIsAGPR &&
+          getRegSeqInit(Defs, UseReg, AMDGPU::OPERAND_REG_INLINE_C_INT32, TII,
+                        *MRI);
+      if (DstIsAGPR && !IsAGPRRegSeqInit &&
+          !TRI->isVGPR(*MRI, OpToFold.getReg())) {
+        LLVM_DEBUG(dbgs() << "Not folding copy to AGPR\n");
+        return;
+      }
+
       UseMI->getOperand(1).setReg(UseReg);
       UseMI->getOperand(1).setSubReg(OpToFold.getSubReg());
       UseMI->getOperand(1).setIsKill(false);
       CopiesToReplace.push_back(UseMI);
       OpToFold.setIsKill(false);
 
       // That is very tricky to store a value into an AGPR. v_accvgpr_write_b32
       // can only accept VGPR or inline immediate. Recreate a reg_sequence with
       // its initializers right here, so we will rematerialize immediates and
       // avoid copies via different reg classes.
-      SmallVector<std::pair<MachineOperand*, unsigned>, 32> Defs;
-      if (Size > 4 && TRI->isAGPR(*MRI, UseMI->getOperand(0).getReg()) &&
-          getRegSeqInit(Defs, UseReg, AMDGPU::OPERAND_REG_INLINE_C_INT32, TII,
-                        *MRI)) {
+      if (IsAGPRRegSeqInit) {
         const DebugLoc &DL = UseMI->getDebugLoc();
         MachineBasicBlock &MBB = *UseMI->getParent();
 
         UseMI->setDesc(TII->get(AMDGPU::REG_SEQUENCE));
         for (unsigned I = UseMI->getNumOperands() - 1; I > 0; --I)
           UseMI->RemoveOperand(I);
 
         MachineInstrBuilder B(*MBB.getParent(), UseMI);
if (UseMI->isCopy() && OpToFold.isReg() &&
           B.addImm(Defs[I].second);
         }
         LLVM_DEBUG(dbgs() << "Folded " << *UseMI << '\n');
         return;
       }
 
       if (Size != 4)
         return;
-      if (TRI->isAGPR(*MRI, UseMI->getOperand(0).getReg()) &&
-          TRI->isVGPR(*MRI, UseMI->getOperand(1).getReg()))
+      if (DstIsAGPR && TRI->isVGPR(*MRI, UseMI->getOperand(1).getReg()))
         UseMI->setDesc(TII->get(AMDGPU::V_ACCVGPR_WRITE_B32));
       else if (TRI->isVGPR(*MRI, UseMI->getOperand(0).getReg()) &&
                TRI->isAGPR(*MRI, UseMI->getOperand(1).getReg()))
         UseMI->setDesc(TII->get(AMDGPU::V_ACCVGPR_READ_B32));
       return;
     }
 
     unsigned UseOpc = UseMI->getOpcode();

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

bool isSGPRReg(const MachineRegisterInfo &MRI, Register Reg) const {
     return isSGPRClass(RC);
   }
 
   /// \returns true if this class contains only AGPR registers
   bool isAGPRClass(const TargetRegisterClass *RC) const {
     return hasAGPRs(RC) && !hasVGPRs(RC);
   }
 
+  /// \returns true if this class contains only VGPR registers FIXME: This works
+  /// correctly for classes obtained VGPR vregs, but will be incorrect for the
+  /// combined VS_* classes.
+  bool isVGPRClass(const TargetRegisterClass *RC) const {
+    return hasVGPRs(RC);
+  }
+
   /// \returns true if this class contains VGPR registers.
   bool hasVGPRs(const TargetRegisterClass *RC) const;
 
   /// \returns true if this class contains AGPR registers.
   bool hasAGPRs(const TargetRegisterClass *RC) const;
 
   /// \returns true if this class contains any vector registers.
   bool hasVectorRegisters(const TargetRegisterClass *RC) const {

llvm/test/CodeGen/AMDGPU/fold-operands-agpr-copies.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 -verify-machineinstrs -run-pass=si-fold-operands,dead-mi-elimination -o - %s | FileCheck -check-prefix=GFX908 %s
+
+---
+name: no_fold_s32_to_v32_to_a32
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: no_fold_s32_to_v32_to_a32
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[S_LOAD_DWORD_IMM:%[0-9]+]]:sreg_32_xm0_xexec = S_LOAD_DWORD_IMM undef %1:sreg_64, 0, 0, 0 :: (load 4, addrspace 4)
+    ; GFX908: [[COPY:%[0-9]+]]:vgpr_32 = COPY [[S_LOAD_DWORD_IMM]]
+    ; GFX908: [[COPY1:%[0-9]+]]:agpr_32 = COPY [[COPY]]
+    ; GFX908: S_ENDPGM 0, implicit [[COPY1]]
+    %0:sreg_32_xm0_xexec = S_LOAD_DWORD_IMM undef %3:sreg_64, 0, 0, 0 :: (load 4, addrspace 4)
+    %1:vgpr_32 = COPY %0
+    %2:agpr_32 = COPY %1
+    S_ENDPGM 0, implicit %2
+
+...
+
+---
+name: no_fold_a32_to_v32_to_a32
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $agpr0
+
+    ; GFX908-LABEL: name: no_fold_a32_to_v32_to_a32
+    ; GFX908: liveins: $agpr0
+    ; GFX908: [[COPY:%[0-9]+]]:agpr_32 = COPY $agpr0
+    ; GFX908: [[COPY1:%[0-9]+]]:vgpr_32 = COPY [[COPY]]
+    ; GFX908: [[COPY2:%[0-9]+]]:agpr_32 = COPY [[COPY1]]
+    ; GFX908: S_ENDPGM 0, implicit [[COPY2]]
+    %0:agpr_32 = COPY $agpr0
+    %1:vgpr_32 = COPY %0
+    %2:agpr_32 = COPY %1
+    S_ENDPGM 0, implicit %2
+
+...
+
+---
+name: fold_imm_a32_to_v32_to_a32
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: fold_imm_a32_to_v32_to_a32
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: S_ENDPGM 0, implicit [[V_ACCVGPR_WRITE_B32_]]
+    %0:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    %1:vgpr_32 = COPY %0
+    %2:agpr_32 = COPY %1
+    S_ENDPGM 0, implicit %2
+
+...
+
+---
+name: fold_imm_s32_to_v32_to_a32
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: fold_imm_s32_to_v32_to_a32
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: S_ENDPGM 0, implicit [[V_ACCVGPR_WRITE_B32_]]
+    %0:sreg_32 = S_MOV_B32 0
+    %1:vgpr_32 = COPY %0
+    %2:agpr_32 = COPY %1
+    S_ENDPGM 0, implicit %2
+
+...
+
+
+# Make sure we still fold immediate materializes into the AGPR moves.
+---
+name: reg_sequence_imm_s64_to_a64_0
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: reg_sequence_imm_s64_to_a64_0
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[S_MOV_B32_:%[0-9]+]]:sreg_32_xm0_xexec = S_MOV_B32 0
+    ; GFX908: [[REG_SEQUENCE:%[0-9]+]]:sreg_64 = REG_SEQUENCE [[S_MOV_B32_]], %subreg.sub0, [[S_MOV_B32_]], %subreg.sub1
+    ; GFX908: [[COPY:%[0-9]+]]:areg_64 = COPY [[REG_SEQUENCE]]
+    ; GFX908: S_ENDPGM 0, implicit [[COPY]]
+    %0:sreg_32_xm0_xexec = S_MOV_B32 0
+    %1:sreg_64 = REG_SEQUENCE %0, %subreg.sub0,  %0, %subreg.sub1
+    %2:areg_64 = COPY %1
+    S_ENDPGM 0, implicit %2
+
+...
+
+---
+name: reg_sequence_imm_s64_to_a64_1
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: reg_sequence_imm_s64_to_a64_1
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_1:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[REG_SEQUENCE:%[0-9]+]]:areg_64 = REG_SEQUENCE [[V_ACCVGPR_WRITE_B32_]], %subreg.sub0, [[V_ACCVGPR_WRITE_B32_1]], %subreg.sub1
+    ; GFX908: S_ENDPGM 0, implicit [[REG_SEQUENCE]]
+    %0:sreg_32_xm0_xexec = S_MOV_B32 0
+    %1:sreg_64 = REG_SEQUENCE %0, %subreg.sub0,  %0, %subreg.sub1
+    %2:sreg_64 = COPY %1
+    %3:areg_64 = COPY %2
+    S_ENDPGM 0, implicit %3
+
+...
+
+---
+name: reg_sequence_imm_s64_to_v64_to_a64
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: reg_sequence_imm_s64_to_v64_to_a64
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_1:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[REG_SEQUENCE:%[0-9]+]]:areg_64 = REG_SEQUENCE [[V_ACCVGPR_WRITE_B32_]], %subreg.sub0, [[V_ACCVGPR_WRITE_B32_1]], %subreg.sub1
+    ; GFX908: S_ENDPGM 0, implicit [[REG_SEQUENCE]]
+    %0:sreg_32_xm0_xexec = S_MOV_B32 0
+    %1:sreg_64 = REG_SEQUENCE %0, %subreg.sub0,  %0, %subreg.sub1
+    %2:vreg_64 = COPY %1
+    %3:areg_64 = COPY %2
+    S_ENDPGM 0, implicit %3
+
+...
+
+---
+name: reg_sequence_imm_v64_to_a64
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: reg_sequence_imm_v64_to_a64
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_1:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[REG_SEQUENCE:%[0-9]+]]:areg_64 = REG_SEQUENCE [[V_ACCVGPR_WRITE_B32_]], %subreg.sub0, [[V_ACCVGPR_WRITE_B32_1]], %subreg.sub1
+    ; GFX908: S_ENDPGM 0, implicit [[REG_SEQUENCE]]
+    %0:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    %1:vreg_64 = REG_SEQUENCE %0, %subreg.sub0,  %0, %subreg.sub1
+    %2:vreg_64 = COPY %1
+    %3:areg_64 = COPY %2
+    S_ENDPGM 0, implicit %3
+
+...
+
+---
+name: reg_squence_imm_a64_to_v64_to_a64
+tracksRegLiveness: true
+body:               |
+  bb.0:
+    liveins: $sgpr0
+
+    ; GFX908-LABEL: name: reg_squence_imm_a64_to_v64_to_a64
+    ; GFX908: liveins: $sgpr0
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[V_ACCVGPR_WRITE_B32_1:%[0-9]+]]:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    ; GFX908: [[REG_SEQUENCE:%[0-9]+]]:agpr_32 = REG_SEQUENCE [[V_ACCVGPR_WRITE_B32_]], %subreg.sub0, [[V_ACCVGPR_WRITE_B32_1]], %subreg.sub1
+    ; GFX908: S_ENDPGM 0, implicit [[REG_SEQUENCE]]
+    %0:agpr_32 = V_ACCVGPR_WRITE_B32 0, implicit $exec
+    %1:areg_64 = REG_SEQUENCE %0, %subreg.sub0,  %0, %subreg.sub1
+    %2:vreg_64 = COPY %1
+    %3:agpr_32 = COPY %2
+    S_ENDPGM 0, implicit %3
+
+...