[ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation) (#21243)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #20842 by
@JCNTH
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/JCNTH/17/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/17/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/17/orig

@diff-train-skip-merge

---------

Co-authored-by: Julian Ng-Thow-Hing <juliannth@meta.com>

Author: pytorchbot

backends/webgpu/CMakeLists.txt

@@ -109,6 +109,7 @@ function(add_webgpu_native_test test_name test_src)
   target_link_libraries(
     ${test_name}
     PRIVATE webgpu_backend
+            vulkan_schema
             ${WEBGPU_GPU_LIB}
             executorch_core
             extension_module_static
@@ -153,10 +154,11 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
     )
 
     # Device-free util unit test: no backend/Dawn link (pure manifest/tolerance
-    # helpers), so it does NOT use the native-test helper.
+    # + dispatch-grid-math helpers), so it does NOT use the native-test helper.
     add_executable(
-      webgpu_op_test_util_test test/op_tests/test_driver_util.cpp
-                               test/op_tests/driver_util.cpp
+      webgpu_op_test_util_test
+      test/op_tests/test_driver_util.cpp test/op_tests/driver_util.cpp
+      test/native/test_webgpu_utils.cpp
     )
     target_include_directories(
       webgpu_op_test_util_test

backends/webgpu/runtime/WebGPUDispatchMath.h

@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// Pure dispatch-grid math with zero WebGPU/Dawn dependency, so it is
+// unit-testable without a WGPUDevice (split out of WebGPUUtils.h, which
+// requires <webgpu/webgpu.h> for its device-facing functions).
+
+#include <cmath>
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+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;
+}
+
+// Product of a tensor's dims; the same accumulation was duplicated per-op.
+inline uint64_t numel(const std::vector<int64_t>& dims) {
+  uint64_t n = 1;
+  for (int64_t d : dims) {
+    if (d < 0) {
+      throw std::runtime_error("numel: negative dimension");
+    }
+    uint64_t ud = static_cast<uint64_t>(d);
+    if (ud != 0 && n > UINT64_MAX / ud) {
+      throw std::runtime_error("numel: element count overflow");
+    }
+    n *= ud;
+  }
+  return n;
+}
+
+// Broadcasts a 1- or 2-element int list to (h, w); PyTorch's convention for
+// kernel_size/stride/padding/dilation args. Was duplicated as a local `hw`
+// lambda in conv2d/conv_transpose2d/max_pool2d.
+inline void parse_hw(
+    const std::vector<int64_t>& v,
+    uint32_t& h,
+    uint32_t& w,
+    const char* op_name,
+    const char* arg_name) {
+  if (v.size() == 1) {
+    h = w = static_cast<uint32_t>(v[0]);
+  } else if (v.size() == 2) {
+    h = static_cast<uint32_t>(v[0]);
+    w = static_cast<uint32_t>(v[1]);
+  } else {
+    throw std::runtime_error(
+        std::string("WebGPU ") + op_name + ": " + arg_name +
+        " must be 1 or 2 elements");
+  }
+}
+
+// Adaptive 1D->2D dispatch grid. `count_x`/`count_y` are the dispatch dims;
+// `stride_x` (= count_x * wg_size) lets the shader decode a flat index as
+// `gid.y * stride_x + gid.x`. Used by ops whose element count can exceed the
+// 65535 per-dimension ceiling that compute_1d_workgroup_count throws on.
+struct DispatchGrid {
+  uint32_t wg_size;
+  uint32_t count_x;
+  uint32_t count_y;
+  uint32_t stride_x;
+};
+
+// 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
+// mostly-idle workgroups (up to ~2x the needed launch) when total isn't a
+// clean multiple of max_dim.
+inline DispatchGrid compute_dispatch_grid_from_limits(
+    uint32_t total, // workgroups needed (1D)
+    uint32_t wg_size,
+    uint32_t max_dim,
+    const char* op_name) {
+  DispatchGrid g;
+  g.wg_size = wg_size;
+  if (total <= max_dim) {
+    g.count_x = total;
+    g.count_y = 1;
+  } else {
+    uint32_t sq =
+        static_cast<uint32_t>(std::ceil(std::sqrt(static_cast<double>(total))));
+    g.count_x = sq < max_dim ? sq : max_dim;
+    g.count_y = div_up(total, g.count_x);
+    if (g.count_y >
+        max_dim) { // > max_dim^2 * wg threads — astronomically large
+      throw std::runtime_error(
+          std::string("WebGPU ") + op_name +
+          ": dispatch exceeds 2D grid capacity");
+    }
+  }
+  g.stride_x = g.count_x * g.wg_size;
+  return g;
+}
+
+} // namespace executorch::backends::webgpu::utils

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -418,6 +418,7 @@ void WebGPUGraph::build(
   value_lists_.resize(num_vals);
   doubles_.resize(num_vals, 0.0);
   bools_.resize(num_vals, false);
+  strings_.resize(num_vals);
 
   // Pre-scan the op chain: a constant may be DEFERRED (no eager GPU buffer; the
   // prepack node materializes it once) only if it is a prepack source AND never
@@ -625,6 +626,14 @@ void WebGPUGraph::build(
         bools_[i] = val->value_as_Bool()->bool_val();
         break;
       }
+      case vkgraph::GraphTypes::String: {
+        value_types_[i] = ValueType::String;
+        const auto* sv = val->value_as_String()->string_val();
+        if (sv) {
+          strings_[i] = sv->str();
+        }
+        break;
+      }
       case vkgraph::GraphTypes::SymInt: {
         // Live scalar: small Uniform buffer the CPU rewrites per execute.
         value_types_[i] = ValueType::SymInt;

backends/webgpu/runtime/WebGPUGraph.h

@@ -152,6 +152,10 @@ class WebGPUGraph {
   bool get_bool(int id) const {
     return bools_[id];
   }
+  // String value (e.g. gelu's `approximate` kwarg).
+  const std::string& get_string(int id) const {
+    return strings_[id];
+  }
 
   // Live-scalar (SymInt) API; mirrors the Vulkan SymInt/ParamsBuffer UBO.
   // set_symint writes the buffer + marks dirty only if the value changed.
@@ -240,6 +244,20 @@ class WebGPUGraph {
     return dispatches_.size() - 1;
   }
 
+  // 2D sibling of add_dispatch (sets workgroup_count_y); returns the index.
+  size_t add_dispatch_2d(
+      WGPUComputePipeline pipeline,
+      WGPUBindGroup bind_group,
+      uint32_t count_x,
+      uint32_t count_y) {
+    WebGPUDispatch d;
+    d.pipeline = pipeline;
+    d.bind_group = bind_group;
+    d.workgroup_count_x = count_x;
+    d.workgroup_count_y = count_y;
+    return add_dispatch(d);
+  }
+
   // In-graph buffer-to-buffer DMA (e.g. flat copy); returns the dispatch index.
   size_t add_buffer_copy(WGPUBuffer src, WGPUBuffer dst, size_t nbytes) {
     WebGPUDispatch d;
@@ -376,6 +394,7 @@ class WebGPUGraph {
   std::vector<std::vector<int>> value_lists_;
   std::vector<double> doubles_;
   std::vector<bool> bools_;
+  std::vector<std::string> strings_;
 
   // SymInt (live scalar): id -> {live Uniform buffer, current value}, sparse.
   struct SymIntSlot {