[AMDGPU] Occupancy w.r.t. workgroup size range is also a range

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

@@ -1175,22 +1175,10 @@ void AMDGPUAsmPrinter::getSIProgramInfo(SIProgramInfo &ProgInfo,
   // Make clamp modifier on NaN input returns 0.
   ProgInfo.DX10Clamp = Mode.DX10Clamp;
 
-  unsigned LDSAlignShift;
-  if (STM.getFeatureBits().test(FeatureAddressableLocalMemorySize163840)) {
-    // LDS is allocated in 320 dword blocks.
-    LDSAlignShift = 11;
-  } else if (STM.getFeatureBits().test(
-                 FeatureAddressableLocalMemorySize65536)) {
-    // LDS is allocated in 128 dword blocks.
-    LDSAlignShift = 9;
-  } else {
-    // LDS is allocated in 64 dword blocks.
-    LDSAlignShift = 8;
-  }
-
   ProgInfo.SGPRSpill = MFI->getNumSpilledSGPRs();
   ProgInfo.VGPRSpill = MFI->getNumSpilledVGPRs();
 
+  unsigned LDSAlignShift = Log2_32_Ceil(STM.getLDSAllocGranularity());
   ProgInfo.LDSSize = MFI->getLDSSize();
   ProgInfo.LDSBlocks =
       alignTo(ProgInfo.LDSSize, 1ULL << LDSAlignShift) >> LDSAlignShift;

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

@@ -1344,7 +1344,7 @@ bool AMDGPUPromoteAllocaImpl::hasSufficientLocalMem(const Function &F) {
   }
 
   unsigned MaxOccupancy =
-      ST.getOccupancyWithWorkGroupSizes(CurrentLocalMemUsage, F).second;
+      ST.getOccupancyWithWorkGroupSizes(CurrentLocalMemUsage, F, TM).second;
 
   // Restrict local memory usage so that we don't drastically reduce occupancy,
   // unless it is already significantly reduced.

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

@@ -55,13 +55,15 @@ AMDGPUSubtarget::getMaxLocalMemSizeWithWaveCount(unsigned NWaves,
   return getLocalMemorySize() / WorkGroupsPerCU;
 }
 
-std::pair<unsigned, unsigned>
-AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(uint32_t LDSBytes,
-                                                const Function &F) const {
-  // FIXME: Is there an allocation granularity for the LDS? If so we would need
-  // to make sure the amount of bytes is aligned on that granularity.
-
+std::pair<unsigned, unsigned> AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(
+    uint32_t LDSBytes, const Function &F, const TargetMachine &TM) const {
   // Compute occupancy restriction based on LDS usage.
+  if (TM.getTargetTriple().getArch() == Triple::amdgcn) {
+    // For GCN subtargets, LDS size must be aligned on allocation granularity.
+    const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);
+    LDSBytes = alignTo(LDSBytes, ST.getLDSAllocGranularity());
+  }
+
   const unsigned MaxWGsLDS = getLocalMemorySize() / std::max(LDSBytes, 1u);
 
   // Queried LDS size may be larger than available on a CU, in which case we
@@ -72,9 +74,8 @@ AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(uint32_t LDSBytes,
     return {1, 1};
 
   const unsigned WaveSize = getWavefrontSize(), WavesPerEU = getMaxWavesPerEU();
-  const unsigned WaveSlotsPerCU = WavesPerEU * getEUsPerCU();
 
-  auto PropsFromWGSize = [&](unsigned WGSize)
+  auto PropsFromWGSize = [=](unsigned WGSize)
       -> std::tuple<const unsigned, const unsigned, unsigned> {
     unsigned WavesPerWG = divideCeil(WGSize, WaveSize);
     unsigned WGsPerCU = std::min(getMaxWorkGroupsPerCU(WGSize), MaxWGsLDS);
@@ -91,10 +92,12 @@ AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(uint32_t LDSBytes,
 
   // It is possible that we end up with flipped minimum and maximum number of
   // waves per CU when the number of minimum/maximum concurrent groups on the CU
-  // is limited by LDS usage or barrier ressources.
+  // is limited by LDS usage or barrier resources.
   if (MinWavesPerCU >= MaxWavesPerCU) {
     std::swap(MinWavesPerCU, MaxWavesPerCU);
   } else {
+    const unsigned WaveSlotsPerCU = WavesPerEU * getEUsPerCU();
+
     // Look for a potential smaller group size than the maximum which decreases
     // the concurrent number of waves on the CU for the same number of
     // concurrent workgroups on the CU.
@@ -140,7 +143,8 @@ AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(uint32_t LDSBytes,
 std::pair<unsigned, unsigned> AMDGPUSubtarget::getOccupancyWithWorkGroupSizes(
     const MachineFunction &MF) const {
   const auto *MFI = MF.getInfo<SIMachineFunctionInfo>();
-  return getOccupancyWithWorkGroupSizes(MFI->getLDSSize(), MF.getFunction());
+  return getOccupancyWithWorkGroupSizes(MFI->getLDSSize(), MF.getFunction(),
+                                        MF.getTarget());
 }
 
 std::pair<unsigned, unsigned>

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h

@@ -133,7 +133,8 @@ class AMDGPUSubtarget {
   /// This notably depends on the range of allowed flat group sizes for the
   /// function and hardware characteristics.
   std::pair<unsigned, unsigned>
-  getOccupancyWithWorkGroupSizes(uint32_t LDSBytes, const Function &F) const;
+  getOccupancyWithWorkGroupSizes(uint32_t LDSBytes, const Function &F,
+                                 const TargetMachine &TM) const;
 
   /// Subtarget's minimum/maximum occupancy, in number of waves per EU, that can
   /// be achieved when the only function running on a CU is \p MF. This notably

llvm/lib/Target/AMDGPU/GCNSubtarget.cpp

@@ -408,7 +408,8 @@ unsigned GCNSubtarget::getReservedNumSGPRs(const Function &F) const {
 std::pair<unsigned, unsigned>
 GCNSubtarget::computeOccupancy(const Function &F, unsigned LDSSize,
                                unsigned NumSGPRs, unsigned NumVGPRs) const {
-  auto [MinOcc, MaxOcc] = getOccupancyWithWorkGroupSizes(LDSSize, F);
+  auto [MinOcc, MaxOcc] =
+      getOccupancyWithWorkGroupSizes(LDSSize, F, TLInfo.getTargetMachine());
   unsigned SGPROcc = getOccupancyWithNumSGPRs(NumSGPRs);
   unsigned VGPROcc = getOccupancyWithNumVGPRs(NumVGPRs);
 
@@ -417,6 +418,14 @@ GCNSubtarget::computeOccupancy(const Function &F, unsigned LDSSize,
   return {std::min(MinOcc, MaxOcc), MaxOcc};
 }
 
+unsigned GCNSubtarget::getLDSAllocGranularity() const {
+  if (getFeatureBits().test(AMDGPU::FeatureAddressableLocalMemorySize163840))
+    return 1280; // LDS is allocated in 320 dword blocks.
+  if (getFeatureBits().test(AMDGPU::FeatureAddressableLocalMemorySize65536))
+    return 512; // LDS is allocated in 128 dword blocks.
+  return 256;   // LDS is allocated in 64 dword blocks.
+}
+
 unsigned GCNSubtarget::getBaseMaxNumSGPRs(
     const Function &F, std::pair<unsigned, unsigned> WavesPerEU,
     unsigned PreloadedSGPRs, unsigned ReservedNumSGPRs) const {

llvm/lib/Target/AMDGPU/GCNSubtarget.h

@@ -1381,6 +1381,9 @@ class GCNSubtarget final : public AMDGPUGenSubtargetInfo,
                                                  unsigned NumSGPRs = 0,
                                                  unsigned NumVGPRs = 0) const;
 
+  /// Returns the LDS's allocation granularity in bytes.
+  unsigned getLDSAllocGranularity() const;
+
   /// \returns true if the flat_scratch register should be initialized with the
   /// pointer to the wave's scratch memory rather than a size and offset.
   bool flatScratchIsPointer() const {

llvm/test/CodeGen/AMDGPU/agpr-copy-no-free-registers.ll

@@ -365,7 +365,10 @@ define amdgpu_kernel void @no_agpr_no_reserve(ptr addrspace(1) %arg) #0 {
   ret void
 }
 
-define void @v32_asm_def_use(float %v0, float %v1) #0 {
+; FIXME: This case is broken. The asm value passed in v32 is live
+; through the range where the reserved def for the copy is introduced,
+; clobbering the user value.
+define void @v32_asm_def_use(float %v0, float %v1) #4 {
 ; GFX908-LABEL: v32_asm_def_use:
 ; GFX908:       ; %bb.0:
 ; GFX908-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
@@ -374,48 +377,57 @@ define void @v32_asm_def_use(float %v0, float %v1) #0 {
 ; GFX908-NEXT:    ;;#ASMSTART
 ; GFX908-NEXT:    ; def v[0:31] a[0:15]
 ; GFX908-NEXT:    ;;#ASMEND
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a15
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a15
+; GFX908-NEXT:    ;;#ASMSTART
+; GFX908-NEXT:    ; def v32
+; GFX908-NEXT:    ;;#ASMEND
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a31, v35
 ; GFX908-NEXT:    v_accvgpr_read_b32 v35, a14
-; GFX908-NEXT:    s_nop 0
-; GFX908-NEXT:    v_accvgpr_write_b32 a31, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a13
+; GFX908-NEXT:    s_nop 1
 ; GFX908-NEXT:    v_accvgpr_write_b32 a30, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a13
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a29, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a12
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a28, v35
 ; GFX908-NEXT:    v_accvgpr_read_b32 v35, a11
-; GFX908-NEXT:    v_accvgpr_write_b32 a29, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a12
+; GFX908-NEXT:    s_nop 1
 ; GFX908-NEXT:    v_accvgpr_write_b32 a27, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a10
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a26, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a9
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a25, v35
 ; GFX908-NEXT:    v_accvgpr_read_b32 v35, a8
-; GFX908-NEXT:    v_accvgpr_write_b32 a28, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a10
+; GFX908-NEXT:    s_nop 1
 ; GFX908-NEXT:    v_accvgpr_write_b32 a24, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a7
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a23, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a6
+; GFX908-NEXT:    s_nop 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a22, v35
 ; GFX908-NEXT:    v_accvgpr_read_b32 v35, a5
-; GFX908-NEXT:    v_accvgpr_write_b32 a26, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a9
+; GFX908-NEXT:    s_nop 1
 ; GFX908-NEXT:    v_accvgpr_write_b32 a21, v35
-; GFX908-NEXT:    v_accvgpr_read_b32 v35, a2
-; GFX908-NEXT:    v_accvgpr_write_b32 a25, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a7
-; GFX908-NEXT:    v_accvgpr_write_b32 a18, v35
-; GFX908-NEXT:    s_nop 0
-; GFX908-NEXT:    v_accvgpr_write_b32 a23, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a6
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a4
 ; GFX908-NEXT:    s_nop 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a22, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a4
+; GFX908-NEXT:    v_accvgpr_write_b32 a20, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a3
 ; GFX908-NEXT:    s_nop 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a20, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a3
+; GFX908-NEXT:    v_accvgpr_write_b32 a19, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a2
 ; GFX908-NEXT:    s_nop 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a19, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a1
+; GFX908-NEXT:    v_accvgpr_write_b32 a18, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a1
 ; GFX908-NEXT:    s_nop 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a17, v32
-; GFX908-NEXT:    v_accvgpr_read_b32 v32, a0
+; GFX908-NEXT:    v_accvgpr_write_b32 a17, v35
+; GFX908-NEXT:    v_accvgpr_read_b32 v35, a0
 ; GFX908-NEXT:    s_nop 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a16, v32
-; GFX908-NEXT:    ;;#ASMSTART
-; GFX908-NEXT:    ; def v32
-; GFX908-NEXT:    ;;#ASMEND
+; GFX908-NEXT:    v_accvgpr_write_b32 a16, v35
 ; GFX908-NEXT:    ;;#ASMSTART
 ; GFX908-NEXT:    ; copy
 ; GFX908-NEXT:    ;;#ASMEND
@@ -1133,3 +1145,4 @@ attributes #0 = { "amdgpu-waves-per-eu"="6,6" }
 attributes #1 = { convergent nounwind readnone willreturn }
 attributes #2 = { nounwind readnone willreturn }
 attributes #3 = { "amdgpu-waves-per-eu"="7,7" }
+attributes #4 = { "amdgpu-waves-per-eu"="6,6" "amdgpu-flat-work-group-size"="1024,1024" }