[ExecuTorch][WebGPU] Lift FastVLM dispatch limits and harden graph construction

backends/webgpu/runtime/WebGPUDispatchMath.h

@@ -24,7 +24,7 @@ namespace executorch::backends::webgpu::utils {
 // Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up).
 template <typename T>
 inline T div_up(T a, T b) {
-  return (a + b - 1) / b;
+  return a / b + (a % b != 0);
 }
 
 // Product of a tensor's dims; the same accumulation was duplicated per-op.

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -1518,17 +1518,10 @@ constexpr uint32_t kRouteK16CausalBound = 1u << 11;
 constexpr uint32_t kRouteBicolSubgroup = 1u << 12;
 constexpr uint32_t kRouteQwen3Q16K16 = 1u << 13;
 constexpr uint32_t kRouteQwen3Q32K16 = 1u << 14;
-#endif // WGPU_BACKEND_ENABLE_PROFILING
-
-// Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the
-// WEBGPU_TIMESTAMP_QUERY env var enables per-pass GPU timestamp queries.
 bool should_timestamp_query() {
-#ifdef WGPU_BACKEND_ENABLE_PROFILING
   return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr;
-#else
-  return false;
-#endif
 }
+#endif // WGPU_BACKEND_ENABLE_PROFILING
 } // namespace
 
 #ifdef WGPU_BACKEND_ENABLE_PROFILING
@@ -1739,12 +1732,13 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) {
     return 1;
   }
 
-  // GPU timestamp queries assume one submit; chunked execute is multi-submit.
+#ifdef WGPU_BACKEND_ENABLE_PROFILING
   if (should_timestamp_query()) {
     throw std::runtime_error(
         "WebGPU: WEBGPU_TIMESTAMP_QUERY is incompatible with chunked execute "
         "(multi-submit); disable chunking to use GPU timestamp queries");
   }
+#endif // WGPU_BACKEND_ENABLE_PROFILING
 
   for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size();
        chunk_index++) {

backends/webgpu/runtime/ops/compare/Compare.cpp

@@ -85,9 +85,7 @@ void compare_impl(
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  WGPUBuffer uniform_buffer =
-      utils::make_uniform(device, &params, sizeof(CompareParams));
-  graph.add_uniform_buffer_bytes(sizeof(CompareParams));
+  WGPUBuffer uniform_buffer = graph.create_params_buffer(params);
 
   // out (rw storage) + in1/in2 (ro storage) + params (uniform).
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
@@ -142,8 +140,6 @@ void compare_impl(
   };
   graph.add_tensor_resize_hook(in1_id, resize);
   graph.add_tensor_resize_hook(in2_id, resize);
-
-  graph.own_uniform_buffer(uniform_buffer);
 }
 
 void eq_op(WebGPUGraph& graph, const std::vector<int>& args) {

backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp

@@ -14,6 +14,7 @@
 
 #include <webgpu/webgpu.h>
 
+#include <limits>
 #include <stdexcept>
 
 namespace executorch::backends::webgpu {
@@ -34,6 +35,22 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const auto& in_tensor = graph.get_tensor(in_id);
   const auto& out_tensor = graph.get_tensor(out_id);
 
+  if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::IntList) {
+    throw std::runtime_error(
+        "WebGPU expand_copy: dynamic target sizes are unsupported");
+  }
+  for (int64_t target_size : graph.get_int_list(args.at(1))) {
+    if (target_size == -1) {
+      throw std::runtime_error(
+          "WebGPU expand_copy: inferred target sizes are unsupported");
+    }
+  }
+  if (graph.tensor_has_dynamic_dims(in_id) ||
+      graph.tensor_has_dynamic_dims(out_id)) {
+    throw std::runtime_error(
+        "WebGPU expand_copy: dynamic shapes are unsupported");
+  }
+
   TensorMeta out_meta;
   TensorMeta in_meta;
   fill_tensor_meta(out_tensor, &out_meta);
@@ -44,21 +61,24 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     throw std::runtime_error(
         "expand_copy: non-fp32 operand (nbytes != numel*4)");
   }
+  if (out_meta.numel >
+      static_cast<uint32_t>(std::numeric_limits<int32_t>::max())) {
+    throw std::runtime_error(
+        "WebGPU expand_copy: element count exceeds the flattened 2D dispatch "
+        "limit");
+  }
 
   uint32_t wg_size =
       utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX);
-  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, out_meta.numel, wg_size, "expand_copy");
 
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
 
-  WGPUBuffer out_meta_buf =
-      utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
-  WGPUBuffer in_meta_buf =
-      utils::make_uniform(device, &in_meta, sizeof(TensorMeta));
-  graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta));
+  WGPUBuffer out_meta_buf = graph.create_params_buffer(out_meta);
+  WGPUBuffer in_meta_buf = graph.create_params_buffer(in_meta);
 
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
       device,
@@ -78,10 +98,8 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       &wg_size_constant,
       1);
 
-  graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
-
-  wgpuBufferRelease(out_meta_buf);
-  wgpuBufferRelease(in_meta_buf);
+  graph.add_dispatch_2d(
+      bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y);
 }
 
 } // namespace

backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl

@@ -13,8 +13,10 @@ struct TensorMeta {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (idx >= out_meta.numel) {
         return;
     }

backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from expand_copy.wgsl - DO NOT EDIT.
-// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370
+// wgsl-sha256: b3c032ab961ffde245fc44289b67df3b5e4ca93eedb9ada2f20a3eaa6f10e9c6
 inline constexpr const char* kExpandCopyWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<f32>;
@@ -30,8 +30,10 @@ struct TensorMeta {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (idx >= out_meta.numel) {
         return;
     }

backends/webgpu/runtime/ops/gelu/Gelu.cpp

@@ -13,6 +13,7 @@
 
 #include <webgpu/webgpu.h>
 
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -42,23 +43,26 @@ void gelu_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const auto& out_tensor = graph.get_tensor(out_id);
   utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu");
 
-  uint32_t num_elements =
-      static_cast<uint32_t>(out_tensor.nbytes / sizeof(float));
+  const uint64_t num_elements64 = out_tensor.nbytes / sizeof(float);
+  if (num_elements64 >
+      static_cast<uint64_t>(std::numeric_limits<int32_t>::max())) {
+    throw std::runtime_error(
+        "WebGPU gelu: element count exceeds the flattened 2D dispatch limit");
+  }
+  const uint32_t num_elements = static_cast<uint32_t>(num_elements64);
 
   // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom).
   uint32_t num_vec4_threads = utils::div_up(num_elements, 4u);
   uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX);
-  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, num_vec4_threads, wg_size, "gelu");
 
   WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size);
 
   GeluParams params = {};
   params.num_elements = num_elements;
 
-  WGPUBuffer uniform_buffer =
-      utils::make_uniform(device, &params, sizeof(GeluParams));
-  graph.add_uniform_buffer_bytes(sizeof(GeluParams));
+  WGPUBuffer uniform_buffer = graph.create_params_buffer(params);
 
   // input (read storage) + output (storage) + params. The exact/approximate
   // choice is baked into the compiled pipeline via the entry point (mirrors
@@ -85,10 +89,38 @@ void gelu_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       1,
       exact ? "main_erf" : "main_tanh");
 
-  graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
-
-  // Drop our ref; the bind group keeps the uniform buffer alive until release.
-  wgpuBufferRelease(uniform_buffer);
+  const size_t dispatch_idx = graph.add_dispatch_2d(
+      bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y);
+
+  WGPUBuffer params_buf = uniform_buffer;
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) {
+        const auto& dims = g.cur_dims(in_id);
+        const uint64_t num_elements64 = utils::numel_of(dims);
+        if (num_elements64 >
+            static_cast<uint64_t>(std::numeric_limits<int32_t>::max())) {
+          throw std::runtime_error(
+              "WebGPU gelu(resize): element count exceeds the flattened 2D "
+              "dispatch limit");
+        }
+        const uint32_t num_elements = static_cast<uint32_t>(num_elements64);
+        g.set_cur_dims(out_id, dims);
+
+        GeluParams params = {};
+        params.num_elements = num_elements;
+        wgpuQueueWriteBuffer(
+            g.queue(), params_buf, 0, &params, sizeof(GeluParams));
+
+        const uint32_t num_vec4_threads = utils::div_up(num_elements, 4u);
+        const utils::WgCount resized_workgroup_count =
+            utils::compute_2d_workgroup_count(
+                g.device(), num_vec4_threads, wg_size, "gelu(resize)");
+        g.dispatch_at(dispatch_idx).workgroup_count_x =
+            resized_workgroup_count.x;
+        g.dispatch_at(dispatch_idx).workgroup_count_y =
+            resized_workgroup_count.y;
+      });
 }
 
 } // namespace

backends/webgpu/runtime/ops/gelu/gelu.wgsl

@@ -33,8 +33,11 @@ fn gelu_erf4(x: vec4<f32>) -> vec4<f32> {
 // before use), computes GELU as one vec4 op, then scatters back only the
 // in-bounds lanes.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main_tanh(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let base = gid.x * 4u;
+fn main_tanh(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    let base = thread_idx * 4u;
     if (base >= params.num_elements) {
         return;
     }
@@ -49,8 +52,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main_erf(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let base = gid.x * 4u;
+fn main_erf(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    let base = thread_idx * 4u;
     if (base >= params.num_elements) {
         return;
     }

backends/webgpu/runtime/ops/gelu/gelu_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from gelu.wgsl - DO NOT EDIT.
-// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f
+// wgsl-sha256: 96570753688590fa009ee5503f754cf3eb572dcb3dcae6818220fe06fe3139ee
 inline constexpr const char* kGeluWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<f32>;
@@ -50,8 +50,11 @@ fn gelu_erf4(x: vec4<f32>) -> vec4<f32> {
 // before use), computes GELU as one vec4 op, then scatters back only the
 // in-bounds lanes.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main_tanh(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let base = gid.x * 4u;
+fn main_tanh(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    let base = thread_idx * 4u;
     if (base >= params.num_elements) {
         return;
     }
@@ -66,8 +69,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main_erf(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let base = gid.x * 4u;
+fn main_erf(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    let base = thread_idx * 4u;
     if (base >= params.num_elements) {
         return;
     }

backends/webgpu/runtime/ops/to_copy/ToCopy.cpp

@@ -70,9 +70,7 @@ void add_convert_op(
   ConvertParams params = {};
   params.num_elements = num_elements;
 
-  WGPUBuffer uniform_buffer =
-      utils::make_uniform(device, &params, sizeof(ConvertParams));
-  graph.add_uniform_buffer_bytes(sizeof(ConvertParams));
+  WGPUBuffer uniform_buffer = graph.create_params_buffer(params);
 
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
       device,
@@ -115,9 +113,6 @@ void add_convert_op(
                 wg_size,
                 "to_copy(resize)");
       });
-
-  // Graph owns it so the resize hook can rewrite it; freed in the dtor.
-  graph.own_uniform_buffer(uniform_buffer);
 }
 
 // Decode byte-packed bool storage into numeric fp32 values.
@@ -160,9 +155,7 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) {
 
   ConvertParams params = {};
   params.num_elements = num_elements;
-  WGPUBuffer uniform_buffer =
-      utils::make_uniform(device, &params, sizeof(ConvertParams));
-  graph.add_uniform_buffer_bytes(sizeof(ConvertParams));
+  WGPUBuffer uniform_buffer = graph.create_params_buffer(params);
 
   utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
       device,
@@ -208,8 +201,6 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) {
                 wg_size,
                 "to_copy_bool_to_float(resize)");
       });
-
-  graph.own_uniform_buffer(uniform_buffer);
 }
 
 void to_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
@@ -223,6 +214,12 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) {
   const auto& in_tensor = graph.get_tensor(in_id);
   const auto& out_tensor = graph.get_tensor(out_id);
 
+  if (in_tensor.is_bool != out_tensor.is_bool && in_tensor.is_int &&
+      out_tensor.is_int) {
+    throw std::runtime_error(
+        "WebGPU to_copy: bool and integer conversions are unsupported");
+  }
+
   // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain.
   if (in_tensor.is_int == out_tensor.is_int &&
       in_tensor.elem_size == out_tensor.elem_size) {