Adding AMDGPU command buffer implementation. (iree-org#21123)

This contains the host-side `iree_hal_amdgpu_command_buffer_t`
responsible for recording and managing the replicas of the device-side
`iree_hal_amdgpu_device_command_buffer_t` instances and the device-side
implementation for issuing and executing command buffer programs.

The comments for `iree_hal_amdgpu_device_command_buffer_t` and its
related commands document how the command buffer is represented in
device memory and what each command does. Some of the commands are not
yet wired up to host recording code as the it is skating ahead of the
HAL command buffer API. Even though unconditional branches cannot be
recorded directly they are used for segmenting large command buffers to
fit within resource limits. Conditional branches are nominally supported
but need an API (similar to indirect dispatches) to be plumbed from the
compiler to the command buffer vtable.

IREE's HAL is intended to optimize for precompiled command buffers that
can amortize the recording cost - this implementation leans that
direction and should be optimal for meaningful (non-benchmark) programs
originating from anything targeting CUDA graphs. There are some
provisions for small command buffers but the intent is that
single-command buffers end up routed towards device queue operations
much earlier.

Since the device-side queue scheduler does not yet exist the two methods
required (`iree_hal_amdgpu_device_queue_scheduler_*`) are stubbed out to
indicate where they'll be used in the future. The API may change in
subsequent changes that implement the behavior.

Author: benvanik

runtime/src/iree/hal/drivers/amdgpu/BUILD.bazel

@@ -24,8 +24,8 @@ iree_runtime_cc_library(
         "buffer_pool.h",
         "channel.c",
         "channel.h",
-        # "command_buffer.c",
-        # "command_buffer.h",
+        "command_buffer.c",
+        "command_buffer.h",
         "driver.c",
         "driver.h",
         "event.c",
@@ -83,7 +83,7 @@ iree_runtime_cc_library(
         "buffer.h",
         "buffer_pool.h",
         "channel.h",
-        # "command_buffer.h",
+        "command_buffer.h",
         "driver.h",
         "event.h",
         "executable.h",

runtime/src/iree/hal/drivers/amdgpu/CMakeLists.txt

@@ -24,6 +24,8 @@ iree_cc_library(
     "buffer_pool.h"
     "channel.c"
     "channel.h"
+    "command_buffer.c"
+    "command_buffer.h"
     "driver.c"
     "driver.h"
     "event.c"
@@ -68,6 +70,7 @@ iree_cc_library(
     "buffer.h"
     "buffer_pool.h"
     "channel.h"
+    "command_buffer.h"
     "driver.h"
     "event.h"
     "executable.h"

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

@@ -0,0 +1,155 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_HAL_DRIVERS_AMDGPU_COMMAND_BUFFER_H_
+#define IREE_HAL_DRIVERS_AMDGPU_COMMAND_BUFFER_H_
+
+#include "iree/base/api.h"
+#include "iree/base/internal/arena.h"
+#include "iree/hal/api.h"
+#include "iree/hal/drivers/amdgpu/device/command_buffer.h"
+#include "iree/hal/drivers/amdgpu/util/affinity.h"
+
+typedef struct iree_hal_amdgpu_block_pools_t iree_hal_amdgpu_block_pools_t;
+
+//===----------------------------------------------------------------------===//
+// iree_hal_amdgpu_command_buffer_options_t
+//===----------------------------------------------------------------------===//
+
+// Determines where and how command buffers are recorded.
+typedef enum iree_hal_amdgpu_command_buffer_recording_flags_t {
+  IREE_HAL_AMDGPU_COMMAND_BUFFER_RECORDING_FLAG_NONE = 0u,
+
+  // TODO(benvanik): support lead-physical-device storage. This would need the
+  // block pool on the lead device to make its blocks accessible to all devices
+  // - today the block pool is device-local only. Produced data is immutable and
+  // PCIe atomics/coherency is not required across devices.
+  //
+  // Allocate embedded data on the lead physical device instead of on each
+  // device the command buffer is recorded for. This reduces overall memory
+  // consumption and recording time at the cost of cross-device transfers.
+  // IREE_HAL_AMDGPU_COMMAND_BUFFER_RECORDING_FLAG_DATA_ON_LEAD_PHYSICAL_DEVICE
+  // = 1u << 0,
+
+  // TODO(benvanik): support compaction. This would require changing the command
+  // buffer to use relative offsets for embedded data and a data table for
+  // indirecting so that we can move around base pointers. A fixup would be
+  // possible as well by launching a kernel that rebased the embedded pointers
+  // (though trickier). For now we assume the block pool block size is a big
+  // enough lever and most programs only use a handful of command buffers so
+  // the waste per command buffer is minimal (compared to a single layer weight
+  // in an ML model).
+  //
+  // Compacts the command buffer when recording ends by reallocating it to the
+  // precise size required and reuploads it to each device. This will return any
+  // block pool blocks back to their respective pool for reuse and ensure
+  // there's no unused device memory - the cost is extra host time to do the
+  // reallocation/copies.
+  // IREE_HAL_AMDGPU_COMMAND_BUFFER_RECORDING_FLAG_COMPACT_ON_FINALIZE
+  // = 1u << 1,
+} iree_hal_amdgpu_command_buffer_recording_flags_t;
+
+// TODO(benvanik): move this someplace common.
+//
+// Block pools for host memory blocks of various sizes.
+typedef struct iree_hal_amdgpu_host_block_pools_t {
+  // Used for small allocations of around 1-4KB.
+  iree_arena_block_pool_t small;
+  // Used for large page-sized allocations of 32-64kB.
+  iree_arena_block_pool_t large;
+} iree_hal_amdgpu_host_block_pools_t;
+
+// Minimum number of AQL packets in a single command buffer block.
+// Any fewer and it's not guaranteed a command buffer can complete execution.
+#define IREE_HAL_AMDGPU_COMMAND_BUFFER_MIN_BLOCK_AQL_PACKET_COUNT (16)
+
+// Maximum number of AQL packets in a single command buffer block.
+// This is currently limited by the `uint16_t packet_offset` in
+// iree_hal_amdgpu_device_cmd_header_t.
+//
+// TODO(benvanik): currently we also limit this by tracy's outstanding GPU event
+// limit. If we made our own timeline (which we really need to for concurrency)
+// then we could eliminate this artificial limit.
+#define IREE_HAL_AMDGPU_COMMAND_BUFFER_MAX_BLOCK_AQL_PACKET_COUNT            \
+  IREE_AMDGPU_MIN(IREE_HAL_AMDGPU_DEVICE_QUERY_RINGBUFFER_CAPACITY,          \
+                  (1u << sizeof(((iree_hal_amdgpu_device_cmd_header_t*)NULL) \
+                                    ->packet_offset) *                       \
+                             8))
+
+// Recording options for a command buffer.
+// Referenced data structures such as block pools must remain live for the
+// lifetime of the command buffer but the options struct and its storage (such
+// as the device block pool list) need not.
+typedef struct iree_hal_amdgpu_command_buffer_options_t {
+  iree_hal_allocator_t* device_allocator;
+  iree_hal_command_buffer_mode_t mode;
+  iree_hal_command_category_t command_categories;
+  iree_hal_queue_affinity_t queue_affinity;
+  iree_host_size_t binding_capacity;
+
+  // Controls recording behavior (placement, optimization, debugging, etc).
+  iree_hal_amdgpu_command_buffer_recording_flags_t recording_flags;
+
+  // Maximum number of AQL packets the command buffer is allowed to issue at
+  // a time. Must be at or under the HSA queue capacity of any execution queue
+  // the command buffer will be scheduled on. The command buffer may decide to
+  // use fewer packets.
+  iree_host_size_t block_aql_packet_count;
+
+  // Block pools for host memory blocks of various sizes.
+  iree_hal_amdgpu_host_block_pools_t* host_block_pools;
+
+  // Bitmap of physical devices that the command buffer will be recorded for.
+  // The command buffer can only be issued on these devices.
+  iree_hal_amdgpu_device_affinity_t device_affinity;
+
+  // Compact list of physical device block pools corresponding to the bits set
+  // in the device_affinity bitmap. A device affinity of 0b110 would lead to two
+  // device block pools in the list at [0] and [1].
+  iree_hal_amdgpu_block_pools_t* const* device_block_pools /*[device_count]*/;
+} iree_hal_amdgpu_command_buffer_options_t;
+
+// Initializes |out_options| to its default values.
+void iree_hal_amdgpu_command_buffer_options_initialize(
+    iree_hal_allocator_t* device_allocator, iree_hal_command_buffer_mode_t mode,
+    iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t binding_capacity,
+    iree_hal_amdgpu_command_buffer_options_t* out_options);
+
+// Verifies command buffer options to ensure they meet the requirements of the
+// devices the command buffer will be scheduled on.
+iree_status_t iree_hal_amdgpu_command_buffer_options_verify(
+    const iree_hal_amdgpu_command_buffer_options_t* options);
+
+//===----------------------------------------------------------------------===//
+// iree_hal_amdgpu_command_buffer_t
+//===----------------------------------------------------------------------===//
+
+// Creates an AMDGPU command buffer with the given |options| controlling how
+// it is recorded and prepared for execution.
+//
+// Referenced data structures in the options such as block pools must remain
+// live for the lifetime of the command buffer.
+iree_status_t iree_hal_amdgpu_command_buffer_create(
+    const iree_hal_amdgpu_command_buffer_options_t* options,
+    iree_allocator_t host_allocator,
+    iree_hal_command_buffer_t** out_command_buffer);
+
+// Returns true if |command_buffer| is a AMDGPU command buffer.
+bool iree_hal_amdgpu_command_buffer_isa(
+    iree_hal_command_buffer_t* command_buffer);
+
+// Queries the device-side command buffer representation for the GPU device
+// agent with |device_ordinal| in the system topology.
+// |out_max_kernarg_capacity| will be set to the minimum required kernarg
+// reservation used by any block in the command buffer.
+iree_status_t iree_hal_amdgpu_command_buffer_query_execution_state(
+    iree_hal_command_buffer_t* command_buffer, iree_host_size_t device_ordinal,
+    IREE_AMDGPU_DEVICE_PTR iree_hal_amdgpu_device_command_buffer_t**
+        out_device_command_buffer,
+    iree_host_size_t* out_max_kernarg_capacity);
+
+#endif  // IREE_HAL_DRIVERS_AMDGPU_COMMAND_BUFFER_H_

runtime/src/iree/hal/drivers/amdgpu/command_buffer.h

@@ -11,13 +11,15 @@
 set(_BITCODE_SRCS
   "blit.c"
   "buffer.c"
+  "command_buffer.c"
   "host_client.c"
   "tracing.c"
 )
 
 set(_BITCODE_HDRS
   "blit.h"
   "buffer.h"
+  "command_buffer.h"
   "kernel_tables.h"
   "kernels.h"
   "host_client.h"

runtime/src/iree/hal/drivers/amdgpu/device/CMakeLists.txt

@@ -63,3 +63,26 @@ void* iree_hal_amdgpu_device_workgroup_count_buffer_ref_resolve(
              ? (uint8_t*)buffer_ref.value.ptr + buffer_ref.offset
              : NULL;
 }
+
+void* iree_hal_amdgpu_device_uint64_buffer_ref_resolve(
+    iree_hal_amdgpu_device_uint64_buffer_ref_t buffer_ref,
+    IREE_AMDGPU_ALIGNAS(64)
+        const iree_hal_amdgpu_device_buffer_ref_t* IREE_AMDGPU_RESTRICT
+            binding_table) {
+  if (buffer_ref.type == IREE_HAL_AMDGPU_DEVICE_BUFFER_TYPE_SLOT) {
+    const iree_hal_amdgpu_device_buffer_ref_t binding =
+        binding_table[buffer_ref.value.slot];
+    const uint64_t offset = buffer_ref.offset + binding.offset;
+    buffer_ref = (iree_hal_amdgpu_device_uint64_buffer_ref_t){
+        .type = binding.type,
+        .offset = offset,
+        .value.bits = binding.value.bits,
+    };
+  }
+  if (buffer_ref.type == IREE_HAL_AMDGPU_DEVICE_BUFFER_TYPE_HANDLE) {
+    buffer_ref.value.ptr = buffer_ref.value.handle->ptr;
+  }
+  return buffer_ref.value.ptr
+             ? (uint8_t*)buffer_ref.value.ptr + buffer_ref.offset
+             : NULL;
+}

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

@@ -131,6 +131,34 @@ static_assert(sizeof(iree_hal_amdgpu_device_workgroup_count_buffer_ref_t) == 16,
 #define iree_hal_amdgpu_device_workgroup_count_buffer_ref_length(buffer_ref) \
   (sizeof(uint32_t) * 3)
 
+// Describes a buffer binding that contains a single uint64_t value.
+// This is a size-optimized version of iree_hal_amdgpu_device_buffer_ref_t so
+// that it will fit in our tiny packets. We know the length is a constant 8 and
+// only need the offset, type, and value.
+typedef struct iree_hal_amdgpu_device_uint64_buffer_ref_t {
+  // Type of the buffer reference used to resolve the device pointer.
+  uint64_t type : 2;  // iree_hal_amdgpu_device_buffer_type_t
+  // Offset, in bytes, into the buffer that the binding starts at.
+  // This will be added to the offset specified on each usage of the slot.
+  uint64_t offset : 62;
+  union {
+    // IREE_HAL_AMDGPU_DEVICE_BUFFER_TYPE_PTR: raw device pointer.
+    void* ptr;
+    // IREE_HAL_AMDGPU_DEVICE_BUFFER_TYPE_HANDLE: queue-ordered allocation
+    // handle.
+    iree_hal_amdgpu_device_allocation_handle_t* handle;
+    // IREE_HAL_AMDGPU_DEVICE_BUFFER_TYPE_SLOT: binding table slot.
+    iree_hal_amdgpu_device_buffer_ordinal_t slot;
+    // Used for setting the value.
+    uint64_t bits;
+  } value;
+} iree_hal_amdgpu_device_uint64_buffer_ref_t;
+static_assert(sizeof(iree_hal_amdgpu_device_uint64_buffer_ref_t) == 16,
+              "binding table entries should be 8 byte aligned and tiny");
+
+#define iree_hal_amdgpu_device_uint64_buffer_ref_length(buffer_ref) \
+  sizeof(uint64_t)
+
 #if defined(IREE_AMDGPU_TARGET_DEVICE)
 
 // Resolves a buffer reference to an absolute device pointer.
@@ -155,6 +183,16 @@ void* iree_hal_amdgpu_device_workgroup_count_buffer_ref_resolve(
         const iree_hal_amdgpu_device_buffer_ref_t* IREE_AMDGPU_RESTRICT
             binding_table);
 
+// Resolves a scalar uint64_t buffer reference to an absolute device pointer.
+// This is equivalent to iree_hal_amdgpu_device_buffer_ref_resolve but for a
+// fixed-size uint64_t value. The returned pointer should have 8-byte
+// alignment.
+void* iree_hal_amdgpu_device_uint64_buffer_ref_resolve(
+    iree_hal_amdgpu_device_uint64_buffer_ref_t buffer_ref,
+    IREE_AMDGPU_ALIGNAS(64)
+        const iree_hal_amdgpu_device_buffer_ref_t* IREE_AMDGPU_RESTRICT
+            binding_table);
+
 #endif  // IREE_AMDGPU_TARGET_DEVICE
 
 #endif  // IREE_HAL_DRIVERS_AMDGPU_DEVICE_BUFFER_H_