Adding iree_hal_device_queue_update and improving queue DMA operations. (#19000)

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.

While updating the queue API I've added placeholder flags to all DMA
operations in preparation for compiler updates that will provide them.
`iree_hal_device_queue_execute` has needed simplification for awhile and
that's done here to allow implementations to not need to worry with
batched command buffer juggling. The unused-since-its-inception
`iree_hal_command_buffer_discard_buffer` API has been renamed to
`iree_hal_command_buffer_advise_buffer` ahead of compiler changes that
will use it for multi-device cache management.

No breaking changes to the compiler here - future PRs will update the
HAL module and ops.

Author: benvanik

experimental/web/sample_webgpu/main.c

@@ -794,8 +794,8 @@ static iree_status_t process_call_outputs(
     };
     status = iree_hal_device_queue_execute(
         device, IREE_HAL_QUEUE_AFFINITY_ANY, iree_hal_semaphore_list_empty(),
-        signal_semaphores, 1, &transfer_command_buffer,
-        /*binding_tables=*/NULL);
+        signal_semaphores, transfer_command_buffer,
+        iree_hal_buffer_binding_table_empty());
   }
   // TODO(scotttodd): Make this async - pass a wait source to iree_loop_wait_one
   //     1. create iree_hal_fence_t, iree_hal_fence_insert(fance, semaphore)

experimental/webgpu/command_buffer.c

@@ -575,9 +575,10 @@ static iree_status_t iree_hal_webgpu_command_buffer_wait_events(
   return iree_ok_status();
 }
 
-static iree_status_t iree_hal_webgpu_command_buffer_discard_buffer(
+static iree_status_t iree_hal_webgpu_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // No-op: though maybe it'd be a useful addition to the spec as otherwise
   // false dependencies can creep in.
   return iree_ok_status();
@@ -608,7 +609,7 @@ static uint32_t iree_hal_webgpu_splat_pattern(const void* pattern,
 static iree_status_t iree_hal_webgpu_command_buffer_fill_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
     iree_hal_buffer_ref_t target_ref, const void* pattern,
-    iree_host_size_t pattern_length) {
+    iree_host_size_t pattern_length, iree_hal_fill_flags_t flags) {
   iree_hal_webgpu_command_buffer_t* command_buffer =
       iree_hal_webgpu_command_buffer_cast(base_command_buffer);
 
@@ -693,7 +694,8 @@ static iree_status_t iree_hal_webgpu_command_buffer_fill_buffer(
 
 static iree_status_t iree_hal_webgpu_command_buffer_update_buffer(
     iree_hal_command_buffer_t* base_command_buffer, const void* source_buffer,
-    iree_host_size_t source_offset, iree_hal_buffer_ref_t target_ref) {
+    iree_host_size_t source_offset, iree_hal_buffer_ref_t target_ref,
+    iree_hal_update_flags_t flags) {
   iree_hal_webgpu_command_buffer_t* command_buffer =
       iree_hal_webgpu_command_buffer_cast(base_command_buffer);
 
@@ -734,7 +736,8 @@ static iree_status_t iree_hal_webgpu_command_buffer_update_buffer(
 
 static iree_status_t iree_hal_webgpu_command_buffer_copy_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t source_ref, iree_hal_buffer_ref_t target_ref) {
+    iree_hal_buffer_ref_t source_ref, iree_hal_buffer_ref_t target_ref,
+    iree_hal_copy_flags_t flags) {
   iree_hal_webgpu_command_buffer_t* command_buffer =
       iree_hal_webgpu_command_buffer_cast(base_command_buffer);
 
@@ -1041,7 +1044,7 @@ const iree_hal_command_buffer_vtable_t iree_hal_webgpu_command_buffer_vtable = {
     .signal_event = iree_hal_webgpu_command_buffer_signal_event,
     .reset_event = iree_hal_webgpu_command_buffer_reset_event,
     .wait_events = iree_hal_webgpu_command_buffer_wait_events,
-    .discard_buffer = iree_hal_webgpu_command_buffer_discard_buffer,
+    .advise_buffer = iree_hal_webgpu_command_buffer_advise_buffer,
     .fill_buffer = iree_hal_webgpu_command_buffer_fill_buffer,
     .update_buffer = iree_hal_webgpu_command_buffer_update_buffer,
     .copy_buffer = iree_hal_webgpu_command_buffer_copy_buffer,

experimental/webgpu/webgpu_device.c

@@ -354,7 +354,7 @@ static iree_status_t iree_hal_webgpu_device_queue_read(
     const iree_hal_semaphore_list_t signal_semaphore_list,
     iree_hal_file_t* source_file, uint64_t source_offset,
     iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
-    iree_device_size_t length, uint32_t flags) {
+    iree_device_size_t length, iree_hal_read_flags_t flags) {
   // TODO: expose streaming chunk count/size options.
   iree_status_t loop_status = iree_ok_status();
   iree_hal_file_transfer_options_t options = {
@@ -376,7 +376,7 @@ static iree_status_t iree_hal_webgpu_device_queue_write(
     const iree_hal_semaphore_list_t signal_semaphore_list,
     iree_hal_buffer_t* source_buffer, iree_device_size_t source_offset,
     iree_hal_file_t* target_file, uint64_t target_offset,
-    iree_device_size_t length, uint32_t flags) {
+    iree_device_size_t length, iree_hal_write_flags_t flags) {
   // TODO: expose streaming chunk count/size options.
   iree_status_t loop_status = iree_ok_status();
   iree_hal_file_transfer_options_t options = {
@@ -396,9 +396,8 @@ static iree_status_t iree_hal_webgpu_device_queue_execute(
     iree_hal_device_t* base_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,
-    iree_host_size_t command_buffer_count,
-    iree_hal_command_buffer_t* const* command_buffers,
-    iree_hal_buffer_binding_table_t const* binding_tables) {
+    iree_hal_command_buffer_t* command_buffer,
+    iree_hal_buffer_binding_table_t binding_table) {
   iree_hal_webgpu_device_t* device = iree_hal_webgpu_device_cast(base_device);
 
   // TODO(benvanik): this currently assumes we are synchronizing on semaphores
@@ -410,11 +409,8 @@ static iree_status_t iree_hal_webgpu_device_queue_execute(
                                                     iree_infinite_timeout()));
 
   // TODO(benvanik): propagate errors to semaphores.
-  for (iree_host_size_t i = 0; i < command_buffer_count; i++) {
-    iree_hal_command_buffer_t* command_buffer = command_buffers[i];
-    IREE_RETURN_IF_ERROR(
-        iree_hal_webgpu_command_buffer_issue(command_buffer, device->queue));
-  }
+  IREE_RETURN_IF_ERROR(
+      iree_hal_webgpu_command_buffer_issue(command_buffer, device->queue));
 
   IREE_RETURN_IF_ERROR(iree_hal_semaphore_list_signal(signal_semaphore_list));
 
@@ -473,6 +469,9 @@ const iree_hal_device_vtable_t iree_hal_webgpu_device_vtable = {
         iree_hal_webgpu_device_query_semaphore_compatibility,
     .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,
     .queue_execute = iree_hal_webgpu_device_queue_execute,

integrations/pjrt/src/iree_pjrt/common/api_impl.cc

@@ -590,8 +590,8 @@ iree_status_t BufferInstance::CopyToHost(void* dst, iree_host_size_t dst_size,
       device_.device(), IREE_HAL_QUEUE_AFFINITY_ANY,
       /*wait_semaphore_list=*/iree_hal_fence_semaphore_list(ready_fence_.get()),
       /*signal_semaphore_list=*/
-      iree_hal_fence_semaphore_list(dst_buffer_ready_fence.get()),
-      /*command_buffer_count=*/1, &transfer_cb, NULL));
+      iree_hal_fence_semaphore_list(dst_buffer_ready_fence.get()), transfer_cb,
+      iree_hal_buffer_binding_table_empty()));
 
   *out_done_event = copy_done_event;
   return iree_ok_status();
@@ -837,7 +837,8 @@ iree_status_t DeviceInstance::HostBufferToDeviceSplat(
   IREE_CHECK_OK(iree_hal_command_buffer_begin(transfer_cb.get()));
   IREE_RETURN_IF_ERROR(iree_hal_command_buffer_fill_buffer(
       transfer_cb.get(), buffer.get(), /*target_offset=*/0,
-      /*target_size=*/byte_length, data, element_type_byte_size));
+      /*target_size=*/byte_length, data, element_type_byte_size,
+      IREE_HAL_FILL_FLAG_NONE));
   IREE_CHECK_OK(iree_hal_command_buffer_end(transfer_cb.get()));
 
   // Execute the enqueued splat:
@@ -846,8 +847,8 @@ iree_status_t DeviceInstance::HostBufferToDeviceSplat(
       /*wait_semaphore_list=*/
       {1, &transfer_timeline_, &signal_alloca_complete},
       /*signal_semaphore_list=*/
-      {1, &transfer_timeline_, &signal_copy_complete},
-      /*command_buffer_count=*/1, &transfer_cb, NULL));
+      {1, &transfer_timeline_, &signal_copy_complete}, transfer_cb,
+      iree_hal_buffer_binding_table_empty()));
 
   // Wrap in a buffer view and return:
   iree::vm::ref<iree_hal_buffer_view_t> result_buffer_view;
@@ -1190,8 +1191,8 @@ iree_status_t DeviceInstance::HostBufferToDevice(
       /*wait_semaphore_list=*/
       {1, &transfer_timeline_, &signal_alloca_complete},
       /*signal_semaphore_list=*/
-      {1, &transfer_timeline_, &signal_copy_complete},
-      /*command_buffer_count=*/1, &transfer_cb, NULL));
+      {1, &transfer_timeline_, &signal_copy_complete}, transfer_cb,
+      iree_hal_buffer_binding_table_empty()));
 
   // Wrap in a buffer view and return.
   iree::vm::ref<iree_hal_buffer_view_t> result_buffer_view;

integrations/pjrt/src/iree_pjrt/common/iree_helpers.h

@@ -139,17 +139,16 @@ iree_status_t hal_device_queue_execute(
     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,
-    iree_host_size_t command_buffer_count,
-    iree_hal_command_buffer_t* const* command_buffers) {
+    iree_hal_command_buffer_t* command_buffer) {
   if (LOGGING_ENABLED) {
     LogInvoke(__func__, "device=%p, wait={%s}, signal={%s}", device,
               SemaphoreListToString(wait_semaphore_list).c_str(),
               SemaphoreListToString(signal_semaphore_list).c_str());
   }
   return HandleStatus(__func__, iree_hal_device_queue_execute(
                                     device, queue_affinity, wait_semaphore_list,
-                                    signal_semaphore_list, command_buffer_count,
-                                    command_buffers, /*binding_tables=*/NULL));
+                                    signal_semaphore_list, command_buffer,
+                                    iree_hal_buffer_binding_table_empty()));
 }
 
 iree_status_t hal_fence_create(iree_host_size_t capacity,

runtime/bindings/python/hal.cc

@@ -499,7 +499,7 @@ void HalDevice::QueueDealloca(HalBuffer& buffer, py::handle wait_semaphores,
       "deallocating memory on queue");
 }
 
-void HalDevice::QueueExecute(py::handle command_buffers,
+void HalDevice::QueueExecute(py::handle command_buffer,
                              py::handle wait_semaphores,
                              py::handle signal_semaphores) {
   iree_hal_semaphore_list_t wait_list;
@@ -548,17 +548,14 @@ void HalDevice::QueueExecute(py::handle command_buffers,
   }
 
   // Unpack command buffers.
-  size_t cb_count = py::len(command_buffers);
-  iree_hal_command_buffer_t** cb_list =
-      static_cast<iree_hal_command_buffer_t**>(
-          alloca(sizeof(iree_hal_command_buffer_t*) * cb_count));
-  for (size_t i = 0; i < cb_count; ++i) {
-    cb_list[i] = py::cast<HalCommandBuffer*>(command_buffers[i])->raw_ptr();
-  }
+  iree_hal_command_buffer_t* cb =
+      !command_buffer.is_none()
+          ? py::cast<HalCommandBuffer*>(command_buffer)->raw_ptr()
+          : NULL;
 
   CheckApiStatus(iree_hal_device_queue_execute(
                      raw_ptr(), IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
-                     signal_list, cb_count, cb_list, /*binding_tables=*/NULL),
+                     signal_list, cb, iree_hal_buffer_binding_table_empty()),
                  "executing command buffers");
 }
 
@@ -619,11 +616,12 @@ void HalDevice::QueueCopy(HalBuffer& source_buffer, HalBuffer& target_buffer,
         "Source and buffer length must be less than the target buffer length "
         "and it does not. Please check allocations");
   }
-  CheckApiStatus(iree_hal_device_queue_copy(
-                     raw_ptr(), IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
-                     signal_list, source_buffer.raw_ptr(), 0,
-                     target_buffer.raw_ptr(), 0, source_length),
-                 "Copying buffer on queue");
+  CheckApiStatus(
+      iree_hal_device_queue_copy(
+          raw_ptr(), IREE_HAL_QUEUE_AFFINITY_ANY, wait_list, signal_list,
+          source_buffer.raw_ptr(), 0, target_buffer.raw_ptr(), 0, source_length,
+          IREE_HAL_COPY_FLAG_NONE),
+      "Copying buffer on queue");
 }
 
 py::object HalDevice::CreateDLPackCapsule(HalBufferView& buffer_view,
@@ -1729,7 +1727,8 @@ void SetupHalBindings(nanobind::module_ m) {
                     iree_hal_make_buffer_ref(source_buffer.raw_ptr(),
                                              source_offset, resolved_length),
                     iree_hal_make_buffer_ref(target_buffer.raw_ptr(),
-                                             target_offset, resolved_length)),
+                                             target_offset, resolved_length),
+                    IREE_HAL_COPY_FLAG_NONE),
                 "copy command");
             if (end) {
               CheckApiStatus(iree_hal_command_buffer_end(self.raw_ptr()),
@@ -1767,7 +1766,8 @@ void SetupHalBindings(nanobind::module_ m) {
                     self.raw_ptr(),
                     iree_hal_make_buffer_ref(target_buffer.raw_ptr(),
                                              target_offset, resolved_length),
-                    pattern_view.buf, pattern_view.len),
+                    pattern_view.buf, pattern_view.len,
+                    IREE_HAL_FILL_FLAG_NONE),
                 "command buffer fill");
             if (end) {
               CheckApiStatus(iree_hal_command_buffer_end(self.raw_ptr()),

runtime/bindings/python/iree/runtime/_binding.pyi

@@ -185,7 +185,7 @@ class HalDevice:
     ) -> None: ...
     def queue_execute(
         self,
-        command_buffers: Sequence[HalCommandBuffer],
+        command_buffer: HalCommandBuffer,
         wait_semaphores: HalSemaphoreList,
         signal_semaphores: HalSemaphoreList,
     ) -> None: ...

runtime/bindings/python/tests/hal_test.py

@@ -463,7 +463,7 @@ def testCommandBufferExecute(self):
 
         sem = self.device.create_semaphore(0)
         self.device.queue_execute(
-            [cb], wait_semaphores=[(sem, 0)], signal_semaphores=[(sem, 1)]
+            cb, wait_semaphores=[(sem, 0)], signal_semaphores=[(sem, 1)]
         )
         iree.runtime.HalFence.create_at(sem, 1).wait()
 
@@ -479,7 +479,7 @@ def testCommandBufferExecuteAcceptsFence(self):
 
         sem = self.device.create_semaphore(0)
         self.device.queue_execute(
-            [cb],
+            cb,
             wait_semaphores=iree.runtime.HalFence.create_at(sem, 0),
             signal_semaphores=iree.runtime.HalFence.create_at(sem, 1),
         )

runtime/src/iree/hal/buffer_transfer.c

@@ -78,7 +78,7 @@ static iree_status_t iree_hal_device_transfer_and_wait(
     };
     status = iree_hal_device_queue_execute(
         device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_semaphores, signal_semaphores,
-        1, &command_buffer, /*binding_tables=*/NULL);
+        command_buffer, iree_hal_buffer_binding_table_empty());
   }
   if (iree_status_is_ok(status)) {
     status = iree_hal_semaphore_wait(fence_semaphore, signal_value, timeout);