[AMDGPU] Ensure there are enough registers for wave dispatch

[tpr]

Summary

This fixes the number of SGPRs and VGPRs in the *_RSRC1 register to
allow for registers set up in wave dispatch, even if those registers are
not used in the shader.

Change-Id: I6575f0e0d2a528d1319d0b289f0ebe4510fa5771

[nhaehnle]

That seems reasonable.

[arsenm]

I think this is the wrong way to solve this. I know I have a patch where I started to fix this before. I don’t like inspecting the IR types in the AsmPrinter to determine the used number of registers. The way this is determined is more complicated than the type size in bits, so this really needs to determine this from what legalization decided to do.

[arsenm]

Arguably we could also just report the number used by the shader and it’s the runtime’s responsibility to round up from the number of registers used, but since we are sort of just emitting the raw config register value that doesn’t really work

[t-tye]

I believe the metadata does have the actual (ie non-granulated rounded up value) register count. So the runtime could access the information from there.

llvm/trunk/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

if (ProgInfo.NumSGPR > MaxAddressableNumSGPRs) {
       ProgInfo.NumSGPR = MaxAddressableNumSGPRs - 1;
     }
   }
 
   // Account for extra SGPRs and VGPRs reserved for debugger use.
   ProgInfo.NumSGPR += ExtraSGPRs;
   ProgInfo.NumVGPR += ExtraVGPRs;
 
+  // Ensure there are enough SGPRs and VGPRs for wave dispatch, where wave
+  // dispatch registers are function args.
+  unsigned WaveDispatchNumSGPR = 0, WaveDispatchNumVGPR = 0;
+  for (auto &Arg : MF.getFunction().args()) {
+    unsigned NumRegs = (Arg.getType()->getPrimitiveSizeInBits() + 31) / 32;
+    if (Arg.hasAttribute(Attribute::InReg))
+      WaveDispatchNumSGPR += NumRegs;
+    else
+      WaveDispatchNumVGPR += NumRegs;
+  }
+  ProgInfo.NumSGPR = std::max(ProgInfo.NumSGPR, WaveDispatchNumSGPR);
+  ProgInfo.NumVGPR = std::max(ProgInfo.NumVGPR, WaveDispatchNumVGPR);
+
   // Adjust number of registers used to meet default/requested minimum/maximum
   // number of waves per execution unit request.
   ProgInfo.NumSGPRsForWavesPerEU = std::max(
     std::max(ProgInfo.NumSGPR, 1u), STM.getMinNumSGPRs(MFI->getMaxWavesPerEU()));
   ProgInfo.NumVGPRsForWavesPerEU = std::max(
     std::max(ProgInfo.NumVGPR, 1u), STM.getMinNumVGPRs(MFI->getMaxWavesPerEU()));
 
   if (STM.getGeneration() <= AMDGPUSubtarget::SEA_ISLANDS ||

llvm/trunk/test/CodeGen/AMDGPU/wave_dispatch_regs.ll

+; RUN: llc -mtriple=amdgcn--amdpal -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=SI -enable-var-scope %s
+; RUN: llc -mtriple=amdgcn--amdpal -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=VI -enable-var-scope %s
+; RUN: llc -mtriple=amdgcn--amdpal -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=GFX9 -enable-var-scope %s
+
+; This compute shader has input args that claim that it has 17 sgprs and 5 vgprs
+; in wave dispatch. Ensure that the sgpr and vgpr counts in COMPUTE_PGM_RSRC1
+; are set to reflect that, even though the registers are not used in the shader.
+
+; GCN-LABEL: {{^}}_amdgpu_cs_main:
+; SI: .amd_amdgpu_pal_metadata{{.*}}0x2e12,0x{{[0-9a-f]*}}81,
+; VI: .amd_amdgpu_pal_metadata{{.*}}0x2e12,0x{{[0-9a-f]*}}c1,
+; GFX9: .amd_amdgpu_pal_metadata{{.*}}0x2e12,0x{{[0-9a-f]*}}81,
+
+define dllexport amdgpu_cs void @_amdgpu_cs_main(i32 inreg, i32 inreg, <2 x i32> inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, i32 inreg, <3 x i32> inreg, i32 inreg, <5 x i32>) {
+.entry:
+  ret void
+}
+