Run WebGPU models in browser Wasm (#23938) (#23938)

Summary:

### 🎯 Why
The opt-in ExecuTorch Wasm build cannot run WebGPU-delegated models in a
browser. This change links the WebGPU backend into that build and
removes two op limits that large vision models hit.

### 🔧 What changed
- Link the WebGPU backend into the opt-in ExecuTorch Wasm build behind
EXECUTORCH_BUILD_WASM_WEBGPU (off by default).
- Initialize one process-wide WebGPU context through a new initWebGPU()
binding.
- Enable Asyncify, exceptions, and a growable heap with a 4 GB maximum
in that configuration only.
- Pin Emscripten 4.0.19, the first pinned version whose emdawnwebgpu
port has the instance-features API that the backend uses.
- Support batched aten.linear inputs by flattening leading dimensions.
- Fold large aten.where dispatches across two workgroup dimensions.
- Build the WebGPU Wasm configuration in pull.yml (build only, because
running a delegated model needs a browser with WebGPU).
- Add configuration tests that check the CMake option, link flags,
bindings, build script, and CI pin.

A build without EXECUTORCH_BUILD_WASM_WEBGPU is unchanged.

Reviewed By: psiddh

Differential Revision: D117468920

Author: raedle

.ci/scripts/setup-emscripten.sh

@@ -15,10 +15,15 @@ install_node() {
 }
 
 install_emscripten() {
+    # The WebGPU backend's instance API (WGPUInstanceFeatureName /
+    # WGPUInstanceDescriptor::requiredFeatures) needs the emdawnwebgpu port
+    # pinned by Emscripten 4.0.13 or newer; 4.0.10 ships the older
+    # capabilities-based descriptor and does not compile it.
+    local emscripten_version="${EMSCRIPTEN_VERSION:-4.0.19}"
     git clone https://github.com/emscripten-core/emsdk.git
     pushd emsdk || return
-    ./emsdk install 4.0.10
-    ./emsdk activate 4.0.10
+    ./emsdk install "${emscripten_version}"
+    ./emsdk activate "${emscripten_version}"
     source ./emsdk_env.sh
     popd || return
 }

.github/workflows/pull.yml

@@ -1581,6 +1581,10 @@ jobs:
         # Test selective build
         PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh
 
+        # Build the WebGPU configuration of the Wasm bindings. Build only:
+        # running a delegated model needs a browser with WebGPU.
+        PYTHON_EXECUTABLE=python bash scripts/build_wasm_tests.sh --webgpu
+
   unittest-wasm-bindings:
     needs: docker-image
     name: unittest-wasm-bindings

CMakeLists.txt

@@ -1543,10 +1543,6 @@ if(EXECUTORCH_BUILD_PYBIND)
   executorch_target_copy_mlx_metallib(portable_lib)
 endif()
 
-if(EXECUTORCH_BUILD_WASM)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/wasm)
-endif()
-
 if(EXECUTORCH_BUILD_TOKENIZERS_WASM)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/wasm/tokenizers)
 endif()
@@ -1602,6 +1598,12 @@ if(EXECUTORCH_BUILD_VGF)
   list(APPEND _executorch_backends vgf_backend)
 endif()
 
+# Added after the backends so EXECUTORCH_BUILD_WASM_WEBGPU can link
+# webgpu_backend, which has to already be defined at that point.
+if(EXECUTORCH_BUILD_WASM)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/wasm)
+endif()
+
 # Top-level interface targets.
 
 # A target containing all configured backends.

backends/webgpu/runtime/ops/linear/Linear.cpp

@@ -16,6 +16,7 @@
 
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <vector>
 
@@ -53,21 +54,39 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const auto& w = graph.get_tensor(w_id);
   const auto& out = graph.get_tensor(out_id);
 
-  if (in.dims.size() != 2 || w.dims.size() != 2) {
-    throw std::runtime_error("WebGPU linear: input/weight must be 2D");
+  if (in.dims.empty() || w.dims.size() != 2) {
+    throw std::runtime_error(
+        "WebGPU linear: input must have rank >= 1 and weight must be 2D");
   }
-  const uint32_t M = static_cast<uint32_t>(in.dims[0]);
-  const uint32_t K = static_cast<uint32_t>(in.dims[1]);
-  const uint32_t N = static_cast<uint32_t>(w.dims[0]);
-  if (static_cast<uint32_t>(w.dims[1]) != K) {
-    throw std::runtime_error("WebGPU linear: weight.dims[1] != input.dims[1]");
+  const int64_t k_dim = in.dims.back();
+  const int64_t n_dim = w.dims[0];
+  if (k_dim <= 0 || n_dim <= 0 || w.dims[1] != k_dim) {
+    throw std::runtime_error(
+        "WebGPU linear: invalid weight or input K dimension");
   }
-  if (M == 0 || N == 0 || K == 0) {
-    throw std::runtime_error("WebGPU linear: M, N, or K == 0");
+  const uint64_t input_numel = utils::numel_of(in.dims);
+  if (input_numel % static_cast<uint64_t>(k_dim) != 0u) {
+    throw std::runtime_error(
+        "WebGPU linear: input numel is not a multiple of K");
+  }
+  const uint64_t m_dim = input_numel / static_cast<uint64_t>(k_dim);
+  if (m_dim == 0u || m_dim > std::numeric_limits<uint32_t>::max() ||
+      static_cast<uint64_t>(n_dim) > std::numeric_limits<uint32_t>::max() ||
+      static_cast<uint64_t>(k_dim) > std::numeric_limits<uint32_t>::max()) {
+    throw std::runtime_error("WebGPU linear: M, N, or K exceeds uint32 range");
+  }
+  const uint32_t M = static_cast<uint32_t>(m_dim);
+  const uint32_t K = static_cast<uint32_t>(k_dim);
+  const uint32_t N = static_cast<uint32_t>(n_dim);
+
+  std::vector<int64_t> expected_out_dims = in.dims;
+  expected_out_dims.back() = n_dim;
+  if (out.dims != expected_out_dims) {
+    throw std::runtime_error("WebGPU linear: output shape mismatch");
   }
 
-  const uint64_t outputs = static_cast<uint64_t>(M) * static_cast<uint64_t>(N);
-  if (in.nbytes != static_cast<uint64_t>(M) * K * sizeof(float) ||
+  const uint64_t outputs = static_cast<uint64_t>(M) * N;
+  if (in.nbytes != input_numel * sizeof(float) ||
       w.nbytes != static_cast<uint64_t>(N) * K * sizeof(float) ||
       out.nbytes != outputs * sizeof(float)) {
     throw std::runtime_error("WebGPU linear: fp32-only (byte-size mismatch)");
@@ -142,6 +161,9 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       [in_id, out_id, M, N, K, has_bias, dispatch_x, dispatch_idx, params_buf](
           WebGPUGraph& g) {
         const auto& d = g.cur_dims(in_id);
+        if (d.empty() || d.back() != static_cast<int64_t>(K)) {
+          throw std::runtime_error("WebGPU linear: live K changed");
+        }
         const uint64_t numel = utils::numel_of(d);
         if (numel % static_cast<uint64_t>(K) != 0u) {
           throw std::runtime_error(
@@ -158,8 +180,9 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         g.dispatch_at(dispatch_idx).workgroup_count_x = dispatch_x;
         g.dispatch_at(dispatch_idx).workgroup_count_y =
             (m + kTile - 1u) / kTile;
-        g.set_cur_dims(
-            out_id, {static_cast<int64_t>(m), static_cast<int64_t>(N)});
+        std::vector<int64_t> output_dims = d;
+        output_dims.back() = static_cast<int64_t>(N);
+        g.set_cur_dims(out_id, output_dims);
       });
 
   graph.own_uniform_buffer(uniform_buffer);

backends/webgpu/runtime/ops/where/Where.cpp

@@ -75,7 +75,7 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const size_t cond_bind_size = (cond_tensor.nbytes + 3) & ~size_t(3);
 
   uint32_t wg_size = utils::clamp_workgroup_size(device, kWhereWorkgroupSizeX);
-  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, out_meta.numel, wg_size, "where");
 
   WGPUConstantEntry wg_size_constant = {};
@@ -120,8 +120,12 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       &wg_size_constant,
       1);
 
-  const size_t dispatch_idx =
-      graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
+  const size_t dispatch_idx = graph.add_dispatch(
+      {bundle.pipeline,
+       bundle.bind_group,
+       workgroup_count.x,
+       "where",
+       workgroup_count.y});
 
   // Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count.
   WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf,
@@ -170,9 +174,10 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     wgpuQueueWriteBuffer(g.queue(), c_buf, 0, &cm, sizeof(cm));
     wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
     wgpuQueueWriteBuffer(g.queue(), bb_buf, 0, &bm, sizeof(bm));
-    g.dispatch_at(dispatch_idx).workgroup_count_x =
-        utils::compute_1d_workgroup_count(
-            g.device(), om.numel, wg_size, "where(resize)");
+    const utils::WgCount wgc = utils::compute_2d_workgroup_count(
+        g.device(), om.numel, wg_size, "where(resize)");
+    g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
+    g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
   };
   graph.add_tensor_resize_hook(cond_id, where_resize);
   graph.add_tensor_resize_hook(a_id, where_resize);

backends/webgpu/runtime/ops/where/where.wgsl

@@ -23,8 +23,10 @@ fn cond_is_true(i: u32) -> bool {
 }
 
 @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/where/where_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from where.wgsl - DO NOT EDIT.
-// wgsl-sha256: 2c5c6491e95822f767920ad82b81b005808ad5550567b85fa98c1521f188dbfc
+// wgsl-sha256: f254e27661bfaa523237e63f9ccd90665daa436e45f769903ee581b006b06576
 inline constexpr const char* kWhereWGSL = R"(
 @group(0) @binding(0) var<storage, read> cond: array<u32>;
 @group(0) @binding(1) var<storage, read> input_a: array<f32>;
@@ -40,8 +40,10 @@ fn cond_is_true(i: u32) -> bool {
 }
 
 @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/test/op_tests/cases.py

@@ -192,6 +192,10 @@
     CONFIGS as _VIEW_CONFIGS,
     ViewModule,
 )
+from executorch.backends.webgpu.test.ops.test_where import (
+    _fp64_golden as _where_fp64_golden,
+    WhereModule,
+)
 
 # rms_norm coverage is exactly the 15 cases the native test covered.
 RMS_NORM_CASES = _CASES
@@ -1332,6 +1336,47 @@ def _layer_norm_suite() -> WebGPUTestSuite:
     )
 
 
+def _where_cond(shape):
+    return (torch.arange(torch.Size(shape).numel()).reshape(shape) % 2) == 0
+
+
+def _where_a(shape):
+    return torch.linspace(-1.0, 1.0, torch.Size(shape).numel()).reshape(shape)
+
+
+def _where_b(shape):
+    return torch.linspace(1.0, -1.0, torch.Size(shape).numel()).reshape(shape)
+
+
+def _where_golden(_module, inputs):
+    return _where_fp64_golden(*inputs)
+
+
+@register_op_test("where")
+def _where_suite() -> WebGPUTestSuite:
+    def case(name, shape):
+        return Case(
+            name=name,
+            construct={},
+            inputs=(
+                InputSpec(shape=shape, gen=_where_cond),
+                InputSpec(shape=shape, gen=_where_a),
+                InputSpec(shape=shape, gen=_where_b),
+            ),
+            golden_fn=_where_golden,
+        )
+
+    return WebGPUTestSuite(
+        module_factory=WhereModule,
+        cases=[
+            case("equal", (4, 8)),
+            case("large_2d_dispatch", (16, 576, 576)),
+        ],
+        atol=1e-4,
+        rtol=1e-3,
+    )
+
+
 from executorch.backends.webgpu.test.ops.test_linear_fp32 import (
     _ramp as _lin_ramp,
     make_linear,
@@ -1340,8 +1385,10 @@ def _layer_norm_suite() -> WebGPUTestSuite:
 
 @register_op_test("linear_fp32")
 def _linear_fp32_suite() -> WebGPUTestSuite:
-    # fp32 linear (BART + DaViT projections); bias + no-bias, and shapes whose
-    # M*N exceeds the 65535 1D ceiling to exercise the 2D-dispatch spill.
+    # fp32 linear (BART, DaViT, and vision-encoder projections); bias +
+    # no-bias, rank-2/3/4 inputs, and shapes whose M*N exceeds the 65535 1D
+    # ceiling to exercise the 2D-dispatch spill.
+    dynamic_batch = torch.export.Dim("linear_rank4_batch", min=1, max=2)
     return WebGPUTestSuite(
         module_factory=make_linear,
         cases=[
@@ -1360,6 +1407,23 @@ def _linear_fp32_suite() -> WebGPUTestSuite:
                 construct={"in_features": 768, "out_features": 768},
                 inputs=(InputSpec(shape=(1, 16, 768), gen=_lin_ramp),),
             ),
+            Case(
+                name="rank4_bias",
+                construct={"in_features": 64, "out_features": 96},
+                inputs=(InputSpec(shape=(2, 3, 4, 64), gen=_lin_ramp),),
+            ),
+            Case(
+                name="rank4_no_bias",
+                construct={"in_features": 64, "out_features": 96, "bias": False},
+                inputs=(InputSpec(shape=(2, 3, 4, 64), gen=_lin_ramp),),
+            ),
+            Case(
+                name="dynamic_rank4_batch_2_to_1",
+                construct={"in_features": 64, "out_features": 96},
+                export_inputs=(InputSpec(shape=(2, 3, 4, 64), gen=_lin_ramp),),
+                inputs=(InputSpec(shape=(1, 3, 4, 64), gen=_lin_ramp),),
+                dynamic_shapes=({0: dynamic_batch},),
+            ),
             Case(
                 name="tall_m",
                 construct={"in_features": 128, "out_features": 64},

extension/wasm/CMakeLists.txt

@@ -50,6 +50,18 @@ list(
   extension_runner_util
 )
 
+if(EXECUTORCH_BUILD_WASM_WEBGPU)
+  if(NOT TARGET webgpu_backend)
+    message(
+      FATAL_ERROR
+        "EXECUTORCH_BUILD_WASM_WEBGPU requires the webgpu_backend target. "
+        "Enable EXECUTORCH_BUILD_WEBGPU and add backends/webgpu before "
+        "extension/wasm."
+    )
+  endif()
+  list(APPEND link_libraries webgpu_backend)
+endif()
+
 add_library(executorch_wasm OBJECT wasm_bindings.cpp)
 
 if(EXECUTORCH_ENABLE_EVENT_TRACER)
@@ -69,6 +81,26 @@ target_link_libraries(
   INTERFACE executorch_kernels
 )
 
+if(EXECUTORCH_BUILD_WASM_WEBGPU)
+  target_compile_definitions(executorch_wasm PUBLIC EXECUTORCH_WASM_WEBGPU)
+  # Delegated execution blocks on wgpuInstanceWaitAny, which emdawnwebgpu
+  # implements with Asyncify.handleAsync; the port aborts in that call when
+  # Asyncify is off. ASYNCIFY_IGNORE_INDIRECT must stay unset: the wait is
+  # reached through the virtual backend interface, an indirect call. Model data
+  # lives in the Wasm heap before delegate upload, so permit growth beyond
+  # Emscripten's 16 MiB default and retain the full wasm32 address-space
+  # ceiling.
+  target_link_options(
+    executorch_wasm PUBLIC -sASYNCIFY -sASYNCIFY_STACK_SIZE=65536
+    -sALLOW_MEMORY_GROWTH=1 -sMAXIMUM_MEMORY=4GB
+  )
+  # The backend reports device and graph failures by throwing, and catches them
+  # at its own boundary; Emscripten needs the flag when compiling and when
+  # linking.
+  target_compile_options(executorch_wasm PUBLIC -fexceptions)
+  target_link_options(executorch_wasm PUBLIC -fexceptions)
+endif()
+
 if(BUILD_TESTING)
   add_subdirectory(test)
 endif()