[AMDGPU] Calculate number of min/max SGPRs/VGPRs for WavesPerEU instead of using switch statement

[arsenm]

LGTM

[tony-tye]

lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

211–220

Should this also have a calculation that takes into account if a trap handler is present?

232

Shouldn't this take into account the presence of a trap handler?

return std::min(MaxNumSGPRs - (TRAP_ENABLED ? TRAP_NUM_SGPRS : 0),
                AddressableNumSGPRs);

251–260

Since getMaxNumVGPRs() was updated to use a calculation should this one too?

[kzhuravl]

Address review comments.

lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

211–220

I do not think so. getMinNumSGPRs is used to bump up granulated sgpr count in pgm rsrc1 register to match maximum number of requested waves per execution unit.

232

It should. But it depends on trap handler subtarget feature from D26010, which has not landed yet.

[tony-tye]

lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

211–220

Right, but the value required to achieve that is based on determining how many registers the kernel must use. That value is reduced if there is a trap handler as the hardware will automatically allocate registers for the trap handler. It would also seem better if the hardcoded numbers here were replaced by the calculation that produces them which can then use the already existing target specific constants. If future hardware changes those limits then the calculation will still be correct.

[tony-tye]

lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

211–220

Sorry, missed that the function was updated with the calculation. So ignore that part: of the comment-)

[tony-tye]

LGTM

lib/Target/AMDGPU/AMDGPUSubtarget.cpp

if (Requested.second > getMaxFlatWorkGroupSize())
     return Default;
 
   return Requested;
 }
 
 std::pair<unsigned, unsigned> AMDGPUSubtarget::getWavesPerEU(
   const Function &F) const {
   // Default minimum/maximum number of waves per execution unit.
-  std::pair<unsigned, unsigned> Default(1, 0);
+  std::pair<unsigned, unsigned> Default(1, getMaxWavesPerEU());
 
   // Default/requested minimum/maximum flat work group sizes.
   std::pair<unsigned, unsigned> FlatWorkGroupSizes = getFlatWorkGroupSizes(F);
 
   // If minimum/maximum flat work group sizes were explicitly requested using
   // "amdgpu-flat-work-group-size" attribute, then set default minimum/maximum
   // number of waves per execution unit to values implied by requested
   // minimum/maximum flat work group sizes.

lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

if (Features.test(FeatureSGPRInitBug))
     return FIXED_NUM_SGPRS_FOR_INIT_BUG;
 
   IsaVersion Version = getIsaVersion(Features);
   if (Version.Major >= 8)
     return 102;
   return 104;
 }
 
 unsigned getMinNumSGPRs(const FeatureBitset &Features, unsigned WavesPerEU) {
-  IsaVersion Version = getIsaVersion(Features);
-  if (Version.Major >= 8) {
-    switch (WavesPerEU) {
-      case 0:  return 0;
-      case 10: return 0;
-      case 9:  return 0;
-      case 8:  return 81;
-      default: return 97;
-    }
-  } else {
-    switch (WavesPerEU) {
-      case 0:  return 0;
-      case 10: return 0;
-      case 9:  return 49;
-      case 8:  return 57;
-      case 7:  return 65;
-      case 6:  return 73;
-      case 5:  return 81;
-      default: return 97;
-    }
-  }
+  assert(WavesPerEU != 0);
+
+  if (WavesPerEU >= getMaxWavesPerEU(Features))
+    return 0;
+  unsigned MinNumSGPRs =
+      alignDown(getTotalNumSGPRs(Features) / (WavesPerEU + 1),
+                getSGPRAllocGranule(Features)) + 1;
+  return std::min(MinNumSGPRs, getAddressableNumSGPRs(Features));
 }
 
 unsigned getMaxNumSGPRs(const FeatureBitset &Features, unsigned WavesPerEU,
                         bool Addressable) {
+  assert(WavesPerEU != 0);
+
   IsaVersion Version = getIsaVersion(Features);
-  if (Version.Major >= 8) {
-    switch (WavesPerEU) {
-      case 0:  return 80;
-      case 10: return 80;
-      case 9:  return 80;
-      case 8:  return 96;
-      default: return Addressable ? getAddressableNumSGPRs(Features) : 112;
-    }
-  } else {
-    switch (WavesPerEU) {
-      case 0:  return 48;
-      case 10: return 48;
-      case 9:  return 56;
-      case 8:  return 64;
-      case 7:  return 72;
-      case 6:  return 80;
-      case 5:  return 96;
-      default: return getAddressableNumSGPRs(Features);
-    }
-  }
+  unsigned MaxNumSGPRs = alignDown(getTotalNumSGPRs(Features) / WavesPerEU,
+                                   getSGPRAllocGranule(Features));
+  unsigned AddressableNumSGPRs = getAddressableNumSGPRs(Features);
+  if (Version.Major >= 8 && !Addressable)
+    AddressableNumSGPRs = 112;
+  return std::min(MaxNumSGPRs, AddressableNumSGPRs);
 }
 
 unsigned getVGPRAllocGranule(const FeatureBitset &Features) {
   return 4;
 }
 
 unsigned getVGPREncodingGranule(const FeatureBitset &Features) {
   return getVGPRAllocGranule(Features);
 }
 
 unsigned getTotalNumVGPRs(const FeatureBitset &Features) {
   return 256;
 }
 
 unsigned getAddressableNumVGPRs(const FeatureBitset &Features) {
   return getTotalNumVGPRs(Features);
 }
 
 unsigned getMinNumVGPRs(const FeatureBitset &Features, unsigned WavesPerEU) {
-  switch (WavesPerEU) {
-    case 0:  return 0;
-    case 10: return 0;
-    case 9:  return 25;
-    case 8:  return 29;
-    case 7:  return 33;
-    case 6:  return 37;
-    case 5:  return 41;
-    case 4:  return 49;
-    case 3:  return 65;
-    case 2:  return 85;
-    default: return 129;
-  }
+  assert(WavesPerEU != 0);
+
+  if (WavesPerEU >= getMaxWavesPerEU(Features))
+    return 0;
+  unsigned MinNumVGPRs =
+      alignDown(getTotalNumVGPRs(Features) / (WavesPerEU + 1),
+                getVGPRAllocGranule(Features)) + 1;
+  return std::min(MinNumVGPRs, getAddressableNumVGPRs(Features));
 }
 
 unsigned getMaxNumVGPRs(const FeatureBitset &Features, unsigned WavesPerEU) {
-  switch (WavesPerEU) {
-    case 0:  return 24;
-    case 10: return 24;
-    case 9:  return 28;
-    case 8:  return 32;
-    case 7:  return 36;
-    case 6:  return 40;
-    case 5:  return 48;
-    case 4:  return 64;
-    case 3:  return 84;
-    case 2:  return 128;
-    default: return getTotalNumVGPRs(Features);
-  }
+  assert(WavesPerEU != 0);
+
+  unsigned MaxNumVGPRs = alignDown(getTotalNumVGPRs(Features) / WavesPerEU,
+                                   getVGPRAllocGranule(Features));
+  unsigned AddressableNumVGPRs = getAddressableNumVGPRs(Features);
+  return std::min(MaxNumVGPRs, AddressableNumVGPRs);
 }
 
 } // namespace IsaInfo
 
 void initDefaultAMDKernelCodeT(amd_kernel_code_t &Header,
                                const FeatureBitset &Features) {
   IsaInfo::IsaVersion ISA = IsaInfo::getIsaVersion(Features);
 

test/CodeGen/AMDGPU/attr-amdgpu-waves-per-eu.ll

 ; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=fiji -verify-machineinstrs < %s | FileCheck %s
 
 ; Exactly 1 wave per execution unit.
 ; CHECK-LABEL: {{^}}empty_exactly_1:
 ; CHECK: SGPRBlocks: 12
 ; CHECK: VGPRBlocks: 32
-; CHECK: NumSGPRsForWavesPerEU: 97
+; CHECK: NumSGPRsForWavesPerEU: 102
 ; CHECK: NumVGPRsForWavesPerEU: 129
 define void @empty_exactly_1() #0 {
 entry:
   ret void
 }
 attributes #0 = {"amdgpu-waves-per-eu"="1,1"}
 
 ; Exactly 5 waves per execution unit.
 ; CHECK-LABEL: {{^}}empty_exactly_5:
 ; CHECK: SGPRBlocks: 12
 ; CHECK: VGPRBlocks: 10
-; CHECK: NumSGPRsForWavesPerEU: 97
+; CHECK: NumSGPRsForWavesPerEU: 102
 ; CHECK: NumVGPRsForWavesPerEU: 41
 define void @empty_exactly_5() #1 {
 entry:
   ret void
 }
 attributes #1 = {"amdgpu-waves-per-eu"="5,5"}
 
 ; Exactly 10 waves per execution unit.

 attributes #5 = {"amdgpu-waves-per-eu"="10"}
 
 ; At most 1 wave per execution unit (same as @empty_exactly_1).
 
 ; At most 5 waves per execution unit.
 ; CHECK-LABEL: {{^}}empty_at_most_5:
 ; CHECK: SGPRBlocks: 12
 ; CHECK: VGPRBlocks: 10
-; CHECK: NumSGPRsForWavesPerEU: 97
+; CHECK: NumSGPRsForWavesPerEU: 102
 ; CHECK: NumVGPRsForWavesPerEU: 41
 define void @empty_at_most_5() #6 {
 entry:
   ret void
 }
 attributes #6 = {"amdgpu-waves-per-eu"="1,5"}
 
 ; At most 10 waves per execution unit.