diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 441b572426b..86054bccec6 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/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}) diff --git a/backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp b/backends/webgpu/runtime/ops/layer_norm/NativeLayerNorm.cpp new file mode 100644 index 00000000000..7551d9dcdcf --- /dev/null +++ b/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 +#include +#include +#include + +#include + +#include +#include +#include + +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& 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& 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::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(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 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(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, ¶ms, 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(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 diff --git a/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl new file mode 100644 index 00000000000..8e2656d7a65 --- /dev/null +++ b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm.wgsl @@ -0,0 +1,116 @@ +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array>; +@group(0) @binding(4) var t_mean: array; +@group(0) @binding(5) var t_rstd: array; + +struct Params { + num_rows: u32, + row_width: u32, + epsilon: f32, + has_affine: u32, +} +@group(0) @binding(6) var 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 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 shared_n: array; +var shared_mean: array; +var shared_m2: array; + +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, + @builtin(local_invocation_id) lid: vec3) { + 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(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; + } +} diff --git a/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h new file mode 100644 index 00000000000..14495b7b40e --- /dev/null +++ b/backends/webgpu/runtime/ops/layer_norm/native_layer_norm_wgsl.h @@ -0,0 +1,140 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from native_layer_norm.wgsl - DO NOT EDIT. +// wgsl-sha256: 135a274c6ee14e47f5de1f1255d88425f2b56cfbfa95c82ddb17aa82d725d88b +inline constexpr const char* kNativeLayerNormWGSL = R"( +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array>; +@group(0) @binding(4) var t_mean: array; +@group(0) @binding(5) var t_rstd: array; + +struct Params { + num_rows: u32, + row_width: u32, + epsilon: f32, + has_affine: u32, +} +@group(0) @binding(6) var 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 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 shared_n: array; +var shared_mean: array; +var shared_m2: array; + +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, + @builtin(local_invocation_id) lid: vec3) { + 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(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; + } +} +)"; + +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeX = 64; +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeY = 1; +inline constexpr uint32_t kNativeLayerNormWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu