[ExecuTorch][WebGPU] Coordinate dynamic dispatch routes

backends/webgpu/CMakeLists.txt

@@ -32,6 +32,7 @@ set(WEBGPU_SRCS
     runtime/WebGPUDelegateHeader.cpp
     runtime/WebGPUDevice.cpp
     runtime/WebGPUQueryPool.cpp
+    runtime/WebGPUShaderRegistry.cpp
     runtime/ops/OperatorRegistry.cpp
 )
 
@@ -222,5 +223,10 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
       webgpu_output_suppression_test test/native/test_output_suppression.cpp
     )
     target_link_libraries(webgpu_output_suppression_test PRIVATE GTest::gtest)
+
+    add_webgpu_native_test(
+      webgpu_compute_dispatch_test test/native/test_compute_dispatch.cpp
+    )
+    target_link_libraries(webgpu_compute_dispatch_test PRIVATE GTest::gtest)
   endif()
 endif()

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -18,6 +18,7 @@
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <executorch/runtime/platform/log.h>
 
+#include <cstring>
 #include <vector>
 
 #include <new>
@@ -43,6 +44,27 @@ using executorch::runtime::resize_tensor;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
 
+Result<WebGPUGraphConfig> parse_webgpu_graph_config(
+    ArrayRef<CompileSpec> compile_specs) {
+  WebGPUGraphConfig config;
+  for (const CompileSpec& spec : compile_specs) {
+    if (spec.key == nullptr ||
+        std::strcmp(spec.key, "webgpu_record_q4gsw_decode_route") != 0) {
+      continue;
+    }
+    if (spec.value.nbytes != sizeof(uint8_t) || spec.value.buffer == nullptr) {
+      ET_LOG(
+          Error,
+          "WebGPU compile option webgpu_record_q4gsw_decode_route must be "
+          "exactly one byte");
+      return Error::DelegateInvalidCompatibility;
+    }
+    config.record_q4gsw_decode_route =
+        *static_cast<const uint8_t*>(spec.value.buffer) != 0;
+  }
+  return config;
+}
+
 bool WebGPUBackend::is_available() const {
   return true;
 }
@@ -51,6 +73,13 @@ Result<DelegateHandle*> WebGPUBackend::init(
     BackendInitContext& context,
     FreeableBuffer* processed,
     ArrayRef<CompileSpec> compile_specs) const {
+  Result<WebGPUGraphConfig> parsed_config =
+      parse_webgpu_graph_config(compile_specs);
+  if (!parsed_config.ok()) {
+    return parsed_config.error();
+  }
+  WebGPUGraphConfig config = parsed_config.get();
+
   // Allocate graph on the runtime allocator
   WebGPUGraph* graph =
       context.get_runtime_allocator()->allocateInstance<WebGPUGraph>();
@@ -83,29 +112,23 @@ Result<DelegateHandle*> WebGPUBackend::init(
   // Load-time backend option (BackendOption / LoadBackendOptionsMap), keyed by
   // the registered backend name; default false. Mirrors the CoreML/XNNPACK
   // runtime-spec pattern -- no compile flag and no .pte re-export needed.
-  bool enable_f16_kv_cache = false;
   {
     Result<bool> spec = context.get_runtime_spec<bool>("enable_f16_kv_cache");
     if (spec.ok()) {
-      enable_f16_kv_cache = spec.get();
+      config.f16_kv_cache = spec.get();
     }
   }
-  bool enable_f16_accumulate_gemm = false;
   {
     Result<bool> spec =
         context.get_runtime_spec<bool>("enable_f16_accumulate_gemm");
     if (spec.ok()) {
-      enable_f16_accumulate_gemm = spec.get();
+      config.f16_accumulate_gemm = spec.get();
     }
   }
 
   try {
     graph->build(
-        flatbuffer_data,
-        constant_data,
-        context.get_named_data_map(),
-        enable_f16_kv_cache,
-        enable_f16_accumulate_gemm);
+        flatbuffer_data, constant_data, context.get_named_data_map(), config);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU graph build failed: %s", e.what());
     graph->~WebGPUGraph();

backends/webgpu/runtime/WebGPUBackend.h

@@ -8,12 +8,17 @@
 
 #pragma once
 
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/runtime/backend/interface.h>
 
 namespace executorch {
 namespace backends {
 namespace webgpu {
 
+executorch::runtime::Result<WebGPUGraphConfig> parse_webgpu_graph_config(
+    executorch::runtime::ArrayRef<executorch::runtime::CompileSpec>
+        compile_specs);
+
 class WebGPUBackend final : public ::executorch::runtime::BackendInterface {
  public:
   ~WebGPUBackend() override = default;

backends/webgpu/runtime/WebGPUDispatchMath.h

@@ -13,6 +13,7 @@
 // requires <webgpu/webgpu.h> for its device-facing functions).
 
 #include <cmath>
+#include <cstddef>
 #include <cstdint>
 #include <stdexcept>
 #include <string>
@@ -74,6 +75,110 @@ struct DispatchGrid {
   uint32_t stride_x;
 };
 
+struct WgCount {
+  uint32_t x;
+  uint32_t y;
+};
+
+struct DispatchRange {
+  size_t begin;
+  size_t end;
+};
+
+constexpr bool should_record_q4gsw_dual_route(
+    uint32_t max_m,
+    bool bicol_eligible,
+    bool has_dynamic_shapes,
+    bool record_q4gsw_decode_route) {
+  return max_m > 1u && bicol_eligible &&
+      (has_dynamic_shapes || record_q4gsw_decode_route);
+}
+
+constexpr bool should_record_sdpa_dual_route(
+    bool fd_eligible,
+    bool has_dynamic_sequence) {
+  return fd_eligible && has_dynamic_sequence;
+}
+
+class DispatchRouteRegistry {
+ public:
+  template <typename IsCompute>
+  size_t register_group(
+      size_t dispatch_count,
+      const std::vector<DispatchRange>& ranges,
+      IsCompute&& is_compute) {
+    if (dispatch_count < owners_.size() || ranges.size() < 2) {
+      throw std::runtime_error("invalid WebGPU dispatch route group");
+    }
+
+    std::vector<bool> claimed(dispatch_count, false);
+    for (const auto& range : ranges) {
+      if (range.begin >= range.end || range.end > dispatch_count) {
+        throw std::runtime_error("invalid WebGPU dispatch route range");
+      }
+      for (size_t i = range.begin; i < range.end; i++) {
+        if (!is_compute(i)) {
+          throw std::runtime_error(
+              "WebGPU dispatch route contains a copy command");
+        }
+        if (claimed[i] || (i < owners_.size() && owners_[i] != kNoOwner)) {
+          throw std::runtime_error("overlapping WebGPU dispatch route ranges");
+        }
+        claimed[i] = true;
+      }
+    }
+
+    const size_t group = groups_.size();
+    owners_.resize(dispatch_count, kNoOwner);
+    for (size_t i = 0; i < claimed.size(); i++) {
+      if (claimed[i]) {
+        owners_[i] = group;
+      }
+    }
+    groups_.push_back(ranges);
+    return group;
+  }
+
+  template <typename SetGrid>
+  void select(
+      size_t group,
+      size_t active_route,
+      const std::vector<WgCount>& active_grids,
+      SetGrid&& set_grid) const {
+    if (group >= groups_.size()) {
+      throw std::runtime_error("invalid WebGPU dispatch route group");
+    }
+    const auto& ranges = groups_[group];
+    if (active_route >= ranges.size()) {
+      throw std::runtime_error("invalid active WebGPU dispatch route");
+    }
+    const auto& active = ranges[active_route];
+    if (active_grids.size() != active.end - active.begin) {
+      throw std::runtime_error("WebGPU dispatch route grid count mismatch");
+    }
+    for (const auto& grid : active_grids) {
+      if (grid.x == 0 || grid.y == 0) {
+        throw std::runtime_error(
+            "active WebGPU dispatch route has a zero grid");
+      }
+    }
+
+    for (const auto& range : ranges) {
+      for (size_t i = range.begin; i < range.end; i++) {
+        set_grid(i, {0, 0});
+      }
+    }
+    for (size_t i = 0; i < active_grids.size(); i++) {
+      set_grid(active.begin + i, active_grids[i]);
+    }
+  }
+
+ private:
+  static constexpr size_t kNoOwner = static_cast<size_t>(-1);
+  std::vector<std::vector<DispatchRange>> groups_;
+  std::vector<size_t> owners_;
+};
+
 // Given the workgroup count needed (1D) and the device's per-dimension
 // dispatch-count ceiling, compute a near-square 2D grid rather than
 // {max_dim, div_up(total, max_dim)} — maxing one dim pads the other with

backends/webgpu/runtime/WebGPUExecutionOptions.cpp

@@ -39,7 +39,12 @@ WebGPUExecutionPlan plan_webgpu_execution(
     size_t output_count,
     ExecuteConfig config,
     const std::vector<SuppressibleOutput>& suppressible_outputs,
-    WebGPUGraphExecutionOptions options) {
+    WebGPUGraphExecutionOptions options,
+    const std::vector<bool>& enabled_dispatches) {
+  if (!enabled_dispatches.empty() &&
+      enabled_dispatches.size() != dispatch_count) {
+    throw std::runtime_error("WebGPU: enabled dispatch count mismatch");
+  }
   std::vector<bool> suppressed_dispatches(dispatch_count, false);
   std::vector<bool> copy_outputs(output_count, true);
   std::vector<bool> seen_output_ordinals(output_count, false);
@@ -78,7 +83,8 @@ WebGPUExecutionPlan plan_webgpu_execution(
     std::vector<size_t> indices;
     indices.reserve(end - begin);
     for (size_t i = begin; i < end; i++) {
-      if (!suppressed_dispatches[i]) {
+      if (!suppressed_dispatches[i] &&
+          (enabled_dispatches.empty() || enabled_dispatches[i])) {
         indices.push_back(i);
       }
     }

backends/webgpu/runtime/WebGPUExecutionOptions.h

@@ -51,7 +51,8 @@ WebGPUExecutionPlan plan_webgpu_execution(
     size_t output_count,
     ExecuteConfig config,
     const std::vector<SuppressibleOutput>& suppressible_outputs,
-    WebGPUGraphExecutionOptions options);
+    WebGPUGraphExecutionOptions options,
+    const std::vector<bool>& enabled_dispatches = {});
 
 WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options(
     const std::vector<const void*>& delegate_outputs,