[ExecuTorch][WebGPU] Add optimized layer_norm op

backends/webgpu/CMakeLists.txt

@@ -53,6 +53,7 @@ set(WEBGPU_SRCS
     runtime/ops/index/Index.cpp
     runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp
     runtime/ops/gelu/Gelu.cpp
+    runtime/ops/layer_norm/NativeLayerNorm.cpp
 )
 
 add_library(webgpu_backend ${WEBGPU_SRCS})

backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp

@@ -0,0 +1,183 @@
+/*
+ * 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/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <limits>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct LayerNormParams {
+  uint32_t num_rows;
+  uint32_t row_width;
+  float epsilon;
+  uint32_t has_affine;
+};
+static_assert(
+    sizeof(LayerNormParams) == 16,
+    "LayerNormParams must be 16 bytes");
+
+// aten.native_layer_norm.default args: [in, normalized_shape, weight, bias,
+// eps, out] where out is a ValueList [out, mean, rstd] (mirrors Vulkan
+// NativeLayerNorm.cpp). normalized_shape (args[1]) only constrains the last
+// dim.
+void native_layer_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int in_id = args.at(0);
+  const int weight_id = args.at(2);
+  const int bias_id = args.at(3);
+  const int eps_id = args.at(4);
+  const int out_list_id = args.at(5);
+
+  if (graph.get_value_type(out_list_id) != WebGPUGraph::ValueType::ValueList) {
+    throw std::runtime_error(
+        "WebGPU native_layer_norm: out is not a ValueList");
+  }
+  const std::vector<int>& outs = graph.get_value_list(out_list_id);
+  if (outs.size() != 3) {
+    throw std::runtime_error(
+        "WebGPU native_layer_norm: expected 3 outputs (out, mean, rstd)");
+  }
+  const int out_id = outs.at(0);
+  const int mean_id = outs.at(1);
+  const int rstd_id = outs.at(2);
+
+  WGPUDevice device = graph.device();
+
+  float epsilon =
+      utils::scalar_or(graph, eps_id, std::numeric_limits<float>::epsilon());
+
+  const auto& in_tensor = graph.get_tensor(in_id);
+  if (in_tensor.dims.empty() || in_tensor.nbytes == 0) {
+    throw std::runtime_error("WebGPU native_layer_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 native_layer_norm: zero row width");
+  }
+  // The shader views t_in/t_out/t_weight/t_bias as vec4<f32> over row_width;
+  // every model in scope (DaViT/BART/Whisper/Voxtral hidden dims 768/1024/
+  // 1280/3072) is always %4==0.
+  if (row_width % 4 != 0) {
+    throw std::runtime_error(
+        "WebGPU native_layer_norm: row_width must be a multiple of 4");
+  }
+  const uint64_t in_numel =
+      utils::check_fp32(in_tensor, "native_layer_norm", "input");
+  const uint32_t num_rows = static_cast<uint32_t>(in_numel / row_width);
+  if (num_rows == 0) {
+    throw std::runtime_error("WebGPU native_layer_norm: zero rows");
+  }
+
+  // Near-square 2D grid of workgroups (1 workgroup = 1 row) past the 65535
+  // per-dimension ceiling; stride_x lets the shader decode a flat row index
+  // as workgroup_id.y * stride_x + workgroup_id.x.
+  utils::DispatchGrid grid =
+      utils::compute_row_dispatch_grid(device, num_rows, "native_layer_norm");
+
+  // weight/bias are optional: aten.native_layer_norm passes None for both when
+  // there is no affine (e.g. the group_norm LN-reframe). When absent, bind
+  // dummy storage buffers and gate the affine in the shader (has_affine == 0).
+  const bool has_affine =
+      graph.get_value_type(weight_id) == WebGPUGraph::ValueType::Tensor &&
+      graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor;
+
+  LayerNormParams params = {};
+  params.num_rows = num_rows;
+  params.row_width = row_width;
+  params.epsilon = epsilon;
+  params.has_affine = has_affine ? 1u : 0u;
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(LayerNormParams));
+  graph.add_uniform_buffer_bytes(sizeof(LayerNormParams));
+
+  const auto& out_tensor = graph.get_tensor(out_id);
+  const auto& mean_tensor = graph.get_tensor(mean_id);
+  const auto& rstd_tensor = graph.get_tensor(rstd_id);
+
+  utils::OptionalBinding weight = utils::make_optional_binding(
+      device,
+      has_affine,
+      has_affine ? graph.get_tensor(weight_id).buffer : nullptr,
+      has_affine ? graph.get_tensor(weight_id).nbytes : 0);
+  utils::OptionalBinding bias = utils::make_optional_binding(
+      device,
+      has_affine,
+      has_affine ? graph.get_tensor(bias_id).buffer : nullptr,
+      has_affine ? graph.get_tensor(bias_id).nbytes : 0);
+
+  WGPUConstantEntry stride_const = {};
+  stride_const.key = {"stride_x", WGPU_STRLEN};
+  stride_const.value = static_cast<double>(grid.stride_x);
+
+  // out(rw,0), in(ro,1), weight(ro,2), bias(ro,3), mean(rw,4), rstd(rw,5),
+  // params(uniform,6).
+  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+      device,
+      kNativeLayerNormWGSL,
+      {
+          {0,
+           WGPUBufferBindingType_Storage,
+           out_tensor.buffer,
+           out_tensor.nbytes},
+          {1,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           in_tensor.buffer,
+           in_tensor.nbytes},
+          {2,
+           WGPUBufferBindingType_ReadOnlyStorage,
+           weight.buffer,
+           weight.nbytes},
+          {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes},
+          {4,
+           WGPUBufferBindingType_Storage,
+           mean_tensor.buffer,
+           mean_tensor.nbytes},
+          {5,
+           WGPUBufferBindingType_Storage,
+           rstd_tensor.buffer,
+           rstd_tensor.nbytes},
+          {6,
+           WGPUBufferBindingType_Uniform,
+           uniform_buffer,
+           sizeof(LayerNormParams)},
+      },
+      &stride_const,
+      1);
+
+  static_assert(
+      kNativeLayerNormWorkgroupSizeX == 64,
+      "must match @workgroup_size and WG_SIZE in native_layer_norm.wgsl");
+  graph.add_dispatch_2d(
+      bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y);
+
+  wgpuBufferRelease(uniform_buffer);
+  if (weight.owned_dummy != nullptr) {
+    wgpuBufferRelease(weight.owned_dummy);
+  }
+  if (bias.owned_dummy != nullptr) {
+    wgpuBufferRelease(bias.owned_dummy);
+  }
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.native_layer_norm.default, native_layer_norm_impl);
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl

@@ -0,0 +1,116 @@
+@group(0) @binding(0) var<storage, read_write> t_out: array<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_weight: array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read> t_bias: array<vec4<f32>>;
+@group(0) @binding(4) var<storage, read_write> t_mean: array<f32>;
+@group(0) @binding(5) var<storage, read_write> t_rstd: array<f32>;
+
+struct Params {
+  num_rows: u32,
+  row_width: u32,
+  epsilon: f32,
+  has_affine: u32,
+}
+@group(0) @binding(6) var<uniform> params: Params;
+
+const WG_SIZE: u32 = 64u;
+
+override stride_x: u32 = 4294967295u; // = count_x; set by host for 2D-spill
+
+// Single-pass, numerically-robust mean+variance via Chan et al.'s parallel
+// Welford: each thread folds its strided vec4<f32> elements into a running
+// (n, mean, M2), then this tree-reduces the WG_SIZE per-thread triples via
+// pairwise merges — no naive E[x^2]-E[x]^2 cancellation risk, de-risked on
+// CPU against large-mean/small-variance activations.
+var<workgroup> shared_n: array<f32, WG_SIZE>;
+var<workgroup> shared_mean: array<f32, WG_SIZE>;
+var<workgroup> shared_m2: array<f32, WG_SIZE>;
+
+fn reduce_shared_welford(worker_id: u32) {
+  workgroupBarrier();
+  var stride: u32 = WG_SIZE / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (worker_id < stride) {
+      let na = shared_n[worker_id];
+      let nb = shared_n[worker_id + stride];
+      let n = na + nb;
+      if (n > 0.0) {
+        let delta = shared_mean[worker_id + stride] - shared_mean[worker_id];
+        shared_mean[worker_id] = shared_mean[worker_id] + delta * (nb / n);
+        shared_m2[worker_id] = shared_m2[worker_id] + shared_m2[worker_id + stride]
+            + delta * delta * (na * nb / n);
+        shared_n[worker_id] = n;
+      }
+    }
+    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.y * stride_x + wid.x;
+  let worker_id = lid.x;
+
+  if (row_idx >= params.num_rows) {
+    return;
+  }
+
+  let row_width4 = params.row_width / 4u;
+  let base4 = row_idx * row_width4;
+
+  var local_n: f32 = 0.0;
+  var local_mean: f32 = 0.0;
+  var local_m2: f32 = 0.0;
+  var x4: u32 = worker_id;
+  loop {
+    if (x4 >= row_width4) {
+      break;
+    }
+    let v4 = t_in[base4 + x4];
+    for (var c: u32 = 0u; c < 4u; c = c + 1u) {
+      local_n = local_n + 1.0;
+      let v = v4[c];
+      let delta = v - local_mean;
+      local_mean = local_mean + delta / local_n;
+      let delta2 = v - local_mean;
+      local_m2 = local_m2 + delta * delta2;
+    }
+    x4 = x4 + WG_SIZE;
+  }
+  shared_n[worker_id] = local_n;
+  shared_mean[worker_id] = local_mean;
+  shared_m2[worker_id] = local_m2;
+  reduce_shared_welford(worker_id);
+  let mean = shared_mean[0];
+  let variance = shared_m2[0] / f32(params.row_width);
+  let rstd = inverseSqrt(variance + params.epsilon);
+
+  if (worker_id == 0u) {
+    t_mean[row_idx] = mean;
+    t_rstd[row_idx] = rstd;
+  }
+  workgroupBarrier();
+
+  x4 = worker_id;
+  loop {
+    if (x4 >= row_width4) {
+      break;
+    }
+    let v4 = t_in[base4 + x4];
+    let normed = (v4 - vec4<f32>(mean)) * rstd;
+    // weight/bias are optional (None when this is the group_norm LN-reframe); the
+    // dummy buffers are not bound then, so gate the affine on has_affine.
+    if (params.has_affine == 1u) {
+      t_out[base4 + x4] = normed * t_weight[x4] + t_bias[x4];
+    } else {
+      t_out[base4 + x4] = normed;
+    }
+    x4 = x4 + WG_SIZE;
+  }
+}