[AMDGPU] Fix hidden kernarg preload count inconsistency

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

@@ -520,6 +520,12 @@ bool AMDGPUCallLowering::lowerFormalArgumentsKernel(
 
   // TODO: Align down to dword alignment and extract bits for extending loads.
   for (auto &Arg : F.args()) {
+    // TODO: Add support for kernarg preload.
+    if (Arg.hasAttribute("amdgpu-hidden-argument")) {
+      LLVM_DEBUG(dbgs() << "Preloading hidden arguments is not supported\n");
+      return false;
+    }
+
     const bool IsByRef = Arg.hasByRefAttr();
     Type *ArgTy = IsByRef ? Arg.getParamByRefType() : Arg.getType();
     unsigned AllocSize = DL.getTypeAllocSize(ArgTy);

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

@@ -144,20 +144,20 @@ class PreloadKernelArgInfo {
   // Returns the maximum number of user SGPRs that we have available to preload
   // arguments.
   void setInitialFreeUserSGPRsCount() {
-    const unsigned MaxUserSGPRs = ST.getMaxNumUserSGPRs();
     GCNUserSGPRUsageInfo UserSGPRInfo(F, ST);
-
-    NumFreeUserSGPRs = MaxUserSGPRs - UserSGPRInfo.getNumUsedUserSGPRs();
+    NumFreeUserSGPRs = UserSGPRInfo.getNumFreeUserSGPRs();
   }
 
   bool tryAllocPreloadSGPRs(unsigned AllocSize, uint64_t ArgOffset,
                             uint64_t LastExplicitArgOffset) {
     //  Check if this argument may be loaded into the same register as the
     //  previous argument.
-    if (!isAligned(Align(4), ArgOffset) && AllocSize < 4)
+    if (ArgOffset - LastExplicitArgOffset < 4 &&
+        !isAligned(Align(4), ArgOffset))
       return true;
 
     // Pad SGPRs for kernarg alignment.
+    ArgOffset = alignDown(ArgOffset, 4);
     unsigned Padding = ArgOffset - LastExplicitArgOffset;
     unsigned PaddingSGPRs = alignTo(Padding, 4) / 4;
     unsigned NumPreloadSGPRs = alignTo(AllocSize, 4) / 4;
@@ -170,6 +170,7 @@ class PreloadKernelArgInfo {
 
   // Try to allocate SGPRs to preload implicit kernel arguments.
   void tryAllocImplicitArgPreloadSGPRs(uint64_t ImplicitArgsBaseOffset,
+                                       uint64_t LastExplicitArgOffset,
                                        IRBuilder<> &Builder) {
     Function *ImplicitArgPtr = Intrinsic::getDeclarationIfExists(
         F.getParent(), Intrinsic::amdgcn_implicitarg_ptr);
@@ -215,7 +216,6 @@ class PreloadKernelArgInfo {
     // argument can actually be preloaded.
     std::sort(ImplicitArgLoads.begin(), ImplicitArgLoads.end(), less_second());
 
-    uint64_t LastExplicitArgOffset = ImplicitArgsBaseOffset;
     // If we fail to preload any implicit argument we know we don't have SGPRs
     // to preload any subsequent ones with larger offsets. Find the first
     // argument that we cannot preload.
@@ -229,7 +229,8 @@ class PreloadKernelArgInfo {
                                     LastExplicitArgOffset))
             return true;
 
-          LastExplicitArgOffset = LoadOffset + LoadSize;
+          LastExplicitArgOffset =
+              ImplicitArgsBaseOffset + LoadOffset + LoadSize;
           return false;
         });
 
@@ -486,7 +487,7 @@ static bool lowerKernelArguments(Function &F, const TargetMachine &TM) {
         alignTo(ExplicitArgOffset, ST.getAlignmentForImplicitArgPtr()) +
         BaseOffset;
     PreloadInfo.tryAllocImplicitArgPreloadSGPRs(ImplicitArgsBaseOffset,
-                                                Builder);
+                                                ExplicitArgOffset, Builder);
   }
 
   return true;

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2545,8 +2545,7 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
       unsigned Padding = ArgOffset - LastExplicitArgOffset;
       unsigned PaddingSGPRs = alignTo(Padding, 4) / 4;
       // Check for free user SGPRs for preloading.
-      if (PaddingSGPRs + NumAllocSGPRs + 1 /*Synthetic SGPRs*/ >
-          SGPRInfo.getNumFreeUserSGPRs()) {
+      if (PaddingSGPRs + NumAllocSGPRs > SGPRInfo.getNumFreeUserSGPRs()) {
         InPreloadSequence = false;
         break;
       }
@@ -3024,6 +3023,20 @@ SDValue SITargetLowering::LowerFormalArguments(
           NewArg = DAG.getMergeValues({NewArg, Chain}, DL);
         }
       } else {
+        // Hidden arguments that are in the kernel signature must be preloaded
+        // to user SGPRs. Print a diagnostic error if a hidden argument is in
+        // the argument list and is not preloaded.
+        if (Arg.isOrigArg()) {
+          Argument *OrigArg = Fn.getArg(Arg.getOrigArgIndex());
+          if (OrigArg->hasAttribute("amdgpu-hidden-argument")) {
+            DiagnosticInfoUnsupported NonPreloadHiddenArg(
+                *OrigArg->getParent(),
+                "hidden argument in kernel signature was not preloaded",
+                DL.getDebugLoc());
+            DAG.getContext()->diagnose(NonPreloadHiddenArg);
+          }
+        }
+
         NewArg =
             lowerKernargMemParameter(DAG, VT, MemVT, DL, Chain, Offset,
                                      Alignment, Ins[i].Flags.isSExt(), &Ins[i]);

llvm/test/CodeGen/AMDGPU/invalid-hidden-kernarg-in-kernel-signature.ll

@@ -0,0 +1,21 @@
+; RUN: not llc -global-isel=1 -global-isel-abort=2 -mtriple=amdgcn--amdhsa -mcpu=gfx942 < %s 2>&1 | FileCheck -check-prefixes=ERROR,GISEL %s
+; RUN: not llc -global-isel=0 -mtriple=amdgcn--amdhsa -mcpu=gfx942 < %s 2>&1 | FileCheck -check-prefix=ERROR %s
+; RUN: not llc -global-isel=1 -global-isel-abort=2 -amdgpu-ir-lower-kernel-arguments=0 -mtriple=amdgcn--amdhsa -mcpu=gfx942 < %s 2>&1 | FileCheck -check-prefixes=ERROR,GISEL %s
+; RUN: not llc -global-isel=0 -amdgpu-ir-lower-kernel-arguments=0 -mtriple=amdgcn--amdhsa -mcpu=gfx942 < %s 2>&1 | FileCheck -check-prefix=ERROR %s
+
+define amdgpu_kernel void @no_free_sgprs_block_count_x_no_preload_diag(ptr addrspace(1) inreg %out, i512 inreg, i32 inreg "amdgpu-hidden-argument" %_hidden_block_count_x) #0 {
+; GISEL: warning: Instruction selection used fallback path for no_free_sgprs_block_count_x_no_preload_diag
+; ERROR: error: <unknown>:0:0: in function no_free_sgprs_block_count_x_no_preload_diag void (ptr addrspace(1), i512, i32): hidden argument in kernel signature was not preloaded
+  store i32 %_hidden_block_count_x, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @preloadremainder_z_no_preload_diag(ptr addrspace(1) inreg %out, i256 inreg, i32 inreg "amdgpu-hidden-argument" %_hidden_block_count_x, i32 inreg "amdgpu-hidden-argument" %_hidden_block_count_y, i32 inreg "amdgpu-hidden-argument" %_hidden_block_count_z, i16 inreg "amdgpu-hidden-argument" %_hidden_group_size_x, i16 inreg "amdgpu-hidden-argument" %_hidden_group_size_y, i16 inreg "amdgpu-hidden-argument" %_hidden_group_size_z, i16 inreg "amdgpu-hidden-argument" %_hidden_remainder_x, i16 inreg "amdgpu-hidden-argument" %_hidden_remainder_y, i16 inreg "amdgpu-hidden-argument" %_hidden_remainder_z) #0 {
+; GISEL: warning: Instruction selection used fallback path for preloadremainder_z_no_preload_diag
+; ERROR: error: <unknown>:0:0: in function preloadremainder_z_no_preload_diag void (ptr addrspace(1), i256, i32, i32, i32, i16, i16, i16, i16, i16, i16): hidden argument in kernel signature was not preloaded
+  %conv = zext i16 %_hidden_remainder_z to i32
+  store i32 %conv, ptr addrspace(1) %out
+  ret void
+}
+
+attributes #0 = { "amdgpu-no-agpr" "amdgpu-no-completion-action" "amdgpu-no-default-queue" "amdgpu-no-dispatch-id" "amdgpu-no-dispatch-ptr" "amdgpu-no-heap-ptr" "amdgpu-no-hostcall-ptr" "amdgpu-no-lds-kernel-id" "amdgpu-no-multigrid-sync-arg" "amdgpu-no-queue-ptr" "amdgpu-no-workgroup-id-x" "amdgpu-no-workgroup-id-y" "amdgpu-no-workgroup-id-z" "amdgpu-no-workitem-id-x" "amdgpu-no-workitem-id-y" "amdgpu-no-workitem-id-z" "uniform-work-group-size"="false" }

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

@@ -599,10 +599,8 @@ define amdgpu_kernel void @no_free_sgprs_preloadremainder_z(ptr addrspace(1) inr
 ; GFX940:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
 ; GFX940-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
 ; GFX940-NEXT:  ; %bb.0:
-; GFX940-NEXT:    s_load_dword s0, s[4:5], 0x1c
+; GFX940-NEXT:    s_lshr_b32 s0, s15, 16
 ; GFX940-NEXT:    v_mov_b32_e32 v0, 0
-; GFX940-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX940-NEXT:    s_lshr_b32 s0, s0, 16
 ; GFX940-NEXT:    v_mov_b32_e32 v1, s0
 ; GFX940-NEXT:    global_store_dword v0, v1, s[8:9] sc0 sc1
 ; GFX940-NEXT:    s_endpgm
@@ -626,4 +624,74 @@ define amdgpu_kernel void @no_free_sgprs_preloadremainder_z(ptr addrspace(1) inr
   ret void
 }
 
+; Check for consistency between isel and earlier passes preload SGPR accounting with max preload SGPRs.
+
+define amdgpu_kernel void @preload_block_max_user_sgprs(ptr addrspace(1) inreg %out, i192 inreg %t0, i32 inreg %t1) #0 {
+; GFX940-LABEL: preload_block_max_user_sgprs:
+; GFX940:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; 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, s12
+; GFX940-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-NEXT:    s_endpgm
+;
+; GFX90a-LABEL: preload_block_max_user_sgprs:
+; GFX90a:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-NEXT:  ; %bb.0:
+; GFX90a-NEXT:    s_load_dword s0, s[4:5], 0x28
+; GFX90a-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX90a-NEXT:    v_mov_b32_e32 v1, s0
+; 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()
+  %load = load i32, ptr addrspace(4) %imp_arg_ptr
+  store i32 %load, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @preload_block_count_z_workgroup_size_z_remainder_z(ptr addrspace(1) inreg %out) #0 {
+; GFX940-LABEL: preload_block_count_z_workgroup_size_z_remainder_z:
+; GFX940:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-NEXT:  ; %bb.0:
+; GFX940-NEXT:    s_lshr_b32 s0, s9, 16
+; GFX940-NEXT:    s_and_b32 s1, s8, 0xffff
+; GFX940-NEXT:    v_mov_b32_e32 v3, 0
+; GFX940-NEXT:    v_mov_b32_e32 v0, s6
+; GFX940-NEXT:    v_mov_b32_e32 v1, s1
+; GFX940-NEXT:    v_mov_b32_e32 v2, s0
+; GFX940-NEXT:    global_store_dwordx3 v3, v[0:2], s[2:3] sc0 sc1
+; GFX940-NEXT:    s_endpgm
+;
+; GFX90a-LABEL: preload_block_count_z_workgroup_size_z_remainder_z:
+; GFX90a:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-NEXT:  ; %bb.0:
+; GFX90a-NEXT:    s_lshr_b32 s0, s13, 16
+; GFX90a-NEXT:    s_and_b32 s1, s12, 0xffff
+; GFX90a-NEXT:    v_mov_b32_e32 v3, 0
+; GFX90a-NEXT:    v_mov_b32_e32 v0, s10
+; GFX90a-NEXT:    v_mov_b32_e32 v1, s1
+; GFX90a-NEXT:    v_mov_b32_e32 v2, s0
+; GFX90a-NEXT:    global_store_dwordx3 v3, v[0:2], s[6:7]
+; GFX90a-NEXT:    s_endpgm
+  %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %gep0 = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 8
+  %gep1 = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 16
+  %gep2 = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 22
+  %load0 = load i32, ptr addrspace(4) %gep0
+  %load1 = load i16, ptr addrspace(4) %gep1
+  %load2 = load i16, ptr addrspace(4) %gep2
+  %conv1 = zext i16 %load1 to i32
+  %conv2 = zext i16 %load2 to i32
+  %ins.0 =  insertelement <3 x i32> poison, i32 %load0, i32 0
+  %ins.1 =  insertelement <3 x i32> %ins.0, i32 %conv1, i32 1
+  %ins.2 =  insertelement <3 x i32> %ins.1, i32 %conv2, i32 2
+  store <3 x i32> %ins.2, ptr addrspace(1) %out
+  ret void
+}
+
 attributes #0 = { "amdgpu-no-agpr" "amdgpu-no-completion-action" "amdgpu-no-default-queue" "amdgpu-no-dispatch-id" "amdgpu-no-dispatch-ptr" "amdgpu-no-heap-ptr" "amdgpu-no-hostcall-ptr" "amdgpu-no-lds-kernel-id" "amdgpu-no-multigrid-sync-arg" "amdgpu-no-queue-ptr" "amdgpu-no-workgroup-id-x" "amdgpu-no-workgroup-id-y" "amdgpu-no-workgroup-id-z" "amdgpu-no-workitem-id-x" "amdgpu-no-workitem-id-y" "amdgpu-no-workitem-id-z" "uniform-work-group-size"="false" }