[AMDGPU][RFC] Improve sgpr function arguments

[sebastian-ne]

Summary

The SGPR layout on functions calls currently looks like this:

s[0:3] SRD	arguments...	s[30:31] return address	s32 stack pointer	s33 frame pointer	s34 base pointer

The return address and stack pointer occupy multiple 4-aligned blocks
of SGPRs.
Large scalar memory reads require a 4-aligned block of SGPRs, so if less
of them are available, register allocation becomes more difficult.

The stack resource descriptor occupies SGPR0-3. If we want to pass
user-data SGPRS to a function, the SGPRs need to be moved from s[0:...]
to s[4:...] before the call.
This is also the case when flat scratch is used instead of the SRD, even
if s[0:4] is unused then, because the same call convention is used.

To improve this, I propose the following layout:

arguments...	s[24:27] SRD	s[28:29] return address	...	s35 stack pointer	...	s40 frame pointer	s41 base pointer

To free s[0:3] for arguments, the SRD is moved to s[24:27]. This has
the effect, that all of s[0:23] can be used for arguments.

The base pointer is not used in the general case, so it is moved to
s41.

[arsenm]

We have another proposal we were working on to rearrange these a bit differently. We need to account for a few more inputs in the layout

[madhur13490]

In D102177#2748278, @arsenm wrote:

We have another proposal we were working on to rearrange these a bit differently. We need to account for a few more inputs in the layout

As long as this remains in GFX land, we should be fine with it because our new proposal is for compute only (as of now).

[sebastian-ne]

In D102177#2750094, @madhur13490 wrote:

In D102177#2748278, @arsenm wrote:

We have another proposal we were working on to rearrange these a bit differently. We need to account for a few more inputs in the layout

As long as this remains in GFX land, we should be fine with it because our new proposal is for compute only (as of now).

I would like to keep the same calling convention in compute and graphics. At least regarding the stack pointer and others, because I don’t see a compelling reason to diverge even more. Actually, I’d like it if they were more common than they are now, because we implement some things twice at the moment.
The compute proposal should work just fine; if we move the stack and frame pointer, we end up with the same benefits as in this patch. I commented on the internal proposal for this (I hope I found the right one?).

[madhur13490]

In D102177#2750256, @sebastian-ne wrote:

In D102177#2750094, @madhur13490 wrote:

In D102177#2748278, @arsenm wrote:

We have another proposal we were working on to rearrange these a bit differently. We need to account for a few more inputs in the layout

As long as this remains in GFX land, we should be fine with it because our new proposal is for compute only (as of now).

I would like to keep the same calling convention in compute and graphics. At least regarding the stack pointer and others, because I don’t see a compelling reason to diverge even more. Actually, I’d like it if they were more common than they are now, because we implement some things twice at the moment.
The compute proposal should work just fine; if we move the stack and frame pointer, we end up with the same benefits as in this patch. I commented on the internal proposal for this (I hope I found the right one?).

Well, then you're saying unification of both ABIs and it is not discussed thoroughly internally. The layout needs to be documented and get reviewed internally before we can proceed with this patch.

[t-tye]

In D102177#2750461, @madhur13490 wrote:

In D102177#2750256, @sebastian-ne wrote:

In D102177#2750094, @madhur13490 wrote:

In D102177#2748278, @arsenm wrote:

We have another proposal we were working on to rearrange these a bit differently. We need to account for a few more inputs in the layout

As long as this remains in GFX land, we should be fine with it because our new proposal is for compute only (as of now).

I would like to keep the same calling convention in compute and graphics. At least regarding the stack pointer and others, because I don’t see a compelling reason to diverge even more. Actually, I’d like it if they were more common than they are now, because we implement some things twice at the moment.
The compute proposal should work just fine; if we move the stack and frame pointer, we end up with the same benefits as in this patch. I commented on the internal proposal for this (I hope I found the right one?).

Well, then you're saying unification of both ABIs and it is not discussed thoroughly internally. The layout needs to be documented and get reviewed internally before we can proceed with this patch.

I think the reason that compute was considering putting FP and BP at the high registers is because they can often be optimized away. So putting them there allows them to be contiguous with other callee user registers rather then leaving an unused hole that is harder for the register allocator to use. Is that reasonable? How does that work for gfx?

I am all for trying to unify the call convention for compute and gfx if possible:-)

[sebastian-ne]

Move registers to be more in line with other plans.

[arsenm]

Current proposal is different values

llvm/docs/AMDGPUUsage.rst

 runtime. It is used, together with Scratch Wavefront Offset as an offset, to
 access the private memory space using a segment address. See
 :ref:`amdgpu-amdhsa-initial-kernel-execution-state`.
 
 The scratch V# is a four-aligned SGPR and always selected for the kernel as
 follows:
 
   - If it is known during instruction selection that there is stack usage,
-    SGPR0-3 is reserved for use as the scratch V#.  Stack usage is assumed if
+    SGPR36-39 is reserved for use as the scratch V#.  Stack usage is assumed if
     optimizations are disabled (``-O0``), if stack objects already exist (for
     locals, etc.), or if there are any function calls.
 
   - Otherwise, four high numbered SGPRs beginning at a four-aligned SGPR index
     are reserved for the tentative scratch V#. These will be used if it is
     determined that spilling is needed.
 
     - If no use is made of the tentative scratch V#, then it is unreserved,

 outer kernel function.
 
 If a kernel has function calls then scratch is always allocated and used for
 the call stack which grows from low address to high address using the swizzled
 scratch address space.
 
 On entry to a function:
 
-1.  SGPR0-3 contain a V# with the following properties (see
+1.  SGPR36-39 contain a V# with the following properties (see
     :ref:`amdgpu-amdhsa-kernel-prolog-private-segment-buffer`):
 
     * Base address pointing to the beginning of the wavefront scratch backing
       memory.
     * Swizzled with dword element size and stride of wavefront size elements.
 
 2.  The FLAT_SCRATCH register pair is setup. See
     :ref:`amdgpu-amdhsa-kernel-prolog-flat-scratch`.
 3.  GFX6-GFX8: M0 register set to the size of LDS in bytes. See
     :ref:`amdgpu-amdhsa-kernel-prolog-m0`.
 4.  The EXEC register is set to the lanes active on entry to the function.
 5.  MODE register: *TBD*
 6.  VGPR0-31 and SGPR4-29 are used to pass function input arguments as described
     below.
-7.  SGPR30-31 return address (RA). The code address that the function must
+7.  SGPR34-35 return address (RA). The code address that the function must
     return to when it completes. The value is undefined if the function is *no
     return*.
 8.  SGPR32 is used for the stack pointer (SP). It is an unswizzled scratch
     offset relative to the beginning of the wavefront scratch backing memory.
 
     The unswizzled SP can be used with buffer instructions as an unswizzled SGPR
-    offset with the scratch V# in SGPR0-3 to access the stack in a swizzled
+    offset with the scratch V# in SGPR36-39 to access the stack in a swizzled
     manner.
 
     The unswizzled SP value can be converted into the swizzled SP value by:
 
       | swizzled SP = unswizzled SP / wavefront size
 
     This may be used to obtain the private address space address of stack
     objects and to convert this address to a flat address by adding the flat
* VGPR arguments are assigned to consecutive VGPRs starting at VGPR0 up to
   arguments are allocated on the stack in order on naturally aligned
   addresses.
 
   .. TODO::
 
     How are overly aligned structures allocated on the stack?
 
 * SGPR arguments are assigned to consecutive SGPRs starting at SGPR0 up to
-  SGPR29.
+  SGPR31.
 
   If there are more arguments than will fit in these registers, the remaining
   arguments are allocated on the stack in order on naturally aligned
   addresses.
 
 Note that decomposed struct type arguments may have some fields passed in
 registers and some in memory.
 
 .. TODO::
 
   So, a struct which can pass some fields as decomposed register arguments, will
   pass the rest as decomposed stack elements? But an argument that will not start
   in registers will not be decomposed and will be passed as a non-decomposed
   stack value?
 
 The following is not part of the AMDGPU function calling convention but
 describes how the AMDGPU implements function calls:
 
 1.  SGPR33 is used as a frame pointer (FP) if necessary. Like the SP it is an
     unswizzled scratch address. It is only needed if runtime sized ``alloca``
     are used, or for the reasons defined in ``SIFrameLowering``.
-2.  Runtime stack alignment is supported. SGPR34 is used as a base pointer (BP)
+2.  Runtime stack alignment is supported. SGPR40 is used as a base pointer (BP)
     to access the incoming stack arguments in the function. The BP is needed
     only when the function requires the runtime stack alignment.
 
 3.  Allocating SGPR arguments on the stack are not supported.
 
 4.  No CFI is currently generated. See
     :ref:`amdgpu-dwarf-call-frame-information`.
 

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

bool AMDGPUCallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
 
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
 
   if (!ST.enableFlatScratch()) {
     // Insert copies for the SRD. In the HSA case, this should be an identity
     // copy.
     auto ScratchRSrcReg = MIRBuilder.buildCopy(LLT::vector(4, 32),
                                                MFI->getScratchRSrcReg());
-    MIRBuilder.buildCopy(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3, ScratchRSrcReg);
-    MIB.addReg(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3, RegState::Implicit);
+    MIRBuilder.buildCopy(TRI->getScratchRSrcReg(), ScratchRSrcReg);
+    MIB.addReg(TRI->getScratchRSrcReg(), RegState::Implicit);
   }
 
   for (std::pair<MCRegister, Register> ArgReg : ImplicitArgRegs) {
     MIRBuilder.buildCopy((Register)ArgReg.first, ArgReg.second);
     MIB.addReg(ArgReg.first, RegState::Implicit);
   }
 
   // Get a count of how many bytes are to be pushed on the stack.

llvm/lib/Target/AMDGPU/AMDGPUCallingConv.td

 
 // Inversion of CCIfInReg
 class CCIfNotInReg<CCAction A> : CCIf<"!ArgFlags.isInReg()", A> {}
 class CCIfExtend<CCAction A>
   : CCIf<"ArgFlags.isSExt() || ArgFlags.isZExt()", A>;
 
 // Calling convention for SI
 def CC_SI_Gfx : CallingConv<[
-  // 0-3 are reserved for the stack buffer descriptor
-  // 30-31 are reserved for the return address
-  // 32 is reserved for the stack pointer
+  // SGPR32 onwards is reserved for the stack pointer, return address, etc.
   CCIfInReg<CCIfType<[f32, i32, f16, i16, v2i16, v2f16] , CCAssignToReg<[
-    SGPR4, SGPR5, SGPR6, SGPR7,
+    SGPR0, SGPR1, SGPR2, SGPR3, SGPR4, SGPR5, SGPR6, SGPR7,
     SGPR8, SGPR9, SGPR10, SGPR11, SGPR12, SGPR13, SGPR14, SGPR15,
     SGPR16, SGPR17, SGPR18, SGPR19, SGPR20, SGPR21, SGPR22, SGPR23,
-    SGPR24, SGPR25, SGPR26, SGPR27, SGPR28, SGPR29,
+    SGPR24, SGPR25, SGPR26, SGPR27, SGPR28, SGPR29, SGPR30, SGPR31,
   ]>>>,
 
   CCIfNotInReg<CCIfType<[f32, i32, f16, i16, v2i16, v2f16] , CCAssignToReg<[
     VGPR0, VGPR1, VGPR2, VGPR3, VGPR4, VGPR5, VGPR6, VGPR7,
     VGPR8, VGPR9, VGPR10, VGPR11, VGPR12, VGPR13, VGPR14, VGPR15,
     VGPR16, VGPR17, VGPR18, VGPR19, VGPR20, VGPR21, VGPR22, VGPR23,
     VGPR24, VGPR25, VGPR26, VGPR27, VGPR28, VGPR29, VGPR30, VGPR31
   ]>>>,
 
   CCIfType<[i32, f32, v2i16, v2f16, i16, f16, i1], CCAssignToStack<4, 4>>
 ]>;
 
 def RetCC_SI_Gfx : CallingConv<[
   CCIfType<[i1], CCPromoteToType<i32>>,
   CCIfType<[i1, i16], CCIfExtend<CCPromoteToType<i32>>>,
 
-  // 0-3 are reserved for the stack buffer descriptor
-  // 32 is reserved for the stack pointer
+  // SGPR32 onwards is reserved for the stack pointer, return address, etc.
   CCIfInReg<CCIfType<[f32, i32, f16, i16, v2i16, v2f16] , CCAssignToReg<[
-    SGPR4, SGPR5, SGPR6, SGPR7,
+    SGPR0, SGPR1, SGPR2, SGPR3, SGPR4, SGPR5, SGPR6, SGPR7,
     SGPR8, SGPR9, SGPR10, SGPR11, SGPR12, SGPR13, SGPR14, SGPR15,
     SGPR16, SGPR17, SGPR18, SGPR19, SGPR20, SGPR21, SGPR22, SGPR23,
     SGPR24, SGPR25, SGPR26, SGPR27, SGPR28, SGPR29, SGPR30, SGPR31,
-    SGPR33, SGPR34, SGPR35, SGPR36, SGPR37, SGPR38, SGPR39,
-    SGPR40, SGPR41, SGPR42, SGPR43
   ]>>>,
 
   CCIfNotInReg<CCIfType<[f32, i32, f16, i16, v2i16, v2f16] , CCAssignToReg<[
     VGPR0, VGPR1, VGPR2, VGPR3, VGPR4, VGPR5, VGPR6, VGPR7,
     VGPR8, VGPR9, VGPR10, VGPR11, VGPR12, VGPR13, VGPR14, VGPR15,
     VGPR16, VGPR17, VGPR18, VGPR19, VGPR20, VGPR21, VGPR22, VGPR23,
     VGPR24, VGPR25, VGPR26, VGPR27, VGPR28, VGPR29, VGPR30, VGPR31,
     VGPR32, VGPR33, VGPR34, VGPR35, VGPR36, VGPR37, VGPR38, VGPR39,

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

if (!IsSibCall) {
     Chain = DAG.getCALLSEQ_START(Chain, 0, 0, DL);
 
     if (!Subtarget->enableFlatScratch()) {
       SmallVector<SDValue, 4> CopyFromChains;
 
       // In the HSA case, this should be an identity copy.
       SDValue ScratchRSrcReg
         = DAG.getCopyFromReg(Chain, DL, Info->getScratchRSrcReg(), MVT::v4i32);
-      RegsToPass.emplace_back(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3, ScratchRSrcReg);
+      const SIRegisterInfo *TRI = getSubtarget()->getRegisterInfo();
+      RegsToPass.emplace_back(TRI->getScratchRSrcReg(), ScratchRSrcReg);
       CopyFromChains.push_back(ScratchRSrcReg.getValue(1));
       Chain = DAG.getTokenFactor(DL, CopyFromChains);
     }
   }
 
   MVT PtrVT = MVT::i32;
 
   // Walk the register/memloc assignments, inserting copies/loads.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

if (!isEntryFunction()) {
 
     // TODO: Pick a high register, and shift down, similar to a kernel.
     FrameOffsetReg = AMDGPU::SGPR33;
     StackPtrOffsetReg = AMDGPU::SGPR32;
 
     if (!ST.enableFlatScratch()) {
       // Non-entry functions have no special inputs for now, other registers
       // required for scratch access.
-      ScratchRSrcReg = AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3;
+      ScratchRSrcReg = AMDGPU::SGPR36_SGPR37_SGPR38_SGPR39;
 
       ArgInfo.PrivateSegmentBuffer =
         ArgDescriptor::createRegister(ScratchRSrcReg);
     }
 
     if (F.hasFnAttribute("amdgpu-implicitarg-ptr"))
       ImplicitArgPtr = true;
   } else {

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

public:
 
   unsigned getRegPressureSetLimit(const MachineFunction &MF,
                                   unsigned Idx) const override;
 
   const int *getRegUnitPressureSets(unsigned RegUnit) const override;
 
   MCRegister getReturnAddressReg(const MachineFunction &MF) const;
 
+  Register getScratchRSrcReg() const;
+
   const TargetRegisterClass *
   getRegClassForSizeOnBank(unsigned Size,
                            const RegisterBank &Bank,
                            const MachineRegisterInfo &MRI) const;
 
   const TargetRegisterClass *
   getRegClassForTypeOnBank(LLT Ty,
                            const RegisterBank &Bank,

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

 
 bool SIRegisterInfo::hasBasePointer(const MachineFunction &MF) const {
   // When we need stack realignment, we can't reference off of the
   // stack pointer, so we reserve a base pointer.
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   return MFI.getNumFixedObjects() && shouldRealignStack(MF);
 }
 
-Register SIRegisterInfo::getBaseRegister() const { return AMDGPU::SGPR34; }
+Register SIRegisterInfo::getBaseRegister() const { return AMDGPU::SGPR40; }
 
 const uint32_t *SIRegisterInfo::getAllVGPRRegMask() const {
   return CSR_AMDGPU_AllVGPRs_RegMask;
 }
 
 const uint32_t *SIRegisterInfo::getAllAGPRRegMask() const {
   return CSR_AMDGPU_AllAGPRs_RegMask;
 }
const int *SIRegisterInfo::getRegUnitPressureSets(unsigned RegUnit) const {
   if (RegPressureIgnoredUnits[RegUnit])
     return Empty;
 
   return AMDGPUGenRegisterInfo::getRegUnitPressureSets(RegUnit);
 }
 
 MCRegister SIRegisterInfo::getReturnAddressReg(const MachineFunction &MF) const {
   // Not a callee saved register.
-  return AMDGPU::SGPR30_SGPR31;
+  return AMDGPU::SGPR34_SGPR35;
+}
+
+Register SIRegisterInfo::getScratchRSrcReg() const {
+  return AMDGPU::SGPR36_SGPR37_SGPR38_SGPR39;
 }
 
 const TargetRegisterClass *
 SIRegisterInfo::getRegClassForSizeOnBank(unsigned Size,
                                          const RegisterBank &RB,
                                          const MachineRegisterInfo &MRI) const {
   switch (RB.getID()) {
   case AMDGPU::VGPRRegBankID: