Adding iree_hal_device_queue_update.

As with all queue DMA operations it's best if things are batched into
command buffers but it's bad to have a command buffer with a single
DMA operation - this completes the set of fill/update/copy operations
at the queue level to match the command buffer DMA operations.
Practically this is useful when combined with reusable/indirect command
buffers for uploading new parameters in queue order prior to issuing
a command buffer that references them. The compiler will use this to
turn push constants into uniform buffers.

An emulated version is added but implementations are encouraged to do
better... they currently don't.

Author: benvanik

experimental/webgpu/webgpu_device.c

@@ -470,6 +470,7 @@ const iree_hal_device_vtable_t iree_hal_webgpu_device_vtable = {
     .queue_alloca = iree_hal_webgpu_device_queue_alloca,
     .queue_dealloca = iree_hal_webgpu_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_webgpu_device_queue_read,
     .queue_write = iree_hal_webgpu_device_queue_write,

runtime/src/iree/hal/command_buffer.c

@@ -662,6 +662,15 @@ IREE_API_EXPORT iree_status_t iree_hal_create_transfer_command_buffer(
               transfer_command->fill.pattern,
               transfer_command->fill.pattern_length, IREE_HAL_FILL_FLAG_NONE);
           break;
+        case IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE:
+          status = iree_hal_command_buffer_update_buffer(
+              command_buffer, transfer_command->update.source_buffer,
+              transfer_command->update.source_offset,
+              iree_hal_make_buffer_ref(transfer_command->update.target_buffer,
+                                       transfer_command->update.target_offset,
+                                       transfer_command->update.length),
+              IREE_HAL_UPDATE_FLAG_NONE);
+          break;
         case IREE_HAL_TRANSFER_COMMAND_TYPE_COPY:
           status = iree_hal_command_buffer_copy_buffer(
               command_buffer,
@@ -673,15 +682,6 @@ IREE_API_EXPORT iree_status_t iree_hal_create_transfer_command_buffer(
                                        transfer_command->copy.length),
               IREE_HAL_COPY_FLAG_NONE);
           break;
-        case IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE:
-          status = iree_hal_command_buffer_update_buffer(
-              command_buffer, transfer_command->update.source_buffer,
-              transfer_command->update.source_offset,
-              iree_hal_make_buffer_ref(transfer_command->update.target_buffer,
-                                       transfer_command->update.target_offset,
-                                       transfer_command->update.length),
-              IREE_HAL_UPDATE_FLAG_NONE);
-          break;
         default:
           status =
               iree_make_status(IREE_STATUS_INVALID_ARGUMENT,

runtime/src/iree/hal/command_buffer.h

@@ -230,7 +230,7 @@ enum iree_hal_fill_flag_bits_t {
   IREE_HAL_FILL_FLAG_NONE = 0,
 };
 
-// Bitfield specifying flags controlling a update operation.
+// Bitfield specifying flags controlling an update operation.
 typedef uint64_t iree_hal_update_flags_t;
 enum iree_hal_update_flag_bits_t {
   IREE_HAL_UPDATE_FLAG_NONE = 0,
@@ -802,10 +802,10 @@ IREE_API_EXPORT iree_status_t iree_hal_command_buffer_validate_submission(
 typedef enum iree_hal_transfer_command_type_t {
   // iree_hal_command_buffer_fill_buffer
   IREE_HAL_TRANSFER_COMMAND_TYPE_FILL = 0u,
-  // iree_hal_command_buffer_copy_buffer
-  IREE_HAL_TRANSFER_COMMAND_TYPE_COPY = 1u,
   // iree_hal_command_buffer_update_buffer
-  IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE = 2u,
+  IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE = 1u,
+  // iree_hal_command_buffer_copy_buffer
+  IREE_HAL_TRANSFER_COMMAND_TYPE_COPY = 2u,
 } iree_hal_transfer_command_type_t;
 
 // Represents a single transfer command within a batch of commands.
@@ -821,14 +821,6 @@ typedef struct iree_hal_transfer_command_t {
       const void* pattern;
       iree_host_size_t pattern_length;
     } fill;
-    // IREE_HAL_TRANSFER_COMMAND_TYPE_COPY
-    struct {
-      iree_hal_buffer_t* source_buffer;
-      iree_device_size_t source_offset;
-      iree_hal_buffer_t* target_buffer;
-      iree_device_size_t target_offset;
-      iree_device_size_t length;
-    } copy;
     // IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE
     struct {
       const void* source_buffer;
@@ -837,6 +829,14 @@ typedef struct iree_hal_transfer_command_t {
       iree_device_size_t target_offset;
       iree_device_size_t length;
     } update;
+    // IREE_HAL_TRANSFER_COMMAND_TYPE_COPY
+    struct {
+      iree_hal_buffer_t* source_buffer;
+      iree_device_size_t source_offset;
+      iree_hal_buffer_t* target_buffer;
+      iree_device_size_t target_offset;
+      iree_device_size_t length;
+    } copy;
   };
 } iree_hal_transfer_command_t;
 

runtime/src/iree/hal/device.c

@@ -196,6 +196,91 @@ IREE_API_EXPORT iree_status_t iree_hal_device_queue_fill(
   return status;
 }
 
+IREE_API_EXPORT iree_status_t iree_hal_device_queue_emulated_update(
+    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
+    const iree_hal_semaphore_list_t wait_semaphore_list,
+    const iree_hal_semaphore_list_t signal_semaphore_list,
+    const void* source_buffer, iree_host_size_t source_offset,
+    iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
+    iree_device_size_t length, iree_hal_update_flags_t flags) {
+  IREE_ASSERT_ARGUMENT(device);
+  IREE_ASSERT_ARGUMENT(source_buffer);
+  IREE_ASSERT_ARGUMENT(target_buffer);
+  IREE_TRACE_ZONE_BEGIN(z0);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)length);
+
+  // If we are starting execution immediately then we can reduce latency by
+  // allowing inline command buffer execution.
+  iree_hal_command_buffer_mode_t command_buffer_mode =
+      IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT;
+  if (wait_semaphore_list.count == 0) {
+    command_buffer_mode |= IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION;
+  }
+
+  // TODO(benvanik): support splitting the update into multiple chunks to fit
+  // under the max command buffer update size limit. This provisional API is
+  // intended only for updating dispatch parameters today.
+  if (length > UINT16_MAX) {
+    return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+                            "queue buffer updates currently limited to 64KB, "
+                            "tried to update %" PRIhsz " bytes",
+                            length);
+  }
+
+  iree_hal_transfer_command_t command = {
+      .type = IREE_HAL_TRANSFER_COMMAND_TYPE_UPDATE,
+      .update =
+          {
+              .source_buffer = source_buffer,
+              .source_offset = source_offset,
+              .target_buffer = target_buffer,
+              .target_offset = target_offset,
+              .length = length,
+          },
+  };
+
+  iree_hal_command_buffer_t* command_buffer = NULL;
+  IREE_RETURN_AND_END_ZONE_IF_ERROR(
+      z0, iree_hal_create_transfer_command_buffer(device, command_buffer_mode,
+                                                  queue_affinity, 1, &command,
+                                                  &command_buffer));
+
+  iree_status_t status = iree_hal_device_queue_execute(
+      device, queue_affinity, wait_semaphore_list, signal_semaphore_list,
+      command_buffer, iree_hal_buffer_binding_table_empty());
+
+  iree_hal_command_buffer_release(command_buffer);
+
+  IREE_TRACE_ZONE_END(z0);
+  return status;
+}
+
+IREE_API_EXPORT iree_status_t iree_hal_device_queue_update(
+    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
+    const iree_hal_semaphore_list_t wait_semaphore_list,
+    const iree_hal_semaphore_list_t signal_semaphore_list,
+    const void* source_buffer, iree_host_size_t source_offset,
+    iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
+    iree_device_size_t length, iree_hal_update_flags_t flags) {
+  IREE_ASSERT_ARGUMENT(device);
+  IREE_ASSERT_ARGUMENT(
+      !wait_semaphore_list.count ||
+      (wait_semaphore_list.semaphores && wait_semaphore_list.payload_values));
+  IREE_ASSERT_ARGUMENT(!signal_semaphore_list.count ||
+                       (signal_semaphore_list.semaphores &&
+                        signal_semaphore_list.payload_values));
+  IREE_ASSERT_ARGUMENT(source_buffer);
+  IREE_ASSERT_ARGUMENT(target_buffer);
+  IREE_TRACE_ZONE_BEGIN(z0);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)length);
+  iree_status_t status = _VTABLE_DISPATCH(device, queue_update)(
+      device, queue_affinity, wait_semaphore_list, signal_semaphore_list,
+      source_buffer, source_offset, target_buffer, target_offset, length,
+      flags);
+  IREE_TRACE_ZONE_END(z0);
+  return status;
+}
+
 IREE_API_EXPORT iree_status_t iree_hal_device_queue_emulated_copy(
     iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
     const iree_hal_semaphore_list_t wait_semaphore_list,

runtime/src/iree/hal/device.h

@@ -312,11 +312,12 @@ IREE_API_EXPORT iree_status_t iree_hal_device_queue_dealloca(
     iree_hal_buffer_t* buffer);
 
 // Enqueues a single queue-ordered fill operation.
+// The |target_buffer| must be visible to the device queue performing the fill.
 //
 // WARNING: individual fills have a high overhead and batching should be
 // performed by the caller instead of calling this multiple times. The
 // iree_hal_create_transfer_command_buffer utility makes it easy to create
-// batches of transfer operations (fill, copy, update) and is only a few lines
+// batches of transfer operations (fill, update, copy) and is only a few lines
 // more code.
 IREE_API_EXPORT iree_status_t iree_hal_device_queue_fill(
     iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
@@ -326,12 +327,36 @@ IREE_API_EXPORT iree_status_t iree_hal_device_queue_fill(
     iree_device_size_t length, const void* pattern,
     iree_host_size_t pattern_length, iree_hal_fill_flags_t flags);
 
+// Enqueues a single queue-ordered buffer update operation.
+// The provided |source_buffer| will be captured and need not remain live or
+// unchanged while the operation is queued. The |target_buffer| must be visible
+// to the device queue performing the update.
+//
+// Some implementations may have limits on the size of the update or may perform
+// poorly if the size is larger than an implementation-defined limit. Updates
+// should be kept as small and infrequent as possible.
+//
+// WARNING: individual copies have a high overhead and batching should be
+// performed by the caller instead of calling this multiple times. The
+// iree_hal_create_transfer_command_buffer utility makes it easy to create
+// batches of transfer operations (fill, update, copy) and is only a few lines
+// more code.
+IREE_API_EXPORT iree_status_t iree_hal_device_queue_update(
+    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
+    const iree_hal_semaphore_list_t wait_semaphore_list,
+    const iree_hal_semaphore_list_t signal_semaphore_list,
+    const void* source_buffer, iree_host_size_t source_offset,
+    iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
+    iree_device_size_t length, iree_hal_update_flags_t flags);
+
 // Enqueues a single queue-ordered copy operation.
+// The |source_buffer| and |target_buffer| must both be visible to the device
+// queue performing the copy.
 //
 // WARNING: individual copies have a high overhead and batching should be
 // performed by the caller instead of calling this multiple times. The
 // iree_hal_create_transfer_command_buffer utility makes it easy to create
-// batches of transfer operations (fill, copy, update) and is only a few lines
+// batches of transfer operations (fill, update, copy) and is only a few lines
 // more code.
 IREE_API_EXPORT iree_status_t iree_hal_device_queue_copy(
     iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
@@ -578,6 +603,14 @@ typedef struct iree_hal_device_vtable_t {
       iree_device_size_t length, const void* pattern,
       iree_host_size_t pattern_length, iree_hal_fill_flags_t flags);
 
+  iree_status_t(IREE_API_PTR* queue_update)(
+      iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
+      const iree_hal_semaphore_list_t wait_semaphore_list,
+      const iree_hal_semaphore_list_t signal_semaphore_list,
+      const void* source_buffer, iree_host_size_t source_offset,
+      iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
+      iree_device_size_t length, iree_hal_update_flags_t flags);
+
   iree_status_t(IREE_API_PTR* queue_copy)(
       iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
       const iree_hal_semaphore_list_t wait_semaphore_list,
@@ -634,6 +667,14 @@ IREE_API_EXPORT iree_status_t iree_hal_device_queue_emulated_fill(
     iree_device_size_t length, const void* pattern,
     iree_host_size_t pattern_length, iree_hal_fill_flags_t flags);
 
+IREE_API_EXPORT iree_status_t iree_hal_device_queue_emulated_update(
+    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
+    const iree_hal_semaphore_list_t wait_semaphore_list,
+    const iree_hal_semaphore_list_t signal_semaphore_list,
+    const void* source_buffer, iree_host_size_t source_offset,
+    iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
+    iree_device_size_t length, iree_hal_update_flags_t flags);
+
 IREE_API_EXPORT iree_status_t iree_hal_device_queue_emulated_copy(
     iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
     const iree_hal_semaphore_list_t wait_semaphore_list,

runtime/src/iree/hal/drivers/cuda/cuda_device.c

@@ -1130,6 +1130,7 @@ static const iree_hal_device_vtable_t iree_hal_cuda_device_vtable = {
     .queue_alloca = iree_hal_cuda_device_queue_alloca,
     .queue_dealloca = iree_hal_cuda_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_cuda_device_queue_read,
     .queue_write = iree_hal_cuda_device_queue_write,

runtime/src/iree/hal/drivers/hip/hip_device.c

@@ -1127,6 +1127,7 @@ static const iree_hal_device_vtable_t iree_hal_hip_device_vtable = {
     .queue_alloca = iree_hal_hip_device_queue_alloca,
     .queue_dealloca = iree_hal_hip_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_hip_device_queue_read,
     .queue_write = iree_hal_hip_device_queue_write,

runtime/src/iree/hal/drivers/local_sync/sync_device.c

@@ -504,6 +504,7 @@ static const iree_hal_device_vtable_t iree_hal_sync_device_vtable = {
     .queue_alloca = iree_hal_sync_device_queue_alloca,
     .queue_dealloca = iree_hal_sync_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_sync_device_queue_read,
     .queue_write = iree_hal_sync_device_queue_write,

runtime/src/iree/hal/drivers/local_task/task_device.c

@@ -539,6 +539,7 @@ static const iree_hal_device_vtable_t iree_hal_task_device_vtable = {
     .queue_alloca = iree_hal_task_device_queue_alloca,
     .queue_dealloca = iree_hal_task_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_task_device_queue_read,
     .queue_write = iree_hal_task_device_queue_write,

runtime/src/iree/hal/drivers/metal/metal_device.m

@@ -618,6 +618,7 @@ static iree_status_t iree_hal_metal_device_profiling_end(iree_hal_device_t* base
     .queue_alloca = iree_hal_metal_device_queue_alloca,
     .queue_dealloca = iree_hal_metal_device_queue_dealloca,
     .queue_fill = iree_hal_device_queue_emulated_fill,
+    .queue_update = iree_hal_device_queue_emulated_update,
     .queue_copy = iree_hal_device_queue_emulated_copy,
     .queue_read = iree_hal_metal_device_queue_read,
     .queue_write = iree_hal_metal_device_queue_write,