[AMDGPU] Machine scheduler should take care of the 'amdgpu-num-vgp' function attribute

[alex-t]

Summary

We currently have the 'amdgpu-num-vgpr' attribute parsed during the TargetRegisterInfo initialization.

GCNSchedStage class method computing occupancy does not read the attribute and consider all the available
 registers for the given target even if the number of VGPRs was explicitly restricted.

[kerbowa]

I thought this already was used in the excess VGPRs pressure and wave-per-eu is that just not true?

[rampitec]

In D159452#4638703, @kerbowa wrote:

I thought this already was used in the excess VGPRs pressure and wave-per-eu is that just not true?

I thought it is buried inside the logic of waves per eu calculation/getEffectiveWavesPerEU/getWavesPerEU, but apparently it is not.

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

677

'llvm::' is not needed.

682

Why not just call getOccupancyWithNumVGPRs here? This logic is already in the getNumWavesPerEUWithNumVGPRs.

[alex-t]

In D159452#4638703, @kerbowa wrote:

I thought this already was used in the excess VGPRs pressure and wave-per-eu is that just not true?

trying to give a detailed response:

We measure the occupancy after the current schedule stage by querying the RegisterPressure class:

unsigned WavesAfter =
      std::min(TargetOccupancy, PressureAfter.getOccupancy(ST));

which does not care about attributes set on function.
Then couple of lines after we query the right method which returns the real number of VGPRs available
for the function and we mark the region as excess RP accordingly:

unsigned MaxVGPRs = ST.getMaxNumVGPRs(MF);
if (PressureAfter.getVGPRNum(false) > MaxVGPRs ||
       *****
    DAG.RegionsWithExcessRP[RegionIdx] = true;
  }

the problem is that we decide if we need to revert the schedule based on both - excess RP and WavesAfter.
The latter is computed w/o the knowledge of the restrictions applied by the attribute:

bool GCNSchedStage::mayCauseSpilling(unsigned WavesAfter) {
  if (WavesAfter <= MFI.getMinWavesPerEU() &&
      !PressureAfter.less(ST, PressureBefore) &&
      isRegionWithExcessRP()) {
    LLVM_DEBUG(dbgs() << "New pressure will result in more spilling.\n");
    return true;
  }

So, we never revert the "bad" schedule!

[alex-t]

Thanks to all who took the time to look at this.
The aim of this review is to discuss and find the right approach.
We have a target description interface that keeps the HW details.
We also have a set of user-defined attributes like: "amdgpu-waves-per-eu", "amdgpu-num-vgpr" etc. which affect the defaults.
We also take into account the LDS size.
Shouldn't we compute the occupancy uniformly accounting for all constraints together?

[alex-t]

One more case that we may want to consider: with "amdgpu-num-vgpr"=8 the difference between 2 RPs of 2 is significant.
The schedule that spends 7 VGPRs does not spill but another one, which spends 9 will.
Currently, we have a "less" operator defined for the GCNRegPressure which only compares the occupancy based on the default target number of available registers and the number of VGPRs from the RP. To this method, RPs in my case are equal as they both have occupancy 10. This is yet another reason why we cannot use the excess RP set to decide if we need to keep or revert the current schedule.

[rampitec]

In D159452#4638728, @alex-t wrote:

Thanks to all who took the time to look at this.
The aim of this review is to discuss and find the right approach.
We have a target description interface that keeps the HW details.
We also have a set of user-defined attributes like: "amdgpu-waves-per-eu", "amdgpu-num-vgpr" etc. which affect the defaults.
We also take into account the LDS size.
Shouldn't we compute the occupancy uniformly accounting for all constraints together?

There is GCNSubtarget::computeOccupancy which aims to do exactly that.

In D159452#4638833, @alex-t wrote:

One more case that we may want to consider: with "amdgpu-num-vgpr"=8 the difference between 2 RPs of 2 is significant.
The schedule that spends 7 VGPRs does not spill but another one, which spends 9 will.
Currently, we have a "less" operator defined for the GCNRegPressure which only compares the occupancy based on the default target number of available registers and the number of VGPRs from the RP. To this method, RPs in my case are equal as they both have occupancy 10. This is yet another reason why we cannot use the excess RP set to decide if we need to keep or revert the current schedule.

8 vgpr case is artificial and only used for stress testing. There is no other reason to use anything below maximum occupancy.

[arsenm]

amdgpu-num-vgprs is deprecated. Do we really need to try handling it?

[alex-t]

In D159452#4638836, @rampitec wrote:

In D159452#4638728, @alex-t wrote:

Thanks to all who took the time to look at this.
The aim of this review is to discuss and find the right approach.
We have a target description interface that keeps the HW details.
We also have a set of user-defined attributes like: "amdgpu-waves-per-eu", "amdgpu-num-vgpr" etc. which affect the defaults.
We also take into account the LDS size.
Shouldn't we compute the occupancy uniformly accounting for all constraints together?

There is GCNSubtarget::computeOccupancy which aims to do exactly that.

No, it does not care of "amdgpu-num-vgpr" attribute.

if (NumVGPRs)
   Occupancy = std::min(Occupancy, getOccupancyWithNumVGPRs(NumVGPRs));

unsigned GCNSubtarget::getOccupancyWithNumVGPRs(unsigned NumVGPRs) const {
  return AMDGPU::IsaInfo::getNumWavesPerEUWithNumVGPRs(this, NumVGPRs);
}

unsigned getNumWavesPerEUWithNumVGPRs(const MCSubtargetInfo *STI,
                                      unsigned NumVGPRs) {
  unsigned MaxWaves = getMaxWavesPerEU(STI);
  unsigned Granule = getVGPRAllocGranule(STI);
  if (NumVGPRs < Granule)
    return MaxWaves;
  unsigned RoundedRegs = alignTo(NumVGPRs, Granule);
  return std::min(std::max(**getTotalNumVGPRs(STI) **/ RoundedRegs, 1u), MaxWaves);
}

getTotalNumVGPRs only takes the defaults.

In D159452#4638833, @alex-t wrote:

One more case that we may want to consider: with "amdgpu-num-vgpr"=8 the difference between 2 RPs of 2 is significant.
The schedule that spends 7 VGPRs does not spill but another one, which spends 9 will.
Currently, we have a "less" operator defined for the GCNRegPressure which only compares the occupancy based on the default target number of available registers and the number of VGPRs from the RP. To this method, RPs in my case are equal as they both have occupancy 10. This is yet another reason why we cannot use the excess RP set to decide if we need to keep or revert the current schedule.

8 vgpr case is artificial and only used for stress testing. There is no other reason to use anything below maximum occupancy.

But it is still a valid input and we don't want that tests to become spilling.
Another question: What this attribute is intended for? If it imposes the constraint we have to take it into account when computing the occupancy, otherwise, we schedule/allocate more registers than was requested by the user and this will be the UB.

[alex-t]

In D159452#4638868, @arsenm wrote:

amdgpu-num-vgprs is deprecated. Do we really need to try handling it?

That is a strong reason to abandon this. But in this case, we need to do something with the tests using it.

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

 //===-- AMDGPUSubtarget.cpp - AMDGPU Subtarget Information ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 using namespace llvm;
 
 #define DEBUG_TYPE "amdgpu-subtarget"
 
 #define GET_SUBTARGETINFO_TARGET_DESC
 #define GET_SUBTARGETINFO_CTOR
 #define AMDGPUSubtarget GCNSubtarget
 #include "AMDGPUGenSubtargetInfo.inc"
+#include "GCNSubtarget.h"
 #undef AMDGPUSubtarget
 
 static cl::opt<bool> EnablePowerSched(
   "amdgpu-enable-power-sched",
   cl::desc("Enable scheduling to minimize mAI power bursts"),
   cl::init(false));
 
 static cl::opt<bool> EnableVGPRIndexMode(
if (SGPRs <= 80)
     return 6;
   return 5;
 }
 
 unsigned GCNSubtarget::getOccupancyWithNumVGPRs(unsigned NumVGPRs) const {
   return AMDGPU::IsaInfo::getNumWavesPerEUWithNumVGPRs(this, NumVGPRs);
 }
 
+unsigned llvm::GCNSubtarget::getOccupancyWithNumVGPRsForFunction(
+    const MachineFunction &MF, unsigned VGPRs) const {
+  const Function &F = MF.getFunction();
+  std::pair<unsigned, unsigned> WavesPerEU = getWavesPerEU(F);
+  unsigned MaxWaves = WavesPerEU.second;
+  unsigned Granule = getVGPRAllocGranule();
+  if (VGPRs < Granule)
+    return MaxWaves;
+  unsigned RoundedRegs = alignTo(VGPRs, Granule);
+  unsigned RequestedPerAttr = getFunctionNVGPRAttr(F, WavesPerEU);
+  unsigned NumVGPRAvailable = getTotalNumVGPRs();
+  if (RequestedPerAttr && RequestedPerAttr < NumVGPRAvailable)
+    NumVGPRAvailable = RequestedPerAttr; 
+  return std::min(std::max(NumVGPRAvailable / RoundedRegs, 1u), MaxWaves);
+}
+
 unsigned
 GCNSubtarget::getBaseReservedNumSGPRs(const bool HasFlatScratch) const {
   if (getGeneration() >= AMDGPUSubtarget::GFX10)
     return 2; // VCC. FLAT_SCRATCH and XNACK are no longer in SGPRs.
 
   if (HasFlatScratch || HasArchitectedFlatScratch) {
     if (getGeneration() >= AMDGPUSubtarget::VOLCANIC_ISLANDS)
       return 6; // FLAT_SCRATCH, XNACK, VCC (in that order).

 unsigned GCNSubtarget::getBaseMaxNumVGPRs(
     const Function &F, std::pair<unsigned, unsigned> WavesPerEU) const {
   // Compute maximum number of VGPRs function can use using default/requested
   // minimum number of waves per execution unit.
   unsigned MaxNumVGPRs = getMaxNumVGPRs(WavesPerEU.first);
 
   // Check if maximum number of VGPRs was explicitly requested using
   // "amdgpu-num-vgpr" attribute.
+  unsigned Requested = getFunctionNVGPRAttr(F, WavesPerEU);
+
+  if (Requested)
+    MaxNumVGPRs = Requested;
+
+  return MaxNumVGPRs;
+}
+
+unsigned llvm::GCNSubtarget::getFunctionNVGPRAttr(
+    const Function &F, std::pair<unsigned, unsigned> WavesPerEU) const {
+
+  unsigned Requested = 0;
   if (F.hasFnAttribute("amdgpu-num-vgpr")) {
-    unsigned Requested =
-        F.getFnAttributeAsParsedInteger("amdgpu-num-vgpr", MaxNumVGPRs);
+    Requested =
+        F.getFnAttributeAsParsedInteger("amdgpu-num-vgpr", WavesPerEU.first);
 
     if (hasGFX90AInsts())
       Requested *= 2;
 
     // Make sure requested value is compatible with values implied by
     // default/requested minimum/maximum number of waves per execution unit.
     if (Requested && Requested > getMaxNumVGPRs(WavesPerEU.first))
       Requested = 0;
-    if (WavesPerEU.second &&
-        Requested && Requested < getMinNumVGPRs(WavesPerEU.second))
+    if (WavesPerEU.second && Requested &&
+        Requested < getMinNumVGPRs(WavesPerEU.second))
       Requested = 0;
-
-    if (Requested)
-      MaxNumVGPRs = Requested;
   }
-
-  return MaxNumVGPRs;
+  return Requested;
 }
 
 unsigned GCNSubtarget::getMaxNumVGPRs(const Function &F) const {
   return getBaseMaxNumVGPRs(F, getWavesPerEU(F));
 }
 
 unsigned GCNSubtarget::getMaxNumVGPRs(const MachineFunction &MF) const {
   const Function &F = MF.getFunction();

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

 //===-- GCNSchedStrategy.cpp - GCN Scheduler Strategy ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (PressureAfter.getSGPRNum() <= S.SGPRCriticalLimit &&
 
     // Early out if we have achieve the occupancy target.
     LLVM_DEBUG(dbgs() << "Pressure in desired limits, done.\n");
     return;
   }
 
   unsigned TargetOccupancy =
       std::min(S.getTargetOccupancy(), ST.getOccupancyWithLocalMemSize(MF));
-  unsigned WavesAfter =
-      std::min(TargetOccupancy, PressureAfter.getOccupancy(ST));
-  unsigned WavesBefore =
-      std::min(TargetOccupancy, PressureBefore.getOccupancy(ST));
+
+  unsigned WavesAfter = std::min(
+      TargetOccupancy,
+      std::min(ST.getOccupancyWithNumSGPRs(PressureAfter.getSGPRNum()),
+               ST.getOccupancyWithNumVGPRsForFunction(
+                   MF, PressureAfter.getVGPRNum(ST.hasGFX90AInsts()))));
+  unsigned WavesBefore = std::min(
+      TargetOccupancy,
+      std::min(ST.getOccupancyWithNumSGPRs(PressureBefore.getSGPRNum()),
+               ST.getOccupancyWithNumVGPRsForFunction(
+                   MF, PressureBefore.getVGPRNum(ST.hasGFX90AInsts()))));
+
   LLVM_DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore
                     << ", after " << WavesAfter << ".\n");
 
   // We may not be able to keep the current target occupancy because of the just
   // scheduled region. We might still be able to revert scheduling if the
   // occupancy before was higher, or if the current schedule has register
   // pressure higher than the excess limits which could lead to more spilling.
   unsigned NewOccupancy = std::max(WavesAfter, WavesBefore);

llvm/lib/Target/AMDGPU/GCNSubtarget.h

 //=====-- GCNSubtarget.h - Define GCN Subtarget for AMDGPU ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //==-----------------------------------------------------------------------===//
 //
public:
   /// Return the maximum number of waves per SIMD for kernels using \p SGPRs
   /// SGPRs
   unsigned getOccupancyWithNumSGPRs(unsigned SGPRs) const;
 
   /// Return the maximum number of waves per SIMD for kernels using \p VGPRs
   /// VGPRs
   unsigned getOccupancyWithNumVGPRs(unsigned VGPRs) const;
 
+  unsigned getOccupancyWithNumVGPRsForFunction(const MachineFunction &MF,
+                                               unsigned VGPRs) const;
+
   /// Return occupancy for the given function. Used LDS and a number of
   /// registers if provided.
   /// Note, occupancy can be affected by the scratch allocation as well, but
   /// we do not have enough information to compute it.
   unsigned computeOccupancy(const Function &F, unsigned LDSSize = 0,
                             unsigned NumSGPRs = 0, unsigned NumVGPRs = 0) const;
 
   /// \returns true if the flat_scratch register should be initialized with the

   unsigned getMaxNumVGPRs(unsigned WavesPerEU) const {
     return AMDGPU::IsaInfo::getMaxNumVGPRs(this, WavesPerEU);
   }
 
   /// \returns max num VGPRs. This is the common utility function
   /// called by MachineFunction and Function variants of getMaxNumVGPRs.
   unsigned getBaseMaxNumVGPRs(const Function &F,
                               std::pair<unsigned, unsigned> WavesPerEU) const;
+  unsigned getFunctionNVGPRAttr(const Function &F,
+                              std::pair<unsigned, unsigned> WavesPerEU) const;
   /// \returns Maximum number of VGPRs that meets number of waves per execution
   /// unit requirement for function \p F, or number of VGPRs explicitly
   /// requested using "amdgpu-num-vgpr" attribute attached to function \p F.
   ///
   /// \returns Value that meets number of waves per execution unit requirement
   /// if explicitly requested value cannot be converted to integer, violates
   /// subtarget's specifications, or does not meet number of waves per execution
   /// unit requirement.

llvm/test/CodeGen/AMDGPU/agpr-remat.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -march=amdgcn -mcpu=gfx908 -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefixes=GFX908 %s
 
 ; Make sure there are no v_accvgpr_read_b32 copying back and forth
 ; between AGPR and VGPR.
 define amdgpu_kernel void @remat_constant_voids_spill(ptr addrspace(1) %p) #1 {
 ; GFX908-LABEL: remat_constant_voids_spill:
 ; GFX908:       ; %bb.0:
-; GFX908-NEXT:    v_accvgpr_write_b32 a1, 1
-; GFX908-NEXT:    v_accvgpr_write_b32 a5, 6
-; GFX908-NEXT:    v_accvgpr_write_b32 a6, 7
-; GFX908-NEXT:    v_accvgpr_write_b32 a7, 8
-; GFX908-NEXT:    v_accvgpr_write_b32 a0, 9
-; GFX908-NEXT:    v_accvgpr_write_b32 a2, 2
-; GFX908-NEXT:    v_accvgpr_write_b32 a3, 3
-; GFX908-NEXT:    v_accvgpr_write_b32 a4, 4
+; GFX908-NEXT:    v_accvgpr_write_b32 a0, 1
+; GFX908-NEXT:    v_accvgpr_write_b32 a1, 2
+; GFX908-NEXT:    v_accvgpr_write_b32 a2, 3
+; GFX908-NEXT:    v_accvgpr_write_b32 a3, 4
 ; GFX908-NEXT:    ;;#ASMSTART
 ; GFX908-NEXT:    ;;#ASMEND
-; GFX908-NEXT:    v_accvgpr_write_b32 a1, 5
+; GFX908-NEXT:    v_accvgpr_write_b32 a0, 5
+; GFX908-NEXT:    v_accvgpr_write_b32 a1, 6
+; GFX908-NEXT:    v_accvgpr_write_b32 a2, 7
+; GFX908-NEXT:    v_accvgpr_write_b32 a3, 8
+; GFX908-NEXT:    v_accvgpr_write_b32 a4, 9
 ; GFX908-NEXT:    ;;#ASMSTART
 ; GFX908-NEXT:    ;;#ASMEND
 ; GFX908-NEXT:    s_endpgm
   call void asm sideeffect "", "a,a,a,a"(i32 1, i32 2, i32 3, i32 4)
   call void asm sideeffect "", "a,a,a,a,a"(i32 5, i32 6, i32 7, i32 8, i32 9)
   ret void
 }
 

llvm/test/CodeGen/AMDGPU/preserve-wwm-copy-dst-reg.ll

 ; GFX906-NEXT:    ;;#ASMEND
 ; GFX906-NEXT:    v_writelane_b32 v40, s20, 21
 ; GFX906-NEXT:    s_waitcnt lgkmcnt(0)
 ; GFX906-NEXT:    v_writelane_b32 v40, s10, 22
 ; GFX906-NEXT:    v_writelane_b32 v40, s11, 23
 ; GFX906-NEXT:    s_or_saveexec_b64 s[34:35], -1
 ; GFX906-NEXT:    s_mov_b64 exec, s[34:35]
 ; GFX906-NEXT:    v_readlane_b32 s16, v40, 22
+; GFX906-NEXT:    v_mov_b32_e32 v31, v32
 ; GFX906-NEXT:    s_mov_b32 s12, s24
 ; GFX906-NEXT:    s_mov_b32 s13, s23
 ; GFX906-NEXT:    s_mov_b32 s14, s22
-; GFX906-NEXT:    v_mov_b32_e32 v31, v32
 ; GFX906-NEXT:    s_mov_b32 s15, s21
 ; GFX906-NEXT:    s_mov_b64 s[10:11], s[26:27]
 ; GFX906-NEXT:    v_readlane_b32 s17, v40, 23
 ; GFX906-NEXT:    buffer_store_dword v32, off, s[0:3], s33 offset:4 ; 4-byte Folded Spill
 ; GFX906-NEXT:    s_swappc_b64 s[30:31], s[16:17]
 ; GFX906-NEXT:    s_or_saveexec_b64 s[34:35], -1
 ; GFX906-NEXT:    s_mov_b64 exec, s[34:35]
 ; GFX906-NEXT:    v_readlane_b32 s11, v40, 12

 ; GFX908-NEXT:    ;;#ASMEND
 ; GFX908-NEXT:    v_writelane_b32 v40, s20, 21
 ; GFX908-NEXT:    s_waitcnt lgkmcnt(0)
 ; GFX908-NEXT:    v_writelane_b32 v40, s10, 22
 ; GFX908-NEXT:    v_writelane_b32 v40, s11, 23
 ; GFX908-NEXT:    s_or_saveexec_b64 s[34:35], -1
 ; GFX908-NEXT:    s_mov_b64 exec, s[34:35]
 ; GFX908-NEXT:    v_readlane_b32 s16, v40, 22
+; GFX908-NEXT:    v_mov_b32_e32 v31, v32
 ; GFX908-NEXT:    s_mov_b32 s12, s24
 ; GFX908-NEXT:    s_mov_b32 s13, s23
 ; GFX908-NEXT:    s_mov_b32 s14, s22
-; GFX908-NEXT:    v_mov_b32_e32 v31, v32
 ; GFX908-NEXT:    s_mov_b32 s15, s21
 ; GFX908-NEXT:    s_mov_b64 s[10:11], s[26:27]
 ; GFX908-NEXT:    v_readlane_b32 s17, v40, 23
 ; GFX908-NEXT:    buffer_store_dword v32, off, s[0:3], s33 offset:8 ; 4-byte Folded Spill
 ; GFX908-NEXT:    s_swappc_b64 s[30:31], s[16:17]
 ; GFX908-NEXT:    s_or_saveexec_b64 s[34:35], -1
 ; GFX908-NEXT:    s_mov_b64 exec, s[34:35]
 ; GFX908-NEXT:    v_readlane_b32 s11, v40, 12