[AMDGPU] Support wwm-reg AV spill pseudos

[cdevadas]

Summary

The wwm register spill pseudos are currently defined for VGPR_32
regclass. It causes a verifier error for gfx908 or above as the
regalloc sometimes restores the values to the vector superclass AV_32.
Fixing it by supporting AV wwm-spill pseudos as well.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

1890

You shouldn't need to consider the target, this should be solely determined from the class you already have

[cdevadas]

Addressed the review comments.
Should add a new test. Working on it.

[arsenm]

Awaiting mentioned test

[cdevadas]

Added lit test.

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
   }
 
   bool isSGPRSpill(uint16_t Opcode) const {
     return get(Opcode).TSFlags & SIInstrFlags::SGPRSpill;
   }
 
   static bool isWWMRegSpillOpcode(uint16_t Opcode) {
     return Opcode == AMDGPU::SI_SPILL_WWM_V32_SAVE ||
-           Opcode == AMDGPU::SI_SPILL_WWM_V32_RESTORE;
+           Opcode == AMDGPU::SI_SPILL_WWM_AV32_SAVE ||
+           Opcode == AMDGPU::SI_SPILL_WWM_V32_RESTORE ||
+           Opcode == AMDGPU::SI_SPILL_WWM_AV32_RESTORE;
   }
 
   static bool isDPP(const MachineInstr &MI) {
     return MI.getDesc().TSFlags & SIInstrFlags::DPP;
   }
 
   bool isDPP(uint16_t Opcode) const {
     return get(Opcode).TSFlags & SIInstrFlags::DPP;

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

case 64:
     return AMDGPU::SI_SPILL_AV512_SAVE;
   case 128:
     return AMDGPU::SI_SPILL_AV1024_SAVE;
   default:
     llvm_unreachable("unknown register size");
   }
 }
 
-static unsigned getWWMRegSpillSaveOpcode(unsigned Size) {
+static unsigned getWWMRegSpillSaveOpcode(unsigned Size,
+                                         bool HasVectorSuperClass) {
   // Currently, there is only 32-bit WWM register spills needed.
   if (Size != 4)
     llvm_unreachable("unknown wwm register spill size");
 
+  // Return the AV pseudo opcode for gfx908+.
+  if (HasVectorSuperClass)
+    return AMDGPU::SI_SPILL_WWM_AV32_SAVE;
+
   return AMDGPU::SI_SPILL_WWM_V32_SAVE;
 }
 
 static unsigned getVectorRegSpillSaveOpcode(Register Reg,
                                             const TargetRegisterClass *RC,
                                             unsigned Size,
                                             const SIRegisterInfo &TRI,
-                                            const SIMachineFunctionInfo &MFI) {
+                                            const SIMachineFunctionInfo &MFI,
+                                            bool HasVectorSuperClass) {
   // Choose the right opcode if spilling a WWM register.
   if (MFI.checkFlag(Reg, AMDGPU::VirtRegFlag::WWM_REG))
-    return getWWMRegSpillSaveOpcode(Size);
+    return getWWMRegSpillSaveOpcode(Size, HasVectorSuperClass);
 
   if (TRI.isVectorSuperClass(RC))
     return getAVSpillSaveOpcode(Size);
 
   return TRI.isAGPRClass(RC) ? getAGPRSpillSaveOpcode(Size)
                              : getVGPRSpillSaveOpcode(Size);
 }
 
BuildMI(MBB, MI, DL, OpDesc)
       .addMemOperand(MMO)
       .addReg(MFI->getStackPtrOffsetReg(), RegState::Implicit);
 
     if (RI.spillSGPRToVGPR())
       FrameInfo.setStackID(FrameIndex, TargetStackID::SGPRSpill);
     return;
   }
 
-  unsigned Opcode = getVectorRegSpillSaveOpcode(VReg ? VReg : SrcReg, RC,
-                                                SpillSize, RI, *MFI);
+  bool HasVectorSuperClass = MF->getSubtarget<GCNSubtarget>().hasMAIInsts();
+  unsigned Opcode = getVectorRegSpillSaveOpcode(
+      VReg ? VReg : SrcReg, RC, SpillSize, RI, *MFI, HasVectorSuperClass);
   MFI->setHasSpilledVGPRs();
 
   BuildMI(MBB, MI, DL, get(Opcode))
     .addReg(SrcReg, getKillRegState(isKill)) // data
     .addFrameIndex(FrameIndex)               // addr
     .addReg(MFI->getStackPtrOffsetReg())     // scratch_offset
     .addImm(0)                               // offset
     .addMemOperand(MMO);
case 64:
     return AMDGPU::SI_SPILL_AV512_RESTORE;
   case 128:
     return AMDGPU::SI_SPILL_AV1024_RESTORE;
   default:
     llvm_unreachable("unknown register size");
   }
 }
 
-static unsigned getWWMRegSpillRestoreOpcode(unsigned Size) {
+static unsigned getWWMRegSpillRestoreOpcode(unsigned Size,
+                                            bool HasVectorSuperClass) {
   // Currently, there is only 32-bit WWM register spills needed.
   if (Size != 4)
     llvm_unreachable("unknown wwm register spill size");
 
+  // Return the AV pseudo opcode for gfx908+.
+  if (HasVectorSuperClass)
+    return AMDGPU::SI_SPILL_WWM_AV32_RESTORE;
+
   return AMDGPU::SI_SPILL_WWM_V32_RESTORE;
 }
 
-static unsigned
-getVectorRegSpillRestoreOpcode(Register Reg, const TargetRegisterClass *RC,
-                               unsigned Size, const SIRegisterInfo &TRI,
-                               const SIMachineFunctionInfo &MFI) {
+static unsigned getVectorRegSpillRestoreOpcode(Register Reg,
+                                               const TargetRegisterClass *RC,
+                                               unsigned Size,
+                                               const SIRegisterInfo &TRI,
+                                               const SIMachineFunctionInfo &MFI,
+                                               bool HasVectorSuperClass) {
   // Choose the right opcode if restoring a WWM register.
   if (MFI.checkFlag(Reg, AMDGPU::VirtRegFlag::WWM_REG))
-    return getWWMRegSpillRestoreOpcode(Size);
+    return getWWMRegSpillRestoreOpcode(Size, HasVectorSuperClass);
 
   if (TRI.isVectorSuperClass(RC))
     return getAVSpillRestoreOpcode(Size);
 
   return TRI.isAGPRClass(RC) ? getAGPRSpillRestoreOpcode(Size)
                              : getVGPRSpillRestoreOpcode(Size);
 }
 
if (RI.isSGPRClass(RC)) {
     BuildMI(MBB, MI, DL, OpDesc, DestReg)
       .addFrameIndex(FrameIndex) // addr
       .addMemOperand(MMO)
       .addReg(MFI->getStackPtrOffsetReg(), RegState::Implicit);
 
     return;
   }
 
-  unsigned Opcode = getVectorRegSpillRestoreOpcode(VReg ? VReg : DestReg, RC,
-                                                   SpillSize, RI, *MFI);
+  bool HasVectorSuperClass = MF->getSubtarget<GCNSubtarget>().hasMAIInsts();
+  unsigned Opcode = getVectorRegSpillRestoreOpcode(
+      VReg ? VReg : DestReg, RC, SpillSize, RI, *MFI, HasVectorSuperClass);
   BuildMI(MBB, MI, DL, get(Opcode), DestReg)
       .addFrameIndex(FrameIndex)           // vaddr
       .addReg(MFI->getStackPtrOffsetReg()) // scratch_offset
       .addImm(0)                           // offset
       .addMemOperand(MMO);
 }
 
 void SIInstrInfo::insertNoop(MachineBasicBlock &MBB,

llvm/lib/Target/AMDGPU/SIInstructions.td

 defm SI_SPILL_AV256 : SI_SPILL_VGPR <AV_256, 1>;
 defm SI_SPILL_AV288 : SI_SPILL_VGPR <AV_288, 1>;
 defm SI_SPILL_AV320 : SI_SPILL_VGPR <AV_320, 1>;
 defm SI_SPILL_AV352 : SI_SPILL_VGPR <AV_352, 1>;
 defm SI_SPILL_AV384 : SI_SPILL_VGPR <AV_384, 1>;
 defm SI_SPILL_AV512 : SI_SPILL_VGPR <AV_512, 1>;
 defm SI_SPILL_AV1024 : SI_SPILL_VGPR <AV_1024, 1>;
 
-let isConvergent = 1 in
-defm SI_SPILL_WWM_V32 : SI_SPILL_VGPR <VGPR_32>;
+let isConvergent = 1 in {
+  defm SI_SPILL_WWM_V32  : SI_SPILL_VGPR <VGPR_32>;
+  defm SI_SPILL_WWM_AV32 : SI_SPILL_VGPR <AV_32, 1>;
+}
 
 def SI_PC_ADD_REL_OFFSET : SPseudoInstSI <
   (outs SReg_64:$dst),
   (ins si_ga:$ptr_lo, si_ga:$ptr_hi),
   [(set SReg_64:$dst,
       (i64 (SIpc_add_rel_offset tglobaladdr:$ptr_lo, tglobaladdr:$ptr_hi)))]> {
   let Defs = [SCC];
 }

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

static unsigned getNumSubRegsForSpillOp(unsigned Op) {
   case AMDGPU::SI_SPILL_V32_SAVE:
   case AMDGPU::SI_SPILL_V32_RESTORE:
   case AMDGPU::SI_SPILL_A32_SAVE:
   case AMDGPU::SI_SPILL_A32_RESTORE:
   case AMDGPU::SI_SPILL_AV32_SAVE:
   case AMDGPU::SI_SPILL_AV32_RESTORE:
   case AMDGPU::SI_SPILL_WWM_V32_SAVE:
   case AMDGPU::SI_SPILL_WWM_V32_RESTORE:
+  case AMDGPU::SI_SPILL_WWM_AV32_SAVE:
+  case AMDGPU::SI_SPILL_WWM_AV32_RESTORE:
     return 1;
   default: llvm_unreachable("Invalid spill opcode");
   }
 }
 
 static int getOffsetMUBUFStore(unsigned Opc) {
   switch (Opc) {
   case AMDGPU::BUFFER_STORE_DWORD_OFFEN:
switch (MI->getOpcode()) {
     case AMDGPU::SI_SPILL_AV256_SAVE:
     case AMDGPU::SI_SPILL_AV224_SAVE:
     case AMDGPU::SI_SPILL_AV192_SAVE:
     case AMDGPU::SI_SPILL_AV160_SAVE:
     case AMDGPU::SI_SPILL_AV128_SAVE:
     case AMDGPU::SI_SPILL_AV96_SAVE:
     case AMDGPU::SI_SPILL_AV64_SAVE:
     case AMDGPU::SI_SPILL_AV32_SAVE:
-    case AMDGPU::SI_SPILL_WWM_V32_SAVE: {
+    case AMDGPU::SI_SPILL_WWM_V32_SAVE:
+    case AMDGPU::SI_SPILL_WWM_AV32_SAVE: {
       const MachineOperand *VData = TII->getNamedOperand(*MI,
                                                          AMDGPU::OpName::vdata);
       assert(TII->getNamedOperand(*MI, AMDGPU::OpName::soffset)->getReg() ==
              MFI->getStackPtrOffsetReg());
 
       unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_STORE_DWORD_SADDR
                                             : AMDGPU::BUFFER_STORE_DWORD_OFFSET;
       auto *MBB = MI->getParent();
switch (MI->getOpcode()) {
     case AMDGPU::SI_SPILL_AV224_RESTORE:
     case AMDGPU::SI_SPILL_AV256_RESTORE:
     case AMDGPU::SI_SPILL_AV288_RESTORE:
     case AMDGPU::SI_SPILL_AV320_RESTORE:
     case AMDGPU::SI_SPILL_AV352_RESTORE:
     case AMDGPU::SI_SPILL_AV384_RESTORE:
     case AMDGPU::SI_SPILL_AV512_RESTORE:
     case AMDGPU::SI_SPILL_AV1024_RESTORE:
-    case AMDGPU::SI_SPILL_WWM_V32_RESTORE: {
+    case AMDGPU::SI_SPILL_WWM_V32_RESTORE:
+    case AMDGPU::SI_SPILL_WWM_AV32_RESTORE: {
       const MachineOperand *VData = TII->getNamedOperand(*MI,
                                                          AMDGPU::OpName::vdata);
       assert(TII->getNamedOperand(*MI, AMDGPU::OpName::soffset)->getReg() ==
              MFI->getStackPtrOffsetReg());
 
       unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_LOAD_DWORD_SADDR
                                             : AMDGPU::BUFFER_LOAD_DWORD_OFFSET;
       auto *MBB = MI->getParent();

llvm/test/CodeGen/AMDGPU/whole-wave-register-spill.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 --stop-after=greedy,1 --verify-machineinstrs -o - %s | FileCheck -check-prefix=GFX906-GREEDY %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 --stop-after=greedy,1 --verify-machineinstrs -o - %s | FileCheck -check-prefix=GFX908-GREEDY %s
+
 ; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 --verify-machineinstrs -o - %s | FileCheck -check-prefix=GCN %s
 ; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 -O0 --verify-machineinstrs -o - %s | FileCheck -check-prefix=GCN-O0 %s
 
 ; Test whole-wave register spilling.
 
 ; In this testcase, the return address registers, PC value (SGPR30_SGPR31) and the scratch SGPR used in
 ; the inline asm statements should be preserved across the call. Since the test limits the VGPR numbers,
 ; the PC will be spilled to the only available CSR VGPR (VGPR40) as we spill CSR SGPRs including the PC
 ; directly to the physical VGPR lane to correctly generate the CFIs. The SGPR20 will get spilled to the
 ; virtual VGPR lane and that would be allocated by regalloc. Since there is no free VGPR to allocate, RA
 ; must spill a scratch VGPR. The writelane/readlane instructions that spill/restore SGPRs into/from VGPR
 ; are whole-wave operations and hence the VGPRs involved in such operations require whole-wave spilling.
 
 define void @test() #0 {
+; GFX906-GREEDY-LABEL: name: test
+; GFX906-GREEDY: bb.0 (%ir-block.0):
+; GFX906-GREEDY: %[[WWMREG:[0-9]+]]:vgpr_32 = V_WRITELANE_B32 killed $sgpr16, 0, %[[WWMREG]]
+; GFX906-GREEDY: SI_SPILL_WWM_V32_SAVE %[[WWMREG]], {{.*}}, $sgpr32, 0, implicit $exec :: (store (s32) into {{.*}}, addrspace 5)
+; GFX906-GREEDY: %{{[0-9]+}}:vgpr_32 = SI_SPILL_WWM_V32_RESTORE {{.*}}, $sgpr32, 0, implicit $exec :: (load (s32) from {{.*}}, addrspace 5)
+
+; GFX908-GREEDY-LABEL: name: test
+; GFX908-GREEDY: bb.0 (%ir-block.0):
+; GFX908-GREEDY: %[[WWMREG:[0-9]+]]:vgpr_32 = V_WRITELANE_B32 killed $sgpr16, 0, %[[WWMREG]]
+; GFX908-GREEDY: SI_SPILL_WWM_AV32_SAVE %[[WWMREG]], {{.*}}, $sgpr32, 0, implicit $exec :: (store (s32) into {{.*}}, addrspace 5)
+; GFX908-GREEDY: %{{[0-9]+}}:vgpr_32 = SI_SPILL_WWM_AV32_RESTORE {{.*}}, $sgpr32, 0, implicit $exec :: (load (s32) from {{.*}}, addrspace 5)
+
 ; GCN-LABEL: test:
 ; GCN:       ; %bb.0:
 ; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GCN-NEXT:    s_mov_b32 s16, s33
 ; GCN-NEXT:    s_mov_b32 s33, s32
 ; GCN-NEXT:    s_xor_saveexec_b64 s[18:19], -1
 ; GCN-NEXT:    buffer_store_dword v0, off, s[0:3], s33 offset:8 ; 4-byte Folded Spill
 ; GCN-NEXT:    buffer_store_dword v1, off, s[0:3], s33 offset:12 ; 4-byte Folded Spill