Unifying dispatch/dispatch_indirect and adding extended configuration. (iree-org#21627)

`iree_hal_dispatch_config_t` now defines the dispatch configuration
(workgroup size, workgroup count, dynamic workgroup local memory, and
anything we want to add in the future). This avoids the extra vtable
entry plumbed through the API for indirect dispatch and will let us add
new dispatch behavior in the future without new APIs. Implementations
can still switch internally based on the flags if they need to.

A new behavior for custom/indirect arguments is defined but not yet
implemented anywhere. This bypasses the IREE HAL ABI and allows users to
directly pass in dispatch arguments either directly inline in the
command buffer (via the constants field) or indirectly via
`bindings[0]`. This can be used to dispatch existing kernels that have
their own ABI (HIP/CUDA/etc). The indirect variant allows for custom
arguments to be provided per submission as the `bindings[0]` buffer can
reference a slot
in the binding table.

Fixes iree-org#21620.
Fixes iree-org#21621.

Signed-off-by: keshavvinayak01 <[email protected]>

Author: benvanik

compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td

@@ -77,7 +77,7 @@ def HAL_BufferUsage_None                         : I32BitEnumAttrCase<"None",
 def HAL_BufferUsage_TransferSource               : I32BitEnumAttrCase<"TransferSource",               0x00000001>;
 def HAL_BufferUsage_TransferTarget               : I32BitEnumAttrCase<"TransferTarget",               0x00000002>;
 def HAL_BufferUsage_Transfer                     : I32BitEnumAttrCase<"Transfer",                     0x00000003>;
-def HAL_BufferUsage_DispatchIndirectParams       : I32BitEnumAttrCase<"DispatchIndirectParams",       0x00000100>;
+def HAL_BufferUsage_DispatchIndirectParameters   : I32BitEnumAttrCase<"DispatchIndirectParameters",   0x00000100>;
 def HAL_BufferUsage_DispatchUniformRead          : I32BitEnumAttrCase<"DispatchUniformRead",          0x00000200>;
 def HAL_BufferUsage_DispatchStorageRead          : I32BitEnumAttrCase<"DispatchStorageRead",          0x00000400>;
 def HAL_BufferUsage_DispatchStorageWrite         : I32BitEnumAttrCase<"DispatchStorageWrite",         0x00000800>;
@@ -101,7 +101,7 @@ def HAL_BufferUsageBitfieldAttr :
       HAL_BufferUsage_TransferSource,
       HAL_BufferUsage_TransferTarget,
       HAL_BufferUsage_Transfer,
-      HAL_BufferUsage_DispatchIndirectParams,
+      HAL_BufferUsage_DispatchIndirectParameters,
       HAL_BufferUsage_DispatchUniformRead,
       HAL_BufferUsage_DispatchStorageRead,
       HAL_BufferUsage_DispatchStorageWrite,

experimental/webgpu/command_buffer.c

@@ -806,123 +806,6 @@ static iree_status_t iree_hal_webgpu_command_buffer_push_descriptor_set(
   return iree_ok_status();
 }
 
-static iree_status_t iree_hal_webgpu_command_buffer_prepare_dispatch(
-    iree_hal_webgpu_command_buffer_t* command_buffer,
-    iree_hal_executable_t* executable, uint32_t ordinal,
-    iree_const_byte_span_t constants, iree_hal_buffer_ref_list_t bindings,
-    iree_hal_dispatch_flags_t flags, WGPUComputePassEncoder* out_compute_pass) {
-  const iree_hal_webgpu_entry_point_t* entry_point =
-      iree_hal_webgpu_executable_lookup_entry_point(executable, ordinal);
-
-  // Upload push constant data - this may incur a segment flush if the staging
-  // buffer is exhausted.
-  iree_host_size_t constant_count =
-      iree_hal_webgpu_pipeline_layout_constant_count(entry_point->layout);
-  iree_const_byte_span_t constant_data = iree_make_const_byte_span(
-      command_buffer->state.constants,
-      constant_count * sizeof(command_buffer->state.constants[0]));
-  uint32_t params_offset = 0;
-  IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_append_parameters(
-      command_buffer, constant_data, &params_offset));
-
-  // Acquire the compute pass we'll encode the dispatch into - this may be
-  // fresh or reused from prior commands.
-  WGPUComputePassEncoder compute_pass = NULL;
-  IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_acquire_compute_pass(
-      command_buffer, &compute_pass));
-  wgpuComputePassEncoderSetPipeline(compute_pass, entry_point->pipeline);
-
-  if (constant_count > 0) {
-    // Bind the push constant emulation bind group at the staging buffer
-    // relative offset for this dispatch.
-    wgpuComputePassEncoderSetBindGroup(
-        compute_pass, IREE_HAL_WEBGPU_PARAMS_BIND_GROUP_INDEX,
-        command_buffer->staging_buffer->bind_group, 1, &params_offset);
-  }
-
-  // Set all bindings.
-  const iree_hal_webgpu_set_binding_info_t* binding_info =
-      iree_hal_webgpu_pipeline_layout_set_binding_info(entry_point->layout);
-  for (iree_host_size_t i = 0; i < binding_info->set_count; ++i) {
-    // If there are no bindings in this set we can skip it.
-    if (binding_info->set_masks[i] == 0) continue;
-
-    // If there is a bind group handle then it means we've done the lookup and
-    // set the bind group on the device already - we can skip.
-    if (command_buffer->state.bind_groups[i].handle) continue;
-
-    // Acquire the bind group to use for the current descriptor set.
-    WGPUBindGroup handle = iree_hal_webgpu_bind_group_cache_acquire(
-        command_buffer->bind_group_cache, binding_info->set_layouts[i],
-        command_buffer->state.bind_groups[i].bindings,
-        binding_info->set_masks[i]);
-
-    // NOTE: today we don't support dynamic offsets for push descriptor sets.
-    // This will be a larger change we'll need to handle in the compiler. If we
-    // wanted to improve caching we could make all the bindings dynamic and then
-    // always cache the base offsets, however
-    // maxDynamicStorageBuffersPerPipelineLayout is minimally 4 and that's not
-    // a lot of bindings.
-    wgpuComputePassEncoderSetBindGroup(compute_pass, (uint32_t)i, handle, 0,
-                                       NULL);
-    command_buffer->state.bind_groups[i].handle = handle;
-    command_buffer->state.bind_groups_empty &= ~(1ull << i);
-  }
-
-  if (constant_count > 0) {
-    // Pad up to IREE_HAL_WEBGPU_PARAMS_BIND_GROUP_INDEX with empty bind groups.
-    WGPUBindGroup empty_handle =
-        command_buffer->staging_buffer->empty_bind_group;
-    for (iree_host_size_t i = binding_info->set_count;
-         i < IREE_HAL_WEBGPU_PARAMS_BIND_GROUP_INDEX; ++i) {
-      // Skip if an empty group is already set at this index.
-      if ((command_buffer->state.bind_groups_empty >> i) & 1ull) continue;
-
-      wgpuComputePassEncoderSetBindGroup(compute_pass, (uint32_t)i,
-                                         empty_handle, 0, NULL);
-      command_buffer->state.bind_groups[i].handle = empty_handle;
-      command_buffer->state.bind_groups_empty |= 1ull << i;
-    }
-  }
-
-  *out_compute_pass = compute_pass;
-  return iree_ok_status();
-}
-
-static iree_status_t iree_hal_webgpu_command_buffer_dispatch(
-    iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_executable_t* executable, int32_t entry_point,
-    uint32_t workgroup_x, uint32_t workgroup_y, uint32_t workgroup_z,
-    iree_hal_dispatch_flags_t flags) {
-  iree_hal_webgpu_command_buffer_t* command_buffer =
-      iree_hal_webgpu_command_buffer_cast(base_command_buffer);
-
-  WGPUComputePassEncoder compute_pass = NULL;
-  IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_prepare_dispatch(
-      command_buffer, executable, entry_point, &compute_pass));
-  wgpuComputePassEncoderDispatchWorkgroups(compute_pass, workgroup_x,
-                                           workgroup_y, workgroup_z);
-
-  return iree_ok_status();
-}
-
-static iree_status_t iree_hal_webgpu_command_buffer_dispatch_indirect(
-    iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_executable_t* executable, int32_t entry_point,
-    iree_hal_buffer_ref_t workgroups_ref, iree_hal_dispatch_flags_t flags) {
-  iree_hal_webgpu_command_buffer_t* command_buffer =
-      iree_hal_webgpu_command_buffer_cast(base_command_buffer);
-
-  WGPUComputePassEncoder compute_pass = NULL;
-  IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_prepare_dispatch(
-      command_buffer, executable, entry_point, &compute_pass));
-  wgpuComputePassEncoderDispatchWorkgroupsIndirect(
-      compute_pass, iree_hal_webgpu_buffer_handle(workgroups_ref.buffer),
-      workgroups_ref.offset);
-
-  return iree_ok_status();
-}
-
 static iree_status_t iree_hal_webgpu_command_buffer_prepare_dispatch(
     iree_hal_webgpu_command_buffer_t* command_buffer,
     iree_hal_executable_t* executable, uint32_t ordinal,
@@ -968,15 +851,16 @@ static iree_status_t iree_hal_webgpu_command_buffer_prepare_dispatch(
     binding_mask |= 1u << i;
     group_bindings[i].type = WGPUBufferBindingType_Storage;
     group_bindings[i].buffer =
-        bindings[i].buffer ? iree_hal_webgpu_buffer_handle(bindings[i].buffer)
-                           : NULL;
-    group_bindings[i] offset = bindings[i].offset;
-    group_bindings[i] length = bindings[i].length;
+        bindings.values[i].buffer
+            ? iree_hal_webgpu_buffer_handle(bindings.values[i].buffer)
+            : NULL;
+    group_bindings[i].offset = bindings.values[i].offset;
+    group_bindings[i].length = bindings.values[i].length;
   }
 
   // Acquire the bind group to use for the current descriptor set.
   WGPUBindGroup handle = iree_hal_webgpu_bind_group_cache_acquire(
-      command_buffer->bind_group_cache, binding_info->set_layout,
+      command_buffer->bind_group_cache, binding_info->set_layouts[0],
       group_bindings, binding_mask);
 
   // NOTE: today we don't support dynamic offsets for push descriptor sets.
@@ -994,36 +878,33 @@ static iree_status_t iree_hal_webgpu_command_buffer_prepare_dispatch(
 static iree_status_t iree_hal_webgpu_command_buffer_dispatch(
     iree_hal_command_buffer_t* base_command_buffer,
     iree_hal_executable_t* executable, int32_t entry_point,
-    const uint32_t workgroup_count[3], iree_const_byte_span_t constants,
-    iree_hal_buffer_ref_list_t bindings, iree_hal_dispatch_flags_t flags) {
+    const iree_hal_dispatch_config_t config, iree_const_byte_span_t constants,
+    const iree_hal_buffer_ref_list_t bindings,
+    iree_hal_dispatch_flags_t flags) {
   iree_hal_webgpu_command_buffer_t* command_buffer =
       iree_hal_webgpu_command_buffer_cast(base_command_buffer);
 
-  WGPUComputePassEncoder compute_pass = NULL;
-  IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_prepare_dispatch(
-      command_buffer, executable, entry_point, constants, bindings, flags,
-      &compute_pass));
-  wgpuComputePassEncoderDispatchWorkgroups(
-      compute_pass, workgroup_count[0], workgroup_count[1], workgroup_count[2]);
-
-  return iree_ok_status();
-}
-
-static iree_status_t iree_hal_webgpu_command_buffer_dispatch_indirect(
-    iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_executable_t* executable, int32_t entry_point,
-    iree_hal_buffer_ref_t workgroups_ref, iree_const_byte_span_t constants,
-    iree_hal_buffer_ref_list_t bindings, iree_hal_dispatch_flags_t flags) {
-  iree_hal_webgpu_command_buffer_t* command_buffer =
-      iree_hal_webgpu_command_buffer_cast(base_command_buffer);
+  if (iree_hal_dispatch_uses_custom_arguments(flags)) {
+    return iree_make_status(
+        IREE_STATUS_UNIMPLEMENTED,
+        "direct/indirect arguments are not supported in WebGPU");
+  }
 
   WGPUComputePassEncoder compute_pass = NULL;
   IREE_RETURN_IF_ERROR(iree_hal_webgpu_command_buffer_prepare_dispatch(
       command_buffer, executable, entry_point, constants, bindings, flags,
       &compute_pass));
-  wgpuComputePassEncoderDispatchWorkgroupsIndirect(
-      compute_pass, iree_hal_webgpu_buffer_handle(workgroups_ref.buffer),
-      workgroups_ref.offset);
+
+  if (iree_hal_dispatch_uses_indirect_parameters(flags)) {
+    wgpuComputePassEncoderDispatchWorkgroupsIndirect(
+        compute_pass,
+        iree_hal_webgpu_buffer_handle(config.workgroup_count_ref.buffer),
+        config.workgroup_count_ref.offset);
+  } else {
+    wgpuComputePassEncoderDispatchWorkgroups(
+        compute_pass, config.workgroup_count[0], config.workgroup_count[1],
+        config.workgroup_count[2]);
+  }
 
   return iree_ok_status();
 }
@@ -1045,7 +926,4 @@ const iree_hal_command_buffer_vtable_t iree_hal_webgpu_command_buffer_vtable = {
     .constants = iree_hal_webgpu_command_buffer_constants,
     .push_descriptor_set = iree_hal_webgpu_command_buffer_push_descriptor_set,
     .dispatch = iree_hal_webgpu_command_buffer_dispatch,
-    .dispatch_indirect = iree_hal_webgpu_command_buffer_dispatch_indirect,
-    .dispatch = iree_hal_webgpu_command_buffer_dispatch,
-    .dispatch_indirect = iree_hal_webgpu_command_buffer_dispatch_indirect,
 };

experimental/webgpu/simple_allocator.c

@@ -176,7 +176,7 @@ static iree_status_t iree_hal_webgpu_simple_allocator_allocate_buffer(
     usage_flags |= WGPUBufferUsage_Uniform;
   }
   if (iree_any_bit_set(params->usage,
-                       IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS)) {
+                       IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMETERS)) {
     usage_flags |= WGPUBufferUsage_Indirect;
   }
 

integrations/pjrt/test/triage_jaxtest.py

@@ -328,7 +328,7 @@ def check_buffer_usage(errortxt, _, __):
     return (
         "requested buffer usage is not supported" in errortxt
         or "tensor requested usage was not specified when the buffer" in errortxt
-        or "PERMISSION_DENIED; requested usage was not specified when the buffer was allocated; buffer allows DISPATCH_INDIRECT_PARAMS"
+        or "PERMISSION_DENIED; requested usage was not specified when the buffer was allocated; buffer allows DISPATCH_INDIRECT_PARAMETERS"
         in errortxt
     )
 

runtime/bindings/python/hal.cc

@@ -1286,8 +1286,8 @@ void SetupHalBindings(nanobind::module_ m) {
       .value("TRANSFER_SOURCE", IREE_HAL_BUFFER_USAGE_TRANSFER_SOURCE)
       .value("TRANSFER_TARGET", IREE_HAL_BUFFER_USAGE_TRANSFER_TARGET)
       .value("TRANSFER", IREE_HAL_BUFFER_USAGE_TRANSFER)
-      .value("DISPATCH_INDIRECT_PARAMS",
-             IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS)
+      .value("DISPATCH_INDIRECT_PARAMETERS",
+             IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMETERS)
       .value("DISPATCH_UNIFORM_READ",
              IREE_HAL_BUFFER_USAGE_DISPATCH_UNIFORM_READ)
       .value("DISPATCH_STORAGE_READ",

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

@@ -42,7 +42,7 @@ class BufferUsage(int):
     DISPATCH_IMAGE: ClassVar[BufferUsage] = ...
     DISPATCH_IMAGE_READ: ClassVar[BufferUsage] = ...
     DISPATCH_IMAGE_WRITE: ClassVar[BufferUsage] = ...
-    DISPATCH_INDIRECT_PARAMS: ClassVar[BufferUsage] = ...
+    DISPATCH_INDIRECT_PARAMETERS: ClassVar[BufferUsage] = ...
     DISPATCH_STORAGE: ClassVar[BufferUsage] = ...
     DISPATCH_STORAGE_READ: ClassVar[BufferUsage] = ...
     DISPATCH_STORAGE_WRITE: ClassVar[BufferUsage] = ...

runtime/src/iree/hal/buffer.c

@@ -86,7 +86,7 @@ static const iree_bitfield_string_mapping_t iree_hal_buffer_usage_mappings[] = {
   // Separate:
   {IREE_HAL_BUFFER_USAGE_TRANSFER_SOURCE, IREE_SVL("TRANSFER_SOURCE")},
   {IREE_HAL_BUFFER_USAGE_TRANSFER_TARGET, IREE_SVL("TRANSFER_TARGET")},
-  {IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS, IREE_SVL("DISPATCH_INDIRECT_PARAMS")},
+  {IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMETERS, IREE_SVL("DISPATCH_INDIRECT_PARAMETERS")},
   {IREE_HAL_BUFFER_USAGE_DISPATCH_UNIFORM_READ, IREE_SVL("DISPATCH_UNIFORM_READ")},
   {IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE_READ, IREE_SVL("DISPATCH_STORAGE_READ")},
   {IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE_WRITE, IREE_SVL("DISPATCH_STORAGE_WRITE")},

runtime/src/iree/hal/buffer.h

@@ -230,7 +230,7 @@ enum iree_hal_buffer_usage_bits_t {
   //  - GPUBufferUsage.INDIRECT
   //  - MTLResourceUsageRead
   //  - VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT
-  IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS = 1u << 8,
+  IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMETERS = 1u << 8,
 
   // Buffer contents are uniformly read by dispatches.
   // These may occasionally be written as storage buffers in cases of
@@ -309,7 +309,7 @@ enum iree_hal_buffer_usage_bits_t {
 
   // Buffer contents are available for use by all dispatch-related operations.
   IREE_HAL_BUFFER_USAGE_DISPATCH =
-      IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS |
+      IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMETERS |
       IREE_HAL_BUFFER_USAGE_DISPATCH_UNIFORM_READ |
       IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE |
       IREE_HAL_BUFFER_USAGE_DISPATCH_IMAGE,

runtime/src/iree/hal/command_buffer.c

@@ -519,12 +519,17 @@ IREE_API_EXPORT iree_status_t iree_hal_command_buffer_collective(
 IREE_API_EXPORT iree_status_t iree_hal_command_buffer_dispatch(
     iree_hal_command_buffer_t* command_buffer,
     iree_hal_executable_t* executable, int32_t entry_point,
-    const uint32_t workgroup_count[3], iree_const_byte_span_t constants,
-    iree_hal_buffer_ref_list_t bindings, iree_hal_dispatch_flags_t flags) {
+    const iree_hal_dispatch_config_t config, iree_const_byte_span_t constants,
+    const iree_hal_buffer_ref_list_t bindings,
+    iree_hal_dispatch_flags_t flags) {
   IREE_ASSERT_ARGUMENT(command_buffer);
   IREE_ASSERT_ARGUMENT(executable);
 
-  if ((workgroup_count[0] | workgroup_count[1] | workgroup_count[2]) == 0) {
+  const bool has_static_workgroup_count =
+      !iree_hal_dispatch_uses_indirect_parameters(flags);
+  if (has_static_workgroup_count &&
+      (config.workgroup_count[0] | config.workgroup_count[1] |
+       config.workgroup_count[2]) == 0) {
     // No-op dispatch. All implementations are expected to do this but we ensure
     // it happens here to avoid the overhead of going all the way down into the
     // device layer for something we know should have no (intentional)
@@ -543,50 +548,34 @@ IREE_API_EXPORT iree_status_t iree_hal_command_buffer_dispatch(
   // slice off a much larger allocation and then suballocate from that ourselves
   // so that we could avoid the tracy_malloc overheads per-dispatch.
   IREE_TRACE({
-    char xyz_string[32];
-    int xyz_string_length =
-        snprintf(xyz_string, IREE_ARRAYSIZE(xyz_string), "%ux%ux%u",
-                 workgroup_count[0], workgroup_count[1], workgroup_count[2]);
-    IREE_TRACE_ZONE_APPEND_TEXT(z0, xyz_string, xyz_string_length);
+    if (has_static_workgroup_count) {
+      char xyz_string[32];
+      int xyz_string_length =
+          snprintf(xyz_string, IREE_ARRAYSIZE(xyz_string), "%ux%ux%u",
+                   config.workgroup_count[0], config.workgroup_count[1],
+                   config.workgroup_count[2]);
+      IREE_TRACE_ZONE_APPEND_TEXT(z0, xyz_string, xyz_string_length);
+    } else {
+      IREE_TRACE_ZONE_APPEND_TEXT(z0, "(indirect)");
+    }
   });
 #endif  // IREE_HAL_VERBOSE_TRACING_ENABLE
 
   IF_VALIDATING(command_buffer, {
     IREE_RETURN_AND_END_ZONE_IF_ERROR(
         z0, iree_hal_command_buffer_dispatch_validation(
                 command_buffer, VALIDATION_STATE(command_buffer), executable,
-                entry_point, workgroup_count, constants, bindings, flags));
+                entry_point, config, constants, bindings, flags));
   });
 
   iree_status_t status = _VTABLE_DISPATCH(command_buffer, dispatch)(
-      command_buffer, executable, entry_point, workgroup_count, constants,
-      bindings, flags);
+      command_buffer, executable, entry_point, config, constants, bindings,
+      flags);
 
   IREE_TRACE_ZONE_END(z0);
   return status;
 }
 
-IREE_API_EXPORT iree_status_t iree_hal_command_buffer_dispatch_indirect(
-    iree_hal_command_buffer_t* command_buffer,
-    iree_hal_executable_t* executable, int32_t entry_point,
-    iree_hal_buffer_ref_t workgroups_ref, iree_const_byte_span_t constants,
-    iree_hal_buffer_ref_list_t bindings, iree_hal_dispatch_flags_t flags) {
-  IREE_ASSERT_ARGUMENT(command_buffer);
-  IREE_ASSERT_ARGUMENT(executable);
-  IREE_TRACE_ZONE_BEGIN(z0);
-  IF_VALIDATING(command_buffer, {
-    IREE_RETURN_AND_END_ZONE_IF_ERROR(
-        z0, iree_hal_command_buffer_dispatch_indirect_validation(
-                command_buffer, VALIDATION_STATE(command_buffer), executable,
-                entry_point, workgroups_ref, constants, bindings, flags));
-  });
-  iree_status_t status = _VTABLE_DISPATCH(command_buffer, dispatch_indirect)(
-      command_buffer, executable, entry_point, workgroups_ref, constants,
-      bindings, flags);
-  IREE_TRACE_ZONE_END(z0);
-  return status;
-}
-
 //===----------------------------------------------------------------------===//
 // Validation support
 //===----------------------------------------------------------------------===//