[AMDGPU] Add dynamic LDS size implicit kernel argument to v5

[skc7]

Summary

Add hidden kernel argument "hidden_dynamic_lds_size" to CO v5.

This argument will have the size of dynamic LDS memory allocated by hip runtime to the kernel.

[arsenm]

Why do we need this? I thought this was already available in the dispatch packet, you just need to subtract the statically known size

llvm/docs/AMDGPUUsage.rst

3964

S/passes/passed

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

1106

This just just be an i32

llvm/test/CodeGen/AMDGPU/hsa-metadata-hidden-args-v5.ll

95

This should have optimized out

[skc7]

Make lds size argument 32 bit int type.

[skc7]

In D158150#4594958, @arsenm wrote:

Why do we need this? I thought this was already available in the dispatch packet, you just need to subtract the statically known size

For ASAN LDS support, we want runtime to place to dynamic LDS size as hidden kernel argument. This is the approach we assume might work. (considering if recycling of dispatch pointer would be done in future runtime version)
cc: @b-sumner

[b-sumner]

In D158150#4595052, @skc7 wrote:

In D158150#4594958, @arsenm wrote:

Why do we need this? I thought this was already available in the dispatch packet, you just need to subtract the statically known size

For ASAN LDS support, we want runtime to place to dynamic LDS size as hidden kernel argument. This is the approach we assume might work. (considering if recycling of dispatch pointer would be done in future runtime version)
cc: @b-sumner

With code object v5 we eliminated the need for the dispatch packet pointer and don't want to bring it back.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

1106

It also doesn't need 8 byte alignment

[skc7]

Update alignment for int32 type lds_size argument

[arsenm]

Discussed not bothering defining metadata when this is implied to the runtime by using dynamic LDS

llvm/docs/AMDGPUUsage.rst

".value_kind"          string         Required  Kernel argument kind that
                                                        and flat addresses. See :ref:`amdgpu-amdhsa-kernel-prolog-flat-scratch`.
 
                                                      "hidden_queue_ptr"
                                                        A global memory address space pointer to the ROCm runtime
                                                        ``struct amd_queue_t`` structure for the HSA queue of the
                                                        associated dispatch AQL packet. It is only required for pre-GFX9
                                                        devices for the trap handler ABI (see :ref:`amdgpu-amdhsa-trap-handler-abi`).
 
+                                                     "hidden_dynamic_lds_size"
+                                                       Size of the dynamically allocated LDS memory is passed in the kernarg.
+
      ====================== ============== ========= ================================
 
 ..
 
 Kernel Dispatch
 ~~~~~~~~~~~~~~~
 
 The HSA architected queuing language (AQL) defines a user space memory interface

llvm/lib/BinaryFormat/AMDGPUMetadataVerifier.cpp

if (!verifyScalarEntry(ArgsMap, ".value_kind", true, msgpack::Type::String,
                                .Case("hidden_hostcall_buffer", true)
                                .Case("hidden_heap_v1", true)
                                .Case("hidden_default_queue", true)
                                .Case("hidden_completion_action", true)
                                .Case("hidden_multigrid_sync_arg", true)
                                .Case("hidden_private_base", true)
                                .Case("hidden_shared_base", true)
                                .Case("hidden_queue_ptr", true)
+                               .Case("hidden_dynamic_lds_size", true)
                                .Default(false);
                          }))
     return false;
   if (!verifyIntegerEntry(ArgsMap, ".pointee_align", false))
     return false;
   if (!verifyScalarEntry(ArgsMap, ".address_space", false,
                          msgpack::Type::String,
                          [](msgpack::DocNode &SNode) {

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

if (!ST.hasApertureRegs()) {
     emitKernelArg(DL, Int32Ty, Align(4), "hidden_private_base", Offset, Args);
     emitKernelArg(DL, Int32Ty, Align(4), "hidden_shared_base", Offset, Args);
   } else {
     Offset += 8; // Skipped.
   }
 
   if (MFI.hasQueuePtr())
     emitKernelArg(DL, Int8PtrTy, Align(8), "hidden_queue_ptr", Offset, Args);
+  else
+    Offset += 8; // Skipped
+
+  // emit argument for hidden dynamic lds size
+  emitKernelArg(DL, Int32Ty, Align(4), "hidden_dynamic_lds_size", Offset, Args);
 }
 
 void MetadataStreamerMsgPackV5::emitKernelAttrs(const Function &Func,
                                                 msgpack::MapDocNode Kern) {
   MetadataStreamerMsgPackV3::emitKernelAttrs(Func, Kern);
 
   if (Func.getFnAttribute("uniform-work-group-size").getValueAsBool())
     Kern[".uniform_work_group_size"] = Kern.getDocument()->getNode(1);
 }
 
 
 } // end namespace HSAMD
 } // end namespace AMDGPU
 } // end namespace llvm

llvm/test/CodeGen/AMDGPU/hsa-metadata-hidden-args-v5.ll

 ; CHECK-NEXT:        .size:           8
 ; CHECK-NEXT:        .value_kind:     hidden_completion_action
 ; GFX8-NEXT:      - .offset:         216
 ; GFX8-NEXT:        .size:           4
 ; GFX8-NEXT:        .value_kind:     hidden_private_base
 ; GFX8-NEXT:      - .offset:         220
 ; GFX8-NEXT:        .size:           4
 ; GFX8-NEXT:        .value_kind:     hidden_shared_base
-; CHECK:      - .offset:         224
+; CHECK:          - .offset:         224
 ; CHECK-NEXT:        .size:           8
 ; CHECK-NEXT:        .value_kind:     hidden_queue_ptr
+; CHECK:          - .offset:         232
+; CHECK-NEXT:        .size:           4
+; CHECK-NEXT:        .value_kind:     hidden_dynamic_lds_size
 
 ; CHECK:          .name:           test_v5
 ; CHECK:          .symbol:         test_v5.kd
 
 ; CHECK:  amdhsa.version:
 ; CHECK-NEXT: - 1
 ; CHECK-NEXT: - 2
 define amdgpu_kernel void @test_v5(