Add rms_norm op + named-data upload + x86 CI

Summary:
Adds the `et_vk.rms_norm.default` operator to the WebGPU backend (a WGSL compute shader using a cooperative tree reduction, one workgroup per row), fixes constant upload so the op's weight loads correctly, and wires the backend into CI.

The Vulkan serializer that the WebGPU backend reuses stores every non-empty constant (e.g. the rms_norm weight) in the PTE's named-data map with `offset == UINT64_MAX` and a `named_key`, rather than inline in the VK00 blob. `WebGPUGraph::build` previously handled only inline constants, so the weight was never uploaded and the op returned all zeros. `build` now also fetches named-data constants via `NamedDataMap::get_data`, mirroring the path `VulkanBackend` already uses. `aten.add` was unaffected since it has no constant tensors.

The shader mirrors the Vulkan implementation (`backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp`, `backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl`); indexing assumes contiguous fp32 inputs. The handler fails loud (throws, mirroring Vulkan's `VK_CHECK_COND`) on invalid shape/dtype/dispatch-limit conditions, and defaults `eps` to the float32 machine epsilon.

Also adds a simple x86 Linux CI job, mirroring the Vulkan delegate: `backends/test/suite/flows/webgpu.py` plus a `WebGPUTester`, run by `oss/.github/workflows/test-backend-webgpu.yml` on SwiftShader (a software Vulkan adapter, via `wgpu-native`, minimal dependencies, no GPU). Two fixes were needed for SwiftShader's downlevel limits: request the adapter's full `requiredLimits` at device creation (software adapters default storage-buffer limits to 0), and lower the `add` op `workgroup_size` from 256 to 64 (256 exceeded SwiftShader's 128-invocation cap; the Vulkan delegate uses 64).

Differential Revision: D106887028

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

@@ -22,13 +22,15 @@ if(NOT TARGET vulkan_schema)
   add_subdirectory(
     ${CMAKE_CURRENT_SOURCE_DIR}/../vulkan
     ${CMAKE_CURRENT_BINARY_DIR}/_vulkan_schema
+    EXCLUDE_FROM_ALL
   )
 endif()
 
 set(WEBGPU_SRCS
     runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp
     runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp
     runtime/ops/OperatorRegistry.cpp runtime/ops/add/BinaryOp.cpp
+    runtime/ops/rms_norm/RmsNorm.cpp
 )
 
 add_library(webgpu_backend ${WEBGPU_SRCS})

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/WebGPUBackend.cpp

@@ -76,7 +76,7 @@ Result<DelegateHandle*> WebGPUBackend::init(
   }
 
   try {
-    graph->build(flatbuffer_data, constant_data);
+    graph->build(flatbuffer_data, constant_data, context.get_named_data_map());
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU graph build failed: %s", e.what());
     graph->~WebGPUGraph();

backends/webgpu/runtime/WebGPUDevice.cpp

@@ -121,7 +121,12 @@ 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);