[AMDGPU] Fix hidden kernarg preload count inconsistency

It is possible that the number of hidden arguments that are selected to
be preloaded in AMDGPULowerKernel arguments and isel can differ. This
isn't an issue with explicit arguments since isel can lower the argument
correctly either way, but with hidden arguments we may have alignment
issues if we try to load these hidden arguments that were added to the
kernel signature.

The reason for the mismatch is that isel reserves an extra synthetic
user SGPR for module LDS.

Instead of teaching lowerFormalArguments how to handle these properly it
makes more sense and is less expensive to fix the mismatch and assert if
we ever run into this issue again. We should never be trying to lower
these in the normal way.

In a future change we probably want to revise how we track "synthetic"
user SGPRs and unify the handling in GCNUserSGPRUsageInfo. Sometimes
synthetic SGPRSs are considered user SGPRs and sometimes they are not.
Until then this patch resolves the inconsistency, fixes the bug, and is
otherwise a NFC.

Author: kerbowa

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

@@ -144,17 +144,17 @@ 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() - 1 /* Synthetic SGPRs*/;
   }
 
   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 && !isAligned(Align(4), ArgOffset) &&
+        AllocSize < 4)
       return true;
 
     // Pad SGPRs for kernarg alignment.
@@ -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.
@@ -485,7 +485,7 @@ static bool lowerKernelArguments(Function &F, const TargetMachine &TM) {
     uint64_t ImplicitArgsBaseOffset =
         alignTo(ExplicitArgOffset, ST.getAlignmentForImplicitArgPtr()) +
         BaseOffset;
-    PreloadInfo.tryAllocImplicitArgPreloadSGPRs(ImplicitArgsBaseOffset,
+    PreloadInfo.tryAllocImplicitArgPreloadSGPRs(ImplicitArgsBaseOffset, ExplicitArgOffset,
                                                 Builder);
   }
 

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -3024,6 +3024,14 @@ SDValue SITargetLowering::LowerFormalArguments(
           NewArg = DAG.getMergeValues({NewArg, Chain}, DL);
         }
       } else {
+#ifndef NDEBUG
+        if (Arg.isOrigArg()) {
+          Argument *OrigArg = Fn.getArg(Arg.getOrigArgIndex());
+          assert(!OrigArg->hasAttribute("amdgpu-hidden-argument") &&
+                 "Hidden arguments should be preloaded");
+        }
+#endif // NDEBUG
+
         NewArg =
             lowerKernargMemParameter(DAG, VT, MemVT, DL, Chain, Offset,
                                      Alignment, Ins[i].Flags.isSExt(), &Ins[i]);

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

@@ -439,13 +439,13 @@ define amdgpu_kernel void @preload_workgroup_size_xyz(ptr addrspace(1) inreg %ou
 ; 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:    v_mov_b32_e32 v3, 0
+; GFX90a-NEXT:    global_load_ushort v2, v3, s[4:5] offset:24
 ; GFX90a-NEXT:    s_lshr_b32 s0, s11, 16
 ; GFX90a-NEXT:    s_and_b32 s1, s11, 0xffff
-; GFX90a-NEXT:    s_and_b32 s2, s12, 0xffff
-; GFX90a-NEXT:    v_mov_b32_e32 v3, 0
 ; GFX90a-NEXT:    v_mov_b32_e32 v0, s1
 ; GFX90a-NEXT:    v_mov_b32_e32 v1, s0
-; GFX90a-NEXT:    v_mov_b32_e32 v2, s2
+; GFX90a-NEXT:    s_waitcnt vmcnt(0)
 ; 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()
@@ -554,27 +554,28 @@ define amdgpu_kernel void @preloadremainder_xyz(ptr addrspace(1) inreg %out) #0
 ; 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_lshr_b32 s1, s8, 16
-; GFX940-NEXT:    s_and_b32 s4, s9, 0xffff
 ; GFX940-NEXT:    v_mov_b32_e32 v3, 0
-; GFX940-NEXT:    v_mov_b32_e32 v0, s1
-; GFX940-NEXT:    v_mov_b32_e32 v1, s4
-; GFX940-NEXT:    v_mov_b32_e32 v2, s0
+; GFX940-NEXT:    global_load_ushort v2, v3, s[0:1] offset:30
+; GFX940-NEXT:    s_lshr_b32 s0, s8, 16
+; GFX940-NEXT:    s_and_b32 s1, s9, 0xffff
+; GFX940-NEXT:    v_mov_b32_e32 v0, s0
+; GFX940-NEXT:    v_mov_b32_e32 v1, s1
+; GFX940-NEXT:    s_waitcnt vmcnt(0)
 ; GFX940-NEXT:    global_store_dwordx3 v3, v[0:2], s[2:3] sc0 sc1
 ; GFX940-NEXT:    s_endpgm
 ;
 ; GFX90a-LABEL: preloadremainder_xyz:
 ; 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_lshr_b32 s1, s12, 16
-; GFX90a-NEXT:    s_and_b32 s2, s13, 0xffff
 ; GFX90a-NEXT:    v_mov_b32_e32 v3, 0
-; GFX90a-NEXT:    v_mov_b32_e32 v0, s1
-; GFX90a-NEXT:    v_mov_b32_e32 v1, s2
-; GFX90a-NEXT:    v_mov_b32_e32 v2, s0
+; GFX90a-NEXT:    global_load_dword v0, v3, s[4:5] offset:26
+; GFX90a-NEXT:    global_load_ushort v2, v3, s[4:5] offset:30
+; GFX90a-NEXT:    s_lshr_b32 s0, s12, 16
+; GFX90a-NEXT:    s_waitcnt vmcnt(1)
+; GFX90a-NEXT:    v_lshrrev_b32_e32 v1, 16, v0
+; GFX90a-NEXT:    v_mov_b32_e32 v0, s0
+; GFX90a-NEXT:    s_waitcnt vmcnt(0)
 ; 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()
@@ -626,4 +627,32 @@ 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.
+
+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
+}
+
 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-kernargs-IR-lowering.ll

@@ -187,15 +187,17 @@ define amdgpu_kernel void @test_preload_IR_lowering_kernel_8(ptr addrspace(1) %i
 ; PRELOAD-8-LABEL: define {{[^@]+}}@test_preload_IR_lowering_kernel_8
 ; PRELOAD-8-SAME: (ptr addrspace(1) inreg [[IN:%.*]], ptr addrspace(1) inreg [[IN1:%.*]], ptr addrspace(1) inreg [[IN2:%.*]], ptr addrspace(1) inreg [[IN3:%.*]], ptr addrspace(1) inreg [[OUT:%.*]], ptr addrspace(1) inreg [[OUT1:%.*]], ptr addrspace(1) inreg [[OUT2:%.*]], ptr addrspace(1) inreg [[OUT3:%.*]]) #[[ATTR0]] {
 ; PRELOAD-8-NEXT:    [[TEST_PRELOAD_IR_LOWERING_KERNEL_8_KERNARG_SEGMENT:%.*]] = call nonnull align 16 dereferenceable(64) ptr addrspace(4) @llvm.amdgcn.kernarg.segment.ptr()
+; PRELOAD-8-NEXT:    [[OUT2_KERNARG_OFFSET:%.*]] = getelementptr inbounds i8, ptr addrspace(4) [[TEST_PRELOAD_IR_LOWERING_KERNEL_8_KERNARG_SEGMENT]], i64 48
+; PRELOAD-8-NEXT:    [[OUT2_LOAD:%.*]] = load ptr addrspace(1), ptr addrspace(4) [[OUT2_KERNARG_OFFSET]], align 16, !invariant.load [[META0:![0-9]+]]
 ; PRELOAD-8-NEXT:    [[OUT3_KERNARG_OFFSET:%.*]] = getelementptr inbounds i8, ptr addrspace(4) [[TEST_PRELOAD_IR_LOWERING_KERNEL_8_KERNARG_SEGMENT]], i64 56
-; PRELOAD-8-NEXT:    [[OUT3_LOAD:%.*]] = load ptr addrspace(1), ptr addrspace(4) [[OUT3_KERNARG_OFFSET]], align 8, !invariant.load [[META0:![0-9]+]]
+; PRELOAD-8-NEXT:    [[OUT3_LOAD:%.*]] = load ptr addrspace(1), ptr addrspace(4) [[OUT3_KERNARG_OFFSET]], align 8, !invariant.load [[META0]]
 ; PRELOAD-8-NEXT:    [[LOAD:%.*]] = load i32, ptr addrspace(1) [[IN]], align 4
 ; PRELOAD-8-NEXT:    [[LOAD1:%.*]] = load i32, ptr addrspace(1) [[IN1]], align 4
 ; PRELOAD-8-NEXT:    [[LOAD2:%.*]] = load i32, ptr addrspace(1) [[IN2]], align 4
 ; PRELOAD-8-NEXT:    [[LOAD3:%.*]] = load i32, ptr addrspace(1) [[IN3]], align 4
 ; PRELOAD-8-NEXT:    store i32 [[LOAD]], ptr addrspace(1) [[OUT]], align 4
 ; PRELOAD-8-NEXT:    store i32 [[LOAD1]], ptr addrspace(1) [[OUT1]], align 4
-; PRELOAD-8-NEXT:    store i32 [[LOAD2]], ptr addrspace(1) [[OUT2]], align 4
+; PRELOAD-8-NEXT:    store i32 [[LOAD2]], ptr addrspace(1) [[OUT2_LOAD]], align 4
 ; PRELOAD-8-NEXT:    store i32 [[LOAD3]], ptr addrspace(1) [[OUT3_LOAD]], align 4
 ; PRELOAD-8-NEXT:    ret void
 ;