AMDGPU: Serialize WWM registers

[arsenm]

5cd17f9d4356caafba9d10b45e929dcaa1acb37a

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

if (MFI->FrameOffsetReg != AMDGPU::FP_REG &&
     return diagnoseRegisterClass(YamlMFI.FrameOffsetReg);
   }
 
   if (MFI->StackPtrOffsetReg != AMDGPU::SP_REG &&
       !AMDGPU::SGPR_32RegClass.contains(MFI->StackPtrOffsetReg)) {
     return diagnoseRegisterClass(YamlMFI.StackPtrOffsetReg);
   }
 
+  for (const auto &YamlReg : YamlMFI.WWMReservedRegs) {
+    Register ParsedReg;
+    if (parseRegister(YamlReg, ParsedReg))
+      return true;
+
+    MFI->reserveWWMRegister(ParsedReg);
+  }
+
   auto parseAndCheckArgument = [&](const Optional<yaml::SIArgument> &A,
                                    const TargetRegisterClass &RC,
                                    ArgDescriptor &Arg, unsigned UserSGPRs,
                                    unsigned SystemSGPRs) {
     // Skip parsing if it's not present.
     if (!A)
       return false;
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

struct SIMachineFunctionInfo final : public yaml::MachineFunctionInfo {
   bool WaveLimiter = false;
   bool HasSpilledSGPRs = false;
   bool HasSpilledVGPRs = false;
   uint32_t HighBitsOf32BitAddress = 0;
 
   // TODO: 10 may be a better default since it's the maximum.
   unsigned Occupancy = 0;
 
+  SmallVector<StringValue> WWMReservedRegs;
+
   StringValue ScratchRSrcReg = "$private_rsrc_reg";
   StringValue FrameOffsetReg = "$fp_reg";
   StringValue StackPtrOffsetReg = "$sp_reg";
 
   Optional<SIArgumentInfo> ArgInfo;
   SIMode Mode;
   Optional<FrameIndex> ScavengeFI;
 
YamlIO.mapOptional("frameOffsetReg", MFI.FrameOffsetReg,
                        StringValue("$fp_reg"));
     YamlIO.mapOptional("stackPtrOffsetReg", MFI.StackPtrOffsetReg,
                        StringValue("$sp_reg"));
     YamlIO.mapOptional("argumentInfo", MFI.ArgInfo);
     YamlIO.mapOptional("mode", MFI.Mode, SIMode());
     YamlIO.mapOptional("highBitsOf32BitAddress",
                        MFI.HighBitsOf32BitAddress, 0u);
     YamlIO.mapOptional("occupancy", MFI.Occupancy, 0);
+    YamlIO.mapOptional("wwmReservedRegs", MFI.WWMReservedRegs);
     YamlIO.mapOptional("scavengeFI", MFI.ScavengeFI);
   }
 };
 
 } // end namespace yaml
 
 /// This class keeps track of the SPI_SP_INPUT_ADDR config register, which
 /// tells the hardware which interpolation parameters to load.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: ExplicitKernArgSize(MFI.getExplicitKernArgSize()),
       HasSpilledSGPRs(MFI.hasSpilledSGPRs()),
       HasSpilledVGPRs(MFI.hasSpilledVGPRs()),
       HighBitsOf32BitAddress(MFI.get32BitAddressHighBits()),
       Occupancy(MFI.getOccupancy()),
       ScratchRSrcReg(regToString(MFI.getScratchRSrcReg(), TRI)),
       FrameOffsetReg(regToString(MFI.getFrameOffsetReg(), TRI)),
       StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)),
       ArgInfo(convertArgumentInfo(MFI.getArgInfo(), TRI)), Mode(MFI.getMode()) {
+  for (Register Reg : MFI.WWMReservedRegs)
+    WWMReservedRegs.push_back(regToString(Reg, TRI));
+
   auto SFI = MFI.getOptionalScavengeFI();
   if (SFI)
     ScavengeFI = yaml::FrameIndex(*SFI, MF.getFrameInfo());
 }
 
 void yaml::SIMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
   MappingTraits<SIMachineFunctionInfo>::mapping(YamlIO, *this);
 }

llvm/test/CodeGen/MIR/AMDGPU/machine-function-info.ll

-; RUN: llc -mtriple=amdgcn-mesa-mesa3d -mcpu=tahiti -stop-after finalize-isel -o %t.mir %s
+; RUN: llc -mtriple=amdgcn-mesa-mesa3d -mcpu=tahiti -stop-after=si-pre-allocate-wwm-regs -o %t.mir %s
 ; RUN: llc -run-pass=none -verify-machineinstrs %t.mir -o - | FileCheck %s
 
 ; Test that SIMachineFunctionInfo can be round trip serialized through
 ; MIR.
 
 @lds = addrspace(3) global [512 x float] undef, align 4
 
 ; CHECK-LABEL: {{^}}name: kernel

 
 ; CHECK-LABEL: {{^}}name: high_address_bits
 ; CHECK: machineFunctionInfo:
 ; CHECK: highBitsOf32BitAddress: 4294934528
 define amdgpu_ps void @high_address_bits() #4 {
   ret void
 }
 
+; CHECK-LABEL: {{^}}name: wwm_reserved_regs
+; CHECK: wwmReservedRegs:
+; CHECK-NEXT: - '$vgpr2'
+; CHECK-NEXT: - '$vgpr3'
+define amdgpu_cs void @wwm_reserved_regs(i32 addrspace(1)* %ptr, <4 x i32> inreg %tmp14) {
+  %ld0 = load volatile i32, i32 addrspace(1)* %ptr
+  %ld1 = load volatile i32, i32 addrspace(1)* %ptr
+  %inactive0 = tail call i32 @llvm.amdgcn.set.inactive.i32(i32 %ld1, i32 0)
+  %inactive1 = tail call i32 @llvm.amdgcn.set.inactive.i32(i32 %ld0, i32 0)
+  store volatile i32 %inactive0, i32 addrspace(1)* %ptr
+  store volatile i32 %inactive1, i32 addrspace(1)* %ptr
+  ret void
+}
+
+declare i32 @llvm.amdgcn.set.inactive.i32(i32, i32) #5
+
 attributes #0 = { "no-signed-zeros-fp-math" = "true" }
 attributes #1 = { "amdgpu-dx10-clamp" = "false" }
 attributes #2 = { "amdgpu-ieee" = "false" }
 attributes #3 = { "amdgpu-dx10-clamp" = "false" "amdgpu-ieee" = "false" }
 attributes #4 = { "amdgpu-32bit-address-high-bits"="0xffff8000" }
+attributes #5 = { convergent nounwind readnone willreturn }

llvm/test/CodeGen/MIR/AMDGPU/wwm-reserved-regs-invalid-reg.mir

+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck -check-prefix=ERR %s
+
+---
+name: invalid_reg
+machineFunctionInfo:
+# ERR: [[@LINE+1]]:32: unknown register name 'notareg'
+  wwmReservedRegs: ['$vgpr0', '$notareg']
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

llvm/test/CodeGen/MIR/AMDGPU/wwm-reserved-regs-not-a-reg.mir

+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck -check-prefix=ERR %s
+
+---
+name: invalid_reg
+machineFunctionInfo:
+# ERR: [[@LINE+1]]:21: expected a named register
+  wwmReservedRegs: [123]
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

llvm/test/CodeGen/MIR/AMDGPU/wwm-reserved-regs.mir

+# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o - | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -simplify-mir -verify-machineinstrs %s -o - | FileCheck %s
+
+# CHECK-LABEL: name: empty_wwm_regs{{$}}
+# CHECK: machineFunctionInfo:
+# CHECK-NOT: wwmReservedRegs
+---
+name: empty_wwm_regs
+machineFunctionInfo:
+  wwmReservedRegs: []
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...
+
+# CHECK-LABEL: name: one_reg{{$}}
+# CHECK: machineFunctionInfo:
+# CHECK: wwmReservedRegs:
+# CHECK-NEXT: - '$vgpr0'
+---
+name: one_reg
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr0']
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...
+
+# CHECK-LABEL: name: two_reg{{$}}
+# CHECK: machineFunctionInfo:
+# CHECK: wwmReservedRegs:
+# CHECK-NEXT: - '$vgpr0'
+# CHECK-NEXT: - '$vgpr1'
+---
+name: two_reg
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr0', '$vgpr1']
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...