Update

[ghstack-poisoned]

Author: JulianCloudNTH

.ci/scripts/setup-webgpu-linux-deps.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# 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.
+
+set -ex
+
+# SwiftShader: software Vulkan adapter for GPU-less CI (LunarG SDK not needed).
+install_swiftshader() {
+  _https_amazon_aws=https://ossci-android.s3.amazonaws.com
+  _swiftshader_archive=swiftshader-abe07b943-prebuilt.tar.gz
+  _swiftshader_dir=/tmp/swiftshader
+  mkdir -p $_swiftshader_dir
+
+  _tmp_archive="/tmp/${_swiftshader_archive}"
+
+  curl --silent --show-error --location --fail --retry 3 --retry-all-errors \
+    --output "${_tmp_archive}" "$_https_amazon_aws/${_swiftshader_archive}"
+
+  tar -C "${_swiftshader_dir}" -xzf "${_tmp_archive}"
+
+  export VK_ICD_FILENAMES="${_swiftshader_dir}/swiftshader/build/Linux/vk_swiftshader_icd.json"
+  export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/"
+  export ETVK_USING_SWIFTSHADER=1
+}
+
+install_swiftshader
+bash backends/webgpu/scripts/setup-wgpu-native.sh

.ci/scripts/test_backend.sh

@@ -57,6 +57,14 @@ if [[ "$FLOW" == *vulkan* ]]; then
     EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_VULKAN=ON"
 fi
 
+if [[ "$FLOW" == *webgpu* ]]; then
+    # Setup swiftshader (software Vulkan adapter for GPU-less runners) and wgpu-native,
+    # which are required to build and run the WebGPU delegate.
+    source .ci/scripts/setup-webgpu-linux-deps.sh
+
+    EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_WEBGPU=ON"
+fi
+
 if [[ "$FLOW" == *arm* ]]; then
     if [[ "$SUITE" == "operators" ]]; then
         PYTEST_RETRY_ARGS=(--reruns 2 --reruns-delay 1)

.github/workflows/test-backend-webgpu.yml

@@ -0,0 +1,27 @@
+name: Test WebGPU Backend
+
+on:
+  schedule:
+    - cron: 0 2 * * *
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+      - ciflow/nightly/*
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}--${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+jobs:
+  test-webgpu:
+    uses: ./.github/workflows/_test_backend.yml
+    with:
+      backend: webgpu
+      flows: '["webgpu"]'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 120
+      run-linux: true

CMakeLists.txt

@@ -1056,6 +1056,10 @@ if(EXECUTORCH_BUILD_PYBIND)
     list(APPEND _dep_libs vulkan_backend)
   endif()
 
+  if(EXECUTORCH_BUILD_WEBGPU)
+    list(APPEND _dep_libs webgpu_backend)
+  endif()
+
   # compile options for pybind
   set(_pybind_compile_options
       $<$<CXX_COMPILER_ID:MSVC>:/EHsc

backends/test/suite/flow.py

@@ -117,6 +117,12 @@ def _load_vulkan() -> list[TestFlow]:
     return [VULKAN_TEST_FLOW, VULKAN_STATIC_INT8_PER_CHANNEL_TEST_FLOW]
 
 
+def _load_webgpu() -> list[TestFlow]:
+    from executorch.backends.test.suite.flows.webgpu import WEBGPU_TEST_FLOW
+
+    return [WEBGPU_TEST_FLOW]
+
+
 def _load_openvino() -> list[TestFlow]:
     from executorch.backends.test.suite.flows.openvino import (
         OPENVINO_INT8_TEST_FLOW,
@@ -178,6 +184,7 @@ def all_flows() -> dict[str, TestFlow]:
         + _register_flow(_load_xnnpack, "XNNPACK")
         + _register_flow(_load_coreml, "Core ML")
         + _register_flow(_load_vulkan, "Vulkan")
+        + _register_flow(_load_webgpu, "WebGPU")
         + _register_flow(_load_openvino, "OpenVINO")
         + _register_flow(_load_qnn, "QNN")
         + _register_flow(_load_arm, "ARM")

backends/test/suite/flows/webgpu.py

@@ -0,0 +1,20 @@
+# 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.
+
+from executorch.backends.test.suite.flow import TestFlow
+from executorch.backends.webgpu.test.tester import WebGPUTester
+
+
+def _create_webgpu_flow() -> TestFlow:
+    return TestFlow(
+        "webgpu",
+        backend="webgpu",
+        tester_factory=WebGPUTester,
+        skip_patterns=["float16", "float64"],  # Not supported in swiftshader
+    )
+
+
+WEBGPU_TEST_FLOW = _create_webgpu_flow()

backends/webgpu/CMakeLists.txt

@@ -21,7 +21,7 @@ if(NOT TARGET vulkan_schema)
   # target), but vulkan_schema is unconditionally defined.
   add_subdirectory(
     ${CMAKE_CURRENT_SOURCE_DIR}/../vulkan
-    ${CMAKE_CURRENT_BINARY_DIR}/_vulkan_schema
+    ${CMAKE_CURRENT_BINARY_DIR}/_vulkan_schema EXCLUDE_FROM_ALL
   )
 endif()
 

backends/webgpu/__init__.py

@@ -0,0 +1,5 @@
+# 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.

backends/webgpu/runtime/WebGPUDevice.cpp

@@ -121,7 +121,13 @@ WebGPUContext create_webgpu_context() {
   device_cb.callback = on_device_request;
   device_cb.userdata1 = &device_result;
 
+  // Request the adapter's full limits; software adapters default many to 0.
+  WGPULimits supported_limits = {};
   WGPUDeviceDescriptor device_desc = {};
+  if (wgpuAdapterGetLimits(ctx.adapter, &supported_limits) ==
+      WGPUStatus_Success) {
+    device_desc.requiredLimits = &supported_limits;
+  }
   device_desc.uncapturedErrorCallbackInfo.callback = on_device_error;
 
   wgpuAdapterRequestDevice(ctx.adapter, &device_desc, device_cb);

backends/webgpu/runtime/ops/add/BinaryOp.cpp

@@ -12,6 +12,7 @@
 
 #include <webgpu/webgpu.h>
 
+#include <algorithm>
 #include <cmath>
 #include <cstring>
 
@@ -50,6 +51,26 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   uint32_t num_elements =
       static_cast<uint32_t>(out_tensor.nbytes / sizeof(float));
 
+  // Clamp the workgroup size to the device limit (SwiftShader caps at 128).
+  WGPULimits limits = {};
+  uint32_t device_max =
+      wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
+          limits.maxComputeInvocationsPerWorkgroup > 0
+      ? limits.maxComputeInvocationsPerWorkgroup
+      : kBinaryAddWorkgroupSize;
+  uint32_t wg_size = std::min(kBinaryAddWorkgroupSize, device_max);
+  uint32_t workgroup_count = (num_elements + wg_size - 1) / wg_size;
+
+  // Validate the 1D dispatch limit before allocating any GPU objects.
+  if (workgroup_count > 65535u) {
+    throw std::runtime_error(
+        "WebGPU add: workgroup count exceeds the 1D dispatch limit (65535)");
+  }
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
   // Create uniform buffer for params
   AddParams params = {};
   params.num_elements = num_elements;
@@ -115,6 +136,8 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   pipeline_desc.layout = pipeline_layout;
   pipeline_desc.compute.module = shader;
   pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
   WGPUComputePipeline pipeline =
       wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
 
@@ -146,16 +169,14 @@ void add_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   bg_desc.entries = bg_entries;
   WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
 
-  uint32_t workgroup_count =
-      (num_elements + kBinaryAddWorkgroupSize - 1) / kBinaryAddWorkgroupSize;
-
   graph.add_dispatch({pipeline, bind_group, workgroup_count});
 
   // Release intermediate objects (pipeline and bind_group are kept by dispatch)
   wgpuShaderModuleRelease(shader);
   wgpuBindGroupLayoutRelease(bgl);
   wgpuPipelineLayoutRelease(pipeline_layout);
-  // uniform_buffer is kept alive by the bind group
+  // Drop our ref; the bind group keeps the uniform buffer alive until release.
+  wgpuBufferRelease(uniform_buffer);
 }
 
 } // namespace