[ExecuTorch][WebGPU] Add rms_norm op

Pull Request resolved: pytorch#19963

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. 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.

The weight constant is uploaded via the named-data path added in the parent diff.
ghstack-source-id: 389206169
@exported-using-ghexport

Differential Revision: [D106887028](https://our.internmc.facebook.com/intern/diff/D106887028/)

Author: JulianCloudNTH

backends/webgpu/CMakeLists.txt

@@ -26,9 +26,13 @@ if(NOT TARGET vulkan_schema)
 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/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})
@@ -116,4 +120,35 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
 
   target_compile_options(webgpu_native_test PRIVATE -fexceptions)
   set_property(TARGET webgpu_native_test PROPERTY CXX_STANDARD 17)
+
+  add_executable(webgpu_rms_norm_test test/native/test_rms_norm.cpp)
+
+  target_include_directories(
+    webgpu_rms_norm_test PRIVATE $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
+                                 "${WGPU_NATIVE_DIR}/include"
+  )
+
+  target_link_libraries(
+    webgpu_rms_norm_test
+    PRIVATE webgpu_backend
+            wgpu_native
+            executorch_core
+            extension_module_static
+            extension_data_loader
+            extension_tensor
+            portable_kernels
+            portable_ops_lib
+  )
+
+  if(APPLE)
+    target_link_libraries(
+      webgpu_rms_norm_test PRIVATE "-framework Metal" "-framework QuartzCore"
+                                   "-framework CoreGraphics"
+    )
+  else()
+    target_link_libraries(webgpu_rms_norm_test PRIVATE dl m pthread)
+  endif()
+
+  target_compile_options(webgpu_rms_norm_test PRIVATE -fexceptions)
+  set_property(TARGET webgpu_rms_norm_test PROPERTY CXX_STANDARD 17)
 endif()

backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp

@@ -0,0 +1,195 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform layout matching the WGSL Params struct (16-byte aligned).
+struct RmsNormParams {
+  uint32_t num_rows;
+  uint32_t row_width;
+  float epsilon;
+  uint32_t _pad;
+};
+static_assert(sizeof(RmsNormParams) == 16, "RmsNormParams must be 16 bytes");
+
+void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  // et_vk.rms_norm.default args: [in, weight, eps, out]
+  const int in_id = args.at(0);
+  const int weight_id = args.at(1);
+  const int eps_id = args.at(2);
+  const int out_id = args.at(3);
+
+  WGPUDevice device = graph.device();
+
+  // Get epsilon (Double from a Python float; defaults to float32 eps)
+  float epsilon = std::numeric_limits<float>::epsilon();
+  if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Double) {
+    epsilon = static_cast<float>(graph.get_double(eps_id));
+  } else if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Int) {
+    epsilon = static_cast<float>(graph.get_int(eps_id));
+  }
+
+  // row_width = last dim; num_rows = product of the rest (PyTorch NCHW order)
+  const auto& in_tensor = graph.get_tensor(in_id);
+  if (in_tensor.dims.empty() || in_tensor.nbytes == 0) {
+    throw std::runtime_error("WebGPU rms_norm: empty input");
+  }
+  const uint32_t row_width = static_cast<uint32_t>(in_tensor.dims.back());
+  if (row_width == 0) {
+    throw std::runtime_error("WebGPU rms_norm: zero row width");
+  }
+  uint64_t in_numel = 1;
+  for (int64_t d : in_tensor.dims) {
+    in_numel *= static_cast<uint64_t>(d);
+  }
+  // fp32-only shader: bail if the bytes don't match an fp32 element count.
+  if (in_tensor.nbytes != in_numel * sizeof(float)) {
+    throw std::runtime_error("WebGPU rms_norm: fp32-only (byte-size mismatch)");
+  }
+  const uint32_t num_rows = static_cast<uint32_t>(in_numel / row_width);
+  if (num_rows == 0) {
+    throw std::runtime_error("WebGPU rms_norm: zero rows");
+  }
+  // Validate the 1D dispatch limit before allocating any GPU objects.
+  if (num_rows > 65535u) {
+    throw std::runtime_error(
+        "WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)");
+  }
+
+  // Create uniform buffer for params
+  RmsNormParams params = {};
+  params.num_rows = num_rows;
+  params.row_width = row_width;
+  params.epsilon = epsilon;
+
+  WGPUBufferDescriptor uniform_desc = {};
+  uniform_desc.size = sizeof(RmsNormParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
+  void* mapped =
+      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RmsNormParams));
+  std::memcpy(mapped, &params, sizeof(RmsNormParams));
+  wgpuBufferUnmap(uniform_buffer);
+
+  graph.add_uniform_buffer_bytes(sizeof(RmsNormParams));
+
+  // Create shader module from built-in WGSL source
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kRmsNormWGSL, WGPU_STRLEN};
+
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  // Create bind group layout: out (rw) + in/weight (ro storage) + params
+  WGPUBindGroupLayoutEntry entries[4] = {};
+
+  // t_out - storage buffer, read-write
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+
+  // t_in - storage buffer, read-only
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+
+  // t_weight - storage buffer, read-only
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+
+  // params - uniform buffer
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  // Create pipeline layout
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  // Create compute pipeline
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  // Create bind group with actual buffers
+  const auto& out_tensor = graph.get_tensor(out_id);
+  const auto& weight_tensor = graph.get_tensor(weight_id);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = out_tensor.buffer;
+  bg_entries[0].size = out_tensor.nbytes;
+
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = in_tensor.buffer;
+  bg_entries[1].size = in_tensor.nbytes;
+
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = weight_tensor.buffer;
+  bg_entries[2].size = weight_tensor.nbytes;
+
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = uniform_buffer;
+  bg_entries[3].size = sizeof(RmsNormParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  // One workgroup per row (kRmsNormWorkgroupSize threads cooperate per row)
+  static_assert(
+      kRmsNormWorkgroupSize == 64,
+      "must match @workgroup_size and WG_SIZE in rms_norm.wgsl");
+  graph.add_dispatch({pipeline, bind_group, num_rows});
+
+  // Release intermediate objects (pipeline and bind_group are kept by dispatch)
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  // Drop our ref; the bind group keeps the uniform buffer alive until release.
+  wgpuBufferRelease(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.rms_norm.default, rms_norm_impl);
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl

@@ -0,0 +1,75 @@
+// NOTE: This file is for editor/tooling support only. The runtime consumes the
+// inline copy of this shader in `rms_norm_wgsl.h` (kRmsNormWGSL). Keep the two
+// in sync by hand — any edit here must be mirrored there.
+@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
+@group(0) @binding(1) var<storage, read> t_in: array<f32>;
+@group(0) @binding(2) var<storage, read> t_weight: array<f32>;
+
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+const WG_SIZE: u32 = 64u;
+
+var<workgroup> shared_sum: array<f32, WG_SIZE>;
+
+fn reduce_shared(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = WG_SIZE / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      shared_sum[worker_id] = shared_sum[worker_id] + shared_sum[worker_id + stride];
+    }
+    workgroupBarrier();
+    stride = stride >> 1u;
+  }
+}
+
+@compute @workgroup_size(64, 1, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row_idx = wid.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let base = row_idx * params.row_width;
+
+  var local_sq_sum: f32 = 0.0;
+  var x: u32 = worker_id;
+  loop {
+    if (x >= params.row_width) {
+      break;
+    }
+    let v = t_in[base + x];
+    local_sq_sum = local_sq_sum + v * v;
+    x = x + WG_SIZE;
+  }
+
+  shared_sum[worker_id] = local_sq_sum;
+  reduce_shared(worker_id);
+
+  let mean_sq = shared_sum[0] / f32(params.row_width);
+  let rstd = inverseSqrt(mean_sq + params.epsilon);
+
+  x = worker_id;
+  loop {
+    if (x >= params.row_width) {
+      break;
+    }
+    let v = t_in[base + x];
+    let w = t_weight[x];
+    t_out[base + x] = v * rstd * w;
+    x = x + WG_SIZE;
+  }
+}