Fix align for hidden args.

Author: kerbowa

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2510,8 +2510,7 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
     const SmallVectorImpl<ISD::InputArg> &Ins, MachineFunction &MF,
     const SIRegisterInfo &TRI, SIMachineFunctionInfo &Info) const {
   Function &F = MF.getFunction();
-  unsigned LastExplicitArgOffset =
-      MF.getSubtarget<GCNSubtarget>().getExplicitKernelArgOffset();
+  unsigned LastExplicitArgOffset = Subtarget->getExplicitKernelArgOffset();
   GCNUserSGPRUsageInfo &SGPRInfo = Info.getUserSGPRInfo();
   bool InPreloadSequence = true;
   unsigned InIdx = 0;
@@ -2539,14 +2538,15 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
           alignTo(ArgLoc.getLocVT().getFixedSizeInBits(), 32) / 32;
 
       // Add padding SPGR to fix alignment for hidden arguments.
-      if (!AlignedForImplictArgs && Arg.hasAttribute("amdgpu-work-group-id")) {
+      if (!AlignedForImplictArgs &&
+          Arg.hasAttribute("amdgpu-hidden-argument")) {
         unsigned OffsetBefore = LastExplicitArgOffset;
         LastExplicitArgOffset = alignTo(
             LastExplicitArgOffset, Subtarget->getAlignmentForImplicitArgPtr());
         if (OffsetBefore != LastExplicitArgOffset) {
           unsigned PaddingSGPRs =
               alignTo(LastExplicitArgOffset - OffsetBefore, 4) / 4;
-          Info.allocateUserSGPRs(PaddingSGPRs);
+          Info.allocateUserSGPRs(*Subtarget, PaddingSGPRs);
           ArgOffset += PaddingSGPRs * 4;
         }
         AlignedForImplictArgs = true;

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -277,11 +277,13 @@ SmallVectorImpl<MCRegister> *SIMachineFunctionInfo::addPreloadedKernArg(
   return &ArgInfo.PreloadKernArgs[KernArgIdx].Regs;
 }
 
-bool SIMachineFunctionInfo::allocateUserSGPRs(unsigned Number) {
-  if (Number <= getNumUserSGPRs())
+bool SIMachineFunctionInfo::allocateUserSGPRs(const GCNSubtarget &ST,
+                                              unsigned Number) {
+  unsigned NewUserSGPRs = NumUserSGPRs + Number;
+  if (NewUserSGPRs > ST.getMaxNumUserSGPRs())
     return false;
 
-  NumUserSGPRs = Number;
+  NumUserSGPRs = NewUserSGPRs;
   return true;
 }
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

@@ -781,7 +781,7 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
                       int PaddingSGPRs);
 
   /// Reserve up to \p Number of user SGPRs.
-  bool allocateUserSGPRs(unsigned Number);
+  bool allocateUserSGPRs(const GCNSubtarget &ST, unsigned Number);
 
   /// Increment user SGPRs used for padding the argument list only.
   Register addReservedUserSGPR() {

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll

@@ -32,7 +32,7 @@ define amdgpu_kernel void @preload_unused_arg_block_count_x(ptr addrspace(1) inr
 ; GFX940-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
 ; GFX940-NEXT:  ; %bb.0:
 ; GFX940-NEXT:    v_mov_b32_e32 v0, 0
-; GFX940-NEXT:    v_mov_b32_e32 v1, s5
+; GFX940-NEXT:    v_mov_b32_e32 v1, s6
 ; GFX940-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
 ; GFX940-NEXT:    s_endpgm
 ;
@@ -41,7 +41,7 @@ define amdgpu_kernel void @preload_unused_arg_block_count_x(ptr addrspace(1) inr
 ; GFX90a-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
 ; GFX90a-NEXT:  ; %bb.0:
 ; GFX90a-NEXT:    v_mov_b32_e32 v0, 0
-; GFX90a-NEXT:    v_mov_b32_e32 v1, s9
+; GFX90a-NEXT:    v_mov_b32_e32 v1, s10
 ; GFX90a-NEXT:    global_store_dword v0, v1, s[6:7]
 ; GFX90a-NEXT:    s_endpgm
   %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
@@ -276,7 +276,7 @@ define amdgpu_kernel void @preload_block_count_x_imparg_align_ptr_i8(ptr addrspa
 ; GFX940-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
 ; GFX940-NEXT:  ; %bb.0:
 ; GFX940-NEXT:    s_and_b32 s0, s4, 0xff
-; GFX940-NEXT:    s_add_i32 s0, s5, s0
+; GFX940-NEXT:    s_add_i32 s0, s6, s0
 ; GFX940-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX940-NEXT:    v_mov_b32_e32 v1, s0
 ; GFX940-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
@@ -287,7 +287,7 @@ define amdgpu_kernel void @preload_block_count_x_imparg_align_ptr_i8(ptr addrspa
 ; GFX90a-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
 ; GFX90a-NEXT:  ; %bb.0:
 ; GFX90a-NEXT:    s_and_b32 s0, s8, 0xff
-; GFX90a-NEXT:    s_add_i32 s0, s9, s0
+; GFX90a-NEXT:    s_add_i32 s0, s10, s0
 ; GFX90a-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX90a-NEXT:    v_mov_b32_e32 v1, s0
 ; GFX90a-NEXT:    global_store_dword v0, v1, s[6:7]