AMDGPU: Make max workgroup size settable by subtarget features

[arsenm]

Summary

This generates 32 subtarget features for each possible multiple
of the wave size. The attribute is named +max-work-group-size=N.
The use of an equal sign might be weird because no other target
seems to do this, but it is prettier.

[• tstellarAMD]

Why isn't this a function attribute?

[arsenm]

In D16834#342478, @tstellarAMD wrote:

Why isn't this a function attribute?

This is a property of the runtime. There should be function attributes for related information, such as required workgroup size. There may also be an attribute for maximum to assume but not required, but the subtarget feature would still be needed to validate that it is not above the maximum supported by the runtime, and would also serve as the default if that attribute is missing.

[arsenm]

In D16834#342542, @arsenm wrote:

In D16834#342478, @tstellarAMD wrote:

Why isn't this a function attribute?

This is a property of the runtime. There should be function attributes for related information, such as required workgroup size. There may also be an attribute for maximum to assume but not required, but the subtarget feature would still be needed to validate that it is not above the maximum supported by the runtime, and would also serve as the default if that attribute is missing.

Actually the higher workgroups might technically work with the runtime, they just aren't supported. I think the way it works in AMD OpenCL is the max 256 is rejected by default, but there is an environment variable to increase the limit for internal development.

[echristo]

How about just making it a normal string attribute similar to use-soft-float or something like that? Then you can just look it up on the function in the backend and won't need to configure a large set of subtarget features for it?

-eric

[arsenm]

This was already done differently

lib/Target/AMDGPU/AMDGPU.td

 
 def FeatureEnableHugeScratchBuffer : SubtargetFeature<
   "huge-scratch-buffer",
   "EnableHugeScratchBuffer",
   "true",
   "Enable scratch buffer sizes greater than 128 GB"
 >;
 
+class MaxWorkGroupSizeFeature<int size> : SubtargetFeature<
+  "max-workgroup-size="#size,
+  "MaxWorkGroupSize",
+   !cast<string>(size),
+  "Set max workgroup size to "#size#" ("#
+    !cast<string>(!srl(size, 6))#" wavefronts)"
+>;
+
+// Theoretically every multiple of 64 up to the maximum 2048 can be
+// used as a workgroup size.
+//XXX - I'm not sure what the maximum size is for R600.
+foreach NumWaves = 1-32 in {
+  def FeatureMaxWorkGroupSize#!shl(NumWaves, 6) :
+    MaxWorkGroupSizeFeature<!shl(NumWaves, 6)>;
+}
+
 def FeatureVGPRSpilling : SubtargetFeature<"vgpr-spilling",
   "EnableVGPRSpilling",
   "true",
   "Enable spilling of VGPRs to scratch memory"
 >;
 
 def FeatureDumpCode : SubtargetFeature <"DumpCode",
   "DumpCode",

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

 // FIXME: This can create globals so should be a module pass.
 class AMDGPUPromoteAlloca : public FunctionPass,
                             public InstVisitor<AMDGPUPromoteAlloca> {
 private:
   const TargetMachine *TM;
   Module *Mod;
   const DataLayout *DL;
   MDNode *MaxWorkGroupSizeRange;
-  int MaxWorkGroupSize;
+  int32_t MaxWorkGroupSize;
 
   // FIXME: This should be per-kernel.
-  int LocalMemAvailable;
+  int32_t LocalMemAvailable;
 
   bool IsAMDGCN;
   bool IsAMDHSA;
 
   std::pair<Value *, Value *> getLocalSizeYZ(IRBuilder<> &Builder);
   Value *getWorkitemID(IRBuilder<> &Builder, unsigned N);
 
 public:
void AMDGPUPromoteAlloca::visitAlloca(AllocaInst &I) {
 
   DEBUG(dbgs() << " alloca is not a candidate for vectorization.\n");
 
   // FIXME: This is the maximum work group size.  We should try to get
   // value from the reqd_work_group_size function attribute if it is
   // available.
   uint64_t AllocaSize = MaxWorkGroupSize * DL->getTypeAllocSize(AllocaTy);
 
-  if (AllocaSize > LocalMemAvailable) {
+  if (AllocaSize > static_cast<uint32_t>(LocalMemAvailable)) {
     DEBUG(dbgs() << " Not enough local memory to promote alloca.\n");
     return;
   }
 
   std::vector<Value*> WorkList;
 
   if (!collectUsesWithPtrTypes(&I, WorkList)) {
     DEBUG(dbgs() << " Do not know how to convert all uses\n");

lib/Target/AMDGPU/AMDGPUSubtarget.h

private:
   bool EnableIRStructurizer;
   bool EnablePromoteAlloca;
   bool EnableIfCvt;
   bool EnableLoadStoreOpt;
   bool EnableUnsafeDSOffsetFolding;
   bool EnableXNACK;
   unsigned WavefrontSize;
   bool CFALUBug;
-  uint32_t LocalMemorySize;
+  int32_t LocalMemorySize;
+  int32_t MaxWorkGroupSize;
   bool EnableVGPRSpilling;
   bool SGPRInitBug;
   bool IsGCN;
   bool GCN1Encoding;
   bool GCN3Encoding;
   bool CIInsts;
   bool FeatureDisable;
   int LDSBankCount;
public:
 
   unsigned getStackEntrySize() const;
 
   bool hasCFAluBug() const {
     assert(getGeneration() <= NORTHERN_ISLANDS);
     return CFALUBug;
   }
 
-  uint32_t getLocalMemorySize() const {
+  int32_t getLocalMemorySize() const {
     return LocalMemorySize;
   }
 
-  int getMaxWorkGroupSize() const {
-    // FIXME: The hardware maximum is 2048. The runtime usually enforces a 256
-    // limit though.
-    return 256;
+  int32_t getMaxWorkGroupSize() const {
+    return MaxWorkGroupSize;
   }
 
   int getHardwareMaxWorkGroupSize() const {
     return 2048;
   }
 
   bool hasSGPRInitBug() const {
     return SGPRInitBug;

lib/Target/AMDGPU/AMDGPUSubtarget.cpp

AMDGPUSubtarget::initializeSubtargetDependencies(const Triple &TT,
 
   SmallString<256> FullFS("+promote-alloca,+fp64-denormals,");
   if (isAmdHsaOS()) // Turn on FlatForGlobal for HSA.
     FullFS += "+flat-for-global,";
   FullFS += FS;
 
   ParseSubtargetFeatures(GPU, FullFS);
 
+  // If no specific size requested, set the default.
+  if (MaxWorkGroupSize == 0)
+    MaxWorkGroupSize = 256;
+
   // FIXME: I don't think think Evergreen has any useful support for
   // denormals, but should be checked. Should we issue a warning somewhere
   // if someone tries to enable these?
   if (getGeneration() <= AMDGPUSubtarget::NORTHERN_ISLANDS) {
     FP32Denormals = false;
     FP64Denormals = false;
   }
   return *this;
 }
 
 AMDGPUSubtarget::AMDGPUSubtarget(const Triple &TT, StringRef GPU, StringRef FS,
                                  TargetMachine &TM)
     : AMDGPUGenSubtargetInfo(TT, GPU, FS),
       DumpCode(false), R600ALUInst(false), HasVertexCache(false),
       TexVTXClauseSize(0), Gen(AMDGPUSubtarget::R600), FP64(false),
       FP64Denormals(false), FP32Denormals(false), FPExceptions(false),
       FastFMAF32(false), HalfRate64Ops(false), CaymanISA(false),
       FlatAddressSpace(false), FlatForGlobal(false), EnableIRStructurizer(true),
       EnablePromoteAlloca(false),
       EnableIfCvt(true), EnableLoadStoreOpt(false),
       EnableUnsafeDSOffsetFolding(false),
       EnableXNACK(false),
       WavefrontSize(0), CFALUBug(false),
-      LocalMemorySize(0),
+      LocalMemorySize(0), MaxWorkGroupSize(0),
       EnableVGPRSpilling(false), SGPRInitBug(false), IsGCN(false),
       GCN1Encoding(false), GCN3Encoding(false), CIInsts(false), LDSBankCount(0),
       IsaVersion(ISAVersion0_0_0), EnableHugeScratchBuffer(false),
       EnableSIScheduler(false), FrameLowering(nullptr),
       InstrItins(getInstrItineraryForCPU(GPU)), TargetTriple(TT) {
 
   initializeSubtargetDependencies(TT, GPU, FS);
 

test/CodeGen/AMDGPU/max-workgroup-size-attr.ll

+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=64" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE64 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=128" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE128 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=192" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE192 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=256" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE256 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=640" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE640 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=512" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE512 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=1024" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE1024 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=1600" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE1600 %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=1664" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=TOOBIG %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=1984" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=TOOBIG %s
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=2048" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=TOOBIG %s
+
+; Test if attribute is repeated with different sizes. Should be max
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -mattr="+max-workgroup-size=128,+max-workgroup-size=192,+max-workgroup-size=64" -amdgpu-promote-alloca < %s | FileCheck -check-prefix=MAXSIZE192 %s
+
+; MAXSIZE64: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [64 x [5 x i32]] undef, align 4
+; MAXSIZE128: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [128 x [5 x i32]] undef, align 4
+; MAXSIZE192: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [192 x [5 x i32]] undef, align 4
+; MAXSIZE256: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [256 x [5 x i32]] undef, align 4
+; MAXSIZE512: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [512 x [5 x i32]] undef, align 4
+; MAXSIZE640: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [640 x [5 x i32]] undef, align 4
+; MAXSIZE1024: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [1024 x [5 x i32]] undef, align 4
+; MAXSIZE1600: @promote_alloca_size.stack = internal unnamed_addr addrspace(3) global [1600 x [5 x i32]] undef, align 4
+
+
+; This exceeds maximum LDS size.
+; TOOBIG: alloca [5 x i32]
+
+; FIXME: We should use subtarget features for each function instead of many opt -S run lines
+define void @promote_alloca_size(i32 addrspace(1)* nocapture %out, i32 addrspace(1)* nocapture %in) #0 {
+entry:
+  %stack = alloca [5 x i32], align 4
+  %0 = load i32, i32 addrspace(1)* %in, align 4
+  %arrayidx1 = getelementptr inbounds [5 x i32], [5 x i32]* %stack, i32 0, i32 %0
+  store i32 4, i32* %arrayidx1, align 4
+  %arrayidx2 = getelementptr inbounds i32, i32 addrspace(1)* %in, i32 1
+  %1 = load i32, i32 addrspace(1)* %arrayidx2, align 4
+  %arrayidx3 = getelementptr inbounds [5 x i32], [5 x i32]* %stack, i32 0, i32 %1
+  store i32 5, i32* %arrayidx3, align 4
+  %arrayidx10 = getelementptr inbounds [5 x i32], [5 x i32]* %stack, i32 0, i32 0
+  %2 = load i32, i32* %arrayidx10, align 4
+  store i32 %2, i32 addrspace(1)* %out, align 4
+  %arrayidx12 = getelementptr inbounds [5 x i32], [5 x i32]* %stack, i32 0, i32 1
+  %3 = load i32, i32* %arrayidx12
+  %arrayidx13 = getelementptr inbounds i32, i32 addrspace(1)* %out, i32 1
+  store i32 %3, i32 addrspace(1)* %arrayidx13
+  ret void
+}
+
+attributes #0 = { nounwind }