[HAL][AMDGPU] Support raw HSACO custom kernargs

Allow AMDGPU raw HSACO modules to expose HIP-visible kernel symbols that have companion kernel descriptors, even when the code object metadata only describes a subset of exports.

Preserve native kernarg layouts for those raw kernels by widening reflected parameter sizes and offsets, accepting pre-packed HIP argument blobs that omit trailing ABI padding, zero-filling missing bytes, and synthesizing HIP implicit args at the metadata-derived or HSA-reported suffix offset.

Use ELF symbol metadata and kernel descriptors to identify additional raw code object exports and keep their dispatch ABI consistent with the loaded HSA executable.

Signed-off-by: Andrew Woloszyn <andrew.woloszyn@gmail.com>

Author: AWoloszyn

runtime/src/iree/hal/drivers/amdgpu/abi/kernel_args.h

@@ -54,8 +54,11 @@ typedef struct iree_hal_amdgpu_device_kernel_args_t {
   uint16_t constant_count;
   // Total number of bindings used by the dispatch (if a HAL dispatch).
   uint16_t binding_count;
+  // Offset of HIP/OpenCL implicit args within custom-direct kernargs, or
+  // UINT16_MAX if the kernel does not use an implicit suffix.
+  uint16_t implicit_args_offset;
   // Reserved for future hot kernel metadata. Must be zero.
-  uint32_t reserved;
+  uint16_t reserved;
 } iree_hal_amdgpu_device_kernel_args_t;
 IREE_AMDGPU_STATIC_ASSERT(
     sizeof(iree_hal_amdgpu_device_kernel_args_t) <= 64,

runtime/src/iree/hal/drivers/amdgpu/aql_command_buffer.c

@@ -1690,11 +1690,32 @@ static iree_status_t iree_hal_amdgpu_aql_command_buffer_write_dispatch_tail(
       }
       return iree_ok_status();
     }
-    case IREE_HAL_AMDGPU_COMMAND_BUFFER_KERNARG_STRATEGY_CUSTOM_DIRECT:
-      if (constants.data_length > 0) {
-        memcpy(tail_payload, constants.data, constants.data_length);
+    case IREE_HAL_AMDGPU_COMMAND_BUFFER_KERNARG_STRATEGY_CUSTOM_DIRECT: {
+      // Custom-direct callers hand us a pre-packed HIP ABI blob. Some raw
+      // HSACO kernels also need an implicit suffix synthesized after that blob,
+      // so zero the full reservation before copying only caller-owned bytes.
+      const iree_host_size_t total_kernarg_size =
+          layout->total_kernarg_size ? layout->total_kernarg_size
+                                     : constants.data_length;
+      memset(tail_payload, 0, total_kernarg_size);
+      const iree_host_size_t explicit_bytes =
+          layout->has_implicit_args ? layout->implicit_args_offset
+                                    : total_kernarg_size;
+      const iree_host_size_t copy_bytes =
+          constants.data_length < explicit_bytes ? constants.data_length
+                                                 : explicit_bytes;
+      if (copy_bytes > 0) {
+        memcpy(tail_payload, constants.data, copy_bytes);
+      }
+      if (layout->has_implicit_args) {
+        iree_amdgpu_kernel_implicit_args_t* implicit_args =
+            (iree_amdgpu_kernel_implicit_args_t*)(tail_payload +
+                                                  layout->implicit_args_offset);
+        iree_hal_amdgpu_aql_command_buffer_write_implicit_args(
+            kernel_args, config, implicit_args);
       }
       return iree_ok_status();
+    }
     case IREE_HAL_AMDGPU_COMMAND_BUFFER_KERNARG_STRATEGY_INDIRECT:
       return iree_make_status(
           IREE_STATUS_UNIMPLEMENTED,
@@ -2017,18 +2038,40 @@ static iree_status_t iree_hal_amdgpu_aql_command_buffer_prepare_dispatch_plan(
 
   if (iree_hal_amdgpu_aql_dispatch_plan_uses_custom_direct_arguments(
           out_plan)) {
-    if (IREE_UNLIKELY(inputs->constants.data_length !=
-                      out_plan->descriptor->kernel_args.kernarg_size)) {
+    // Callers (e.g. rocBLAS/Tensile) sometimes omit trailing ABI padding or pad
+    // beyond the declared kernarg_segment_size with extra trailing scalars. The
+    // kernel only reads its declared size, so trailing bytes are ignored and the
+    // memcpy in write_dispatch_tail clamps to the declared size.
+    //
+    // Validate after 8-byte ABI padding so we accept missing tail padding while
+    // still rejecting truly short pre-packed HIP argument buffers.
+    const uint32_t required_explicit_bytes =
+        (uint32_t)out_plan->descriptor->custom_kernarg_layout
+            .explicit_kernarg_size;
+    const iree_host_size_t padded_constant_length =
+        iree_host_align(inputs->constants.data_length, /*alignment=*/8);
+    if (IREE_UNLIKELY(padded_constant_length < required_explicit_bytes)) {
       return iree_make_status(
           IREE_STATUS_INVALID_ARGUMENT,
-          "custom dispatch argument length mismatch; expected %u but got "
-          "%" PRIhsz,
-          out_plan->descriptor->kernel_args.kernarg_size,
-          inputs->constants.data_length);
+          "custom dispatch argument length too short; expected at least %u "
+          "but got %" PRIhsz " (padded to %" PRIhsz ")",
+          required_explicit_bytes, inputs->constants.data_length,
+          padded_constant_length);
     }
     out_plan->layout = &out_plan->descriptor->custom_kernarg_layout;
     out_plan->kernarg_block_count =
         iree_max(1u, out_plan->descriptor->custom_kernarg_block_count);
+    if (out_plan->layout->total_kernarg_size == 0 &&
+        inputs->constants.data_length > 0) {
+      // Some raw kernels have no reflected kernarg size. In that case the
+      // caller-provided blob is the only reservation size we can trust.
+      const uint32_t provided_kernarg_block_count =
+          (uint32_t)iree_host_size_ceil_div(
+              inputs->constants.data_length,
+              sizeof(iree_hal_amdgpu_kernarg_block_t));
+      out_plan->kernarg_block_count =
+          iree_max(out_plan->kernarg_block_count, provided_kernarg_block_count);
+    }
     out_plan->kernarg_strategy =
         IREE_HAL_AMDGPU_COMMAND_BUFFER_KERNARG_STRATEGY_CUSTOM_DIRECT;
     return iree_ok_status();
@@ -2102,14 +2145,18 @@ iree_hal_amdgpu_aql_command_buffer_calculate_dispatch_layout(
           ? 0
           : (iree_host_size_t)plan->kernel_args->binding_count *
                 sizeof(uint64_t);
-  const iree_host_size_t tail_byte_length =
-      plan->layout->total_kernarg_size - binding_bytes;
+  const iree_host_size_t total_kernarg_size =
+      iree_hal_amdgpu_aql_dispatch_plan_uses_custom_direct_arguments(plan) &&
+              plan->layout->total_kernarg_size == 0
+          ? inputs->constants.data_length
+          : plan->layout->total_kernarg_size;
+  const iree_host_size_t tail_byte_length = total_kernarg_size - binding_bytes;
   IREE_RETURN_IF_ERROR(iree_hal_amdgpu_aql_command_buffer_qword_length(
       tail_byte_length, "dispatch tail payload",
       &out_layout->kernarg.tail_length_qwords,
       &out_layout->kernarg.tail_padded_length));
   IREE_RETURN_IF_ERROR(iree_hal_amdgpu_aql_command_buffer_qword_length(
-      plan->layout->total_kernarg_size, "dispatch kernarg",
+      total_kernarg_size, "dispatch kernarg",
       &out_layout->kernarg.total_length_qwords,
       &out_layout->kernarg.total_padded_length));
   out_layout->kernarg.implicit_args_offset_qwords =

runtime/src/iree/hal/drivers/amdgpu/device/dispatch.c

@@ -96,14 +96,34 @@ void iree_hal_amdgpu_device_dispatch_emplace_hal_kernargs(
 }
 
 void iree_hal_amdgpu_device_dispatch_emplace_custom_kernargs(
+    const iree_hal_amdgpu_device_kernel_args_t* IREE_AMDGPU_RESTRICT
+        kernel_args,
+    const uint32_t workgroup_count[3], uint32_t dynamic_workgroup_local_memory,
     const iree_hal_amdgpu_device_dispatch_kernarg_layout_t* IREE_AMDGPU_RESTRICT
         layout,
     const void* IREE_AMDGPU_RESTRICT custom_kernarg_ptr,
+    size_t custom_kernarg_length,
     void* IREE_AMDGPU_RESTRICT kernarg_ptr) {
-  if (layout->total_kernarg_size > 0) {
-    iree_amdgpu_memcpy(kernarg_ptr, custom_kernarg_ptr,
-                       layout->total_kernarg_size);
+  const size_t total_kernarg_size =
+      layout->total_kernarg_size ? layout->total_kernarg_size
+                                 : custom_kernarg_length;
+  if (total_kernarg_size > 0) {
+    iree_amdgpu_memset(kernarg_ptr, 0, total_kernarg_size);
+    const size_t explicit_bytes =
+        layout->has_implicit_args
+            ? layout->implicit_args_offset
+            : total_kernarg_size;
+    const size_t copy_bytes =
+        custom_kernarg_length < explicit_bytes ? custom_kernarg_length
+                                               : explicit_bytes;
+    if (copy_bytes > 0) {
+      iree_amdgpu_memcpy(kernarg_ptr, custom_kernarg_ptr, copy_bytes);
+    }
   }
+
+  iree_hal_amdgpu_device_dispatch_emplace_implicit_args(
+      kernel_args, workgroup_count, dynamic_workgroup_local_memory, layout,
+      kernarg_ptr);
 }
 
 //===----------------------------------------------------------------------===//

runtime/src/iree/hal/drivers/amdgpu/device/dispatch.h

@@ -178,18 +178,24 @@ void iree_hal_amdgpu_device_dispatch_emplace_hal_kernargs(
 
 // Populates custom direct kernargs in already-reserved storage.
 //
-// |custom_kernarg_ptr| must provide |layout->total_kernarg_size| bytes in the
-// final kernel ABI shape expected by the target kernel.
+// |custom_kernarg_ptr| provides up to |layout->total_kernarg_size| bytes in the
+// final kernel ABI shape expected by the target kernel. Missing trailing padding
+// bytes remain zeroed.
 //
 // Preconditions:
-//   - |layout| and |kernarg_ptr| are non-NULL.
+//   - |kernel_args|, |workgroup_count|, |layout|, and |kernarg_ptr| are
+//     non-NULL.
 //   - |layout| was derived with
 //     iree_hal_amdgpu_device_dispatch_make_custom_kernarg_layout.
-//   - |custom_kernarg_ptr| is non-NULL when |layout->total_kernarg_size| > 0.
+//   - |custom_kernarg_ptr| is non-NULL when |custom_kernarg_length| > 0.
 void iree_hal_amdgpu_device_dispatch_emplace_custom_kernargs(
+    const iree_hal_amdgpu_device_kernel_args_t* IREE_AMDGPU_RESTRICT
+        kernel_args,
+    const uint32_t workgroup_count[3], uint32_t dynamic_workgroup_local_memory,
     const iree_hal_amdgpu_device_dispatch_kernarg_layout_t* IREE_AMDGPU_RESTRICT
         layout,
     const void* IREE_AMDGPU_RESTRICT custom_kernarg_ptr,
+    size_t custom_kernarg_length,
     void* IREE_AMDGPU_RESTRICT kernarg_ptr);
 
 // Populates the builtin patch dispatch that updates an indirect-parameter

runtime/src/iree/hal/drivers/amdgpu/executable.c

@@ -965,6 +965,7 @@ static iree_status_t iree_hal_amdgpu_executable_resolve_kernel_args_from_symbol(
 
   out_kernel_args->binding_count = binding_count;
   out_kernel_args->constant_count = constant_count;
+  out_kernel_args->implicit_args_offset = UINT16_MAX;
 
   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
@@ -1141,6 +1142,21 @@ static iree_status_t iree_hal_amdgpu_executable_initialize_dispatch_descriptor(
   out_descriptor->custom_kernarg_layout =
       iree_hal_amdgpu_device_dispatch_make_custom_kernarg_layout(
           kernel_args->kernarg_size);
+  const uint16_t custom_implicit_args_offset =
+      kernel_args->implicit_args_offset != UINT16_MAX
+          ? kernel_args->implicit_args_offset
+          : kernel_args->kernarg_size;
+  if (custom_implicit_args_offset != UINT16_MAX) {
+    out_descriptor->custom_kernarg_layout.explicit_kernarg_size =
+        custom_implicit_args_offset;
+    out_descriptor->custom_kernarg_layout.implicit_args_offset =
+        custom_implicit_args_offset;
+    out_descriptor->custom_kernarg_layout.total_kernarg_size = iree_max(
+        (size_t)kernel_args->kernarg_size,
+        (size_t)custom_implicit_args_offset +
+            IREE_AMDGPU_KERNEL_IMPLICIT_ARGS_SIZE);
+    out_descriptor->custom_kernarg_layout.has_implicit_args = true;
+  }
   IREE_RETURN_IF_ERROR(iree_hal_amdgpu_executable_calculate_kernarg_block_count(
       &out_descriptor->custom_kernarg_layout,
       &out_descriptor->custom_kernarg_block_count));
@@ -1656,13 +1672,6 @@ iree_hal_amdgpu_executable_calculate_raw_hsaco_reflection_storage(
     iree_host_size_t* out_export_name_storage_size,
     iree_host_size_t* out_export_parameter_count,
     iree_host_size_t* out_export_parameter_name_storage_size) {
-  if (iree_string_view_is_empty(hsaco_metadata->target)) {
-    return iree_make_status(
-        IREE_STATUS_INVALID_ARGUMENT,
-        "raw HSACO metadata is missing `amdhsa.target`; direct loading "
-        "requires the code object to declare its target ISA");
-  }
-
   iree_host_size_t export_name_storage_size = 0;
   iree_host_size_t export_parameter_count = 0;
   iree_host_size_t export_parameter_name_storage_size = 0;
@@ -1922,6 +1931,60 @@ static iree_status_t iree_hal_amdgpu_executable_resolve_raw_hsaco_kernel_args(
             requirements.binding_count, &host_kernel_args[kernel_ordinal]),
         "resolving kernel args for raw kernel `%.*s`", (int)symbol_name.size,
         symbol_name.data);
+    // Raw HSACO metadata is the source of truth for caller-visible kernargs:
+    // some code objects report a smaller HSA symbol size than the metadata
+    // segment needed by pre-packed HIP launch buffers.
+    if (host_kernel_args[kernel_ordinal].kernarg_size <
+        kernel->kernarg_segment_size) {
+      host_kernel_args[kernel_ordinal].kernarg_size =
+          kernel->kernarg_segment_size;
+    }
+    if (host_kernel_args[kernel_ordinal].kernarg_alignment <
+        kernel->kernarg_segment_alignment) {
+      host_kernel_args[kernel_ordinal].kernarg_alignment =
+          kernel->kernarg_segment_alignment;
+    }
+
+    uint16_t implicit_args_offset = UINT16_MAX;
+    uint32_t explicit_args_end = 0;
+    for (iree_host_size_t arg_i = 0; arg_i < kernel->arg_count; ++arg_i) {
+      const iree_hal_amdgpu_hsaco_metadata_arg_t* arg =
+          &kernel->args[arg_i];
+      if (arg->kind == IREE_HAL_AMDGPU_HSACO_METADATA_ARG_KIND_HIDDEN ||
+          arg->kind == IREE_HAL_AMDGPU_HSACO_METADATA_ARG_KIND_HIDDEN_NONE) {
+        if (arg->offset <= UINT16_MAX &&
+            (implicit_args_offset == UINT16_MAX ||
+             arg->offset < implicit_args_offset)) {
+          implicit_args_offset = (uint16_t)arg->offset;
+        }
+      } else {
+        iree_host_size_t arg_end = 0;
+        if (!iree_host_size_checked_add(arg->offset, arg->size, &arg_end) ||
+            arg_end > UINT32_MAX) {
+          return iree_make_status(
+              IREE_STATUS_OUT_OF_RANGE,
+              "AMDGPU kernel `%.*s` argument offset overflows",
+              (int)symbol_name.size, symbol_name.data);
+        }
+        explicit_args_end = iree_max(explicit_args_end, (uint32_t)arg_end);
+      }
+    }
+    if (implicit_args_offset != UINT16_MAX && implicit_args_offset > 0) {
+      host_kernel_args[kernel_ordinal].implicit_args_offset =
+          implicit_args_offset;
+    } else if (explicit_args_end <= UINT16_MAX &&
+               host_kernel_args[kernel_ordinal].kernarg_size >=
+                   explicit_args_end + IREE_AMDGPU_KERNEL_IMPLICIT_ARGS_SIZE) {
+      host_kernel_args[kernel_ordinal].implicit_args_offset =
+          (uint16_t)explicit_args_end;
+    } else if (kernel->uses_borrowed_arg_layout &&
+               host_kernel_args[kernel_ordinal].kernarg_size > 0) {
+      // ELF-synthesized exports borrow a template arg layout, so its hidden arg
+      // offsets may not apply. Preserve the full caller-provided raw kernarg
+      // blob and append the implicit suffix after the HSA-reported segment.
+      host_kernel_args[kernel_ordinal].implicit_args_offset =
+          host_kernel_args[kernel_ordinal].kernarg_size;
+    }
   }
   return iree_ok_status();
 }

runtime/src/iree/hal/drivers/amdgpu/host_queue_dispatch.c

@@ -221,13 +221,17 @@ static iree_status_t iree_hal_amdgpu_host_queue_validate_dispatch_kernargs(
 
   iree_host_size_t operation_resource_count = 1;
   if (iree_any_bit_set(flags, IREE_HAL_DISPATCH_FLAG_CUSTOM_DIRECT_ARGUMENTS)) {
-    if (IREE_UNLIKELY(constants.data_length !=
-                      descriptor->kernel_args.kernarg_size)) {
+    const uint32_t required_explicit_bytes =
+        (uint32_t)descriptor->custom_kernarg_layout.explicit_kernarg_size;
+    const iree_host_size_t padded_constant_length =
+        iree_host_align(constants.data_length, /*alignment=*/8);
+    if (IREE_UNLIKELY(padded_constant_length < required_explicit_bytes)) {
       return iree_make_status(
           IREE_STATUS_INVALID_ARGUMENT,
-          "custom dispatch argument length mismatch; expected %u but got "
-          "%" PRIhsz,
-          descriptor->kernel_args.kernarg_size, constants.data_length);
+          "custom dispatch argument length too short; expected at least %u "
+          "but got %" PRIhsz " (padded to %" PRIhsz ")",
+          required_explicit_bytes, constants.data_length,
+          padded_constant_length);
     }
     if (IREE_UNLIKELY(constants.data_length > 0 && !constants.data)) {
       return iree_make_status(
@@ -238,6 +242,13 @@ static iree_status_t iree_hal_amdgpu_host_queue_validate_dispatch_kernargs(
     *out_layout = &descriptor->custom_kernarg_layout;
     *out_kernarg_block_count =
         iree_max(1u, descriptor->custom_kernarg_block_count);
+    if ((*out_layout)->total_kernarg_size == 0 && constants.data_length > 0) {
+      const uint32_t provided_kernarg_block_count =
+          (uint32_t)iree_host_size_ceil_div(
+              constants.data_length, sizeof(iree_hal_amdgpu_kernarg_block_t));
+      *out_kernarg_block_count =
+          iree_max(*out_kernarg_block_count, provided_kernarg_block_count);
+    }
   } else {
     if (IREE_UNLIKELY(constants.data_length > 0 && !constants.data)) {
       return iree_make_status(
@@ -531,7 +542,9 @@ static iree_status_t iree_hal_amdgpu_host_queue_submit_direct_dispatch(
 
   if (uses_custom_direct_arguments) {
     iree_hal_amdgpu_device_dispatch_emplace_custom_kernargs(
-        plan->layout, constants.data, submission.kernel.kernargs.blocks->data);
+        plan->kernel_args, config.workgroup_count,
+        config.dynamic_workgroup_local_memory, plan->layout, constants.data,
+        constants.data_length, submission.kernel.kernargs.blocks->data);
   } else {
     iree_hal_amdgpu_device_dispatch_emplace_hal_kernargs(
         plan->kernel_args, config.workgroup_count,
@@ -725,7 +738,9 @@ static iree_status_t iree_hal_amdgpu_host_queue_submit_indirect_dispatch(
   const uint32_t placeholder_workgroup_count[3] = {0, 0, 0};
   if (uses_custom_direct_arguments) {
     iree_hal_amdgpu_device_dispatch_emplace_custom_kernargs(
-        plan->layout, constants.data, dispatch_kernarg_data);
+        plan->kernel_args, placeholder_workgroup_count,
+        config.dynamic_workgroup_local_memory, plan->layout, constants.data,
+        constants.data_length, dispatch_kernarg_data);
   } else {
     iree_hal_amdgpu_device_dispatch_emplace_hal_kernargs(
         plan->kernel_args, placeholder_workgroup_count,