From 5686e57a962ba09a94ad9fa2b2f5baa4cb053ad4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 16 Jul 2026 14:49:19 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../runtime/ops/batch_norm/BatchNorm.cpp | 155 ++++++++++++++++++ .../runtime/ops/batch_norm/batch_norm.wgsl | 41 +++++ .../runtime/ops/batch_norm/batch_norm_wgsl.h | 65 ++++++++ 4 files changed, 262 insertions(+) create mode 100644 backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp create mode 100644 backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl create mode 100644 backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 61843a9d71a..8cda8aef997 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -66,6 +66,7 @@ set(WEBGPU_SRCS runtime/ops/max_pool2d/MaxPool2d.cpp runtime/ops/leaky_relu/LeakyRelu.cpp runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp + runtime/ops/batch_norm/BatchNorm.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp b/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp new file mode 100644 index 00000000000..9b2d84c68bd --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/BatchNorm.cpp @@ -0,0 +1,155 @@ +/* + * 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 BNParams { + uint32_t num_elements; + uint32_t C; + uint32_t HW; + uint32_t has_weight; + uint32_t has_bias; + float eps; + uint32_t _p0; + uint32_t _p1; +}; +static_assert(sizeof(BNParams) == 32, "BNParams must be 32 bytes"); + +// _native_batch_norm_legit_no_training: per-channel affine from running stats. +void batch_norm_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 8) { + throw std::runtime_error("WebGPU batch_norm: expected >=8 args"); + } + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int mean_id = args.at(3); + const int var_id = args.at(4); + const double eps = graph.get_double(args.at(6)); + const std::vector& outs = graph.get_value_list(args.at(args.size() - 1)); + if (outs.empty()) { + throw std::runtime_error("WebGPU batch_norm: empty out ValueList"); + } + const int out_id = outs.at(0); + + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(out_id); + const auto& mean = graph.get_tensor(mean_id); + const auto& var = graph.get_tensor(var_id); + + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU batch_norm: expected 4D in/out"); + } + const uint32_t N = static_cast(in.dims[0]); + const uint32_t C = static_cast(in.dims[1]); + const uint32_t HW = + static_cast(in.dims[2]) * static_cast(in.dims[3]); + + uint64_t out_numel = uint64_t(N) * C * HW; + if (out.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error("WebGPU batch_norm: fp32-only (byte mismatch)"); + } + if (out_numel > 0xFFFFFFFFull) { + throw std::runtime_error("WebGPU batch_norm: output too large"); + } + + const bool has_weight = + graph.get_value_type(weight_id) == WebGPUGraph::ValueType::Tensor; + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + static_cast(out_numel), + kBatchNormWorkgroupSizeX, + "batch_norm"); + + BNParams params = {}; + params.num_elements = static_cast(out_numel); + params.C = C; + params.HW = HW; + params.has_weight = has_weight ? 1u : 0u; + params.has_bias = has_bias ? 1u : 0u; + params.eps = static_cast(eps); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(BNParams)); + graph.add_uniform_buffer_bytes(sizeof(BNParams)); + + utils::OptionalBinding wbind = utils::make_optional_binding( + device, + has_weight, + has_weight ? graph.get_tensor(weight_id).buffer : nullptr, + has_weight ? graph.get_tensor(weight_id).nbytes : 0); + utils::OptionalBinding bbind = utils::make_optional_binding( + device, + has_bias, + has_bias ? graph.get_tensor(bias_id).buffer : nullptr, + has_bias ? graph.get_tensor(bias_id).nbytes : 0); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBatchNormWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + wbind.buffer, + wbind.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + bbind.buffer, + bbind.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, mean.buffer, mean.nbytes}, + {5, WGPUBufferBindingType_ReadOnlyStorage, var.buffer, var.nbytes}, + {6, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(BNParams)}, + }, + constants.data(), + constants.size()); + + WebGPUDispatch dispatch{}; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; + dispatch.workgroup_count_x = grid.count_x; + dispatch.workgroup_count_y = grid.count_y; + graph.add_dispatch(dispatch); + + wgpuBufferRelease(uniform_buffer); + if (wbind.owned_dummy != nullptr) { + wgpuBufferRelease(wbind.owned_dummy); + } + if (bbind.owned_dummy != nullptr) { + wgpuBufferRelease(bbind.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + aten._native_batch_norm_legit_no_training.default, batch_norm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl b/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl new file mode 100644 index 00000000000..015f6c7f8a2 --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/batch_norm.wgsl @@ -0,0 +1,41 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var var_: array; + +struct Params { + num_elements: u32, + C: u32, + HW: u32, + has_weight: u32, + has_bias: u32, + eps: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +// Inference batch_norm NCHW fp32: per-channel affine from running stats. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let c = (i / params.HW) % params.C; + let inv = inverseSqrt(var_[c] + params.eps); + var g: f32 = 1.0; + if (params.has_weight == 1u) { + g = weight[c]; + } + var b: f32 = 0.0; + if (params.has_bias == 1u) { + b = bias[c]; + } + out[i] = (inp[i] - mean[c]) * inv * g + b; +} diff --git a/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h b/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h new file mode 100644 index 00000000000..f78d8cdc765 --- /dev/null +++ b/backends/webgpu/runtime/ops/batch_norm/batch_norm_wgsl.h @@ -0,0 +1,65 @@ +/* + * 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 batch_norm.wgsl - DO NOT EDIT. +// wgsl-sha256: ac7208cec0fe1454698b9b16985172c1ef34c994880b568238eaeaf1d62c0342 +inline constexpr const char* kBatchNormWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var var_: array; + +struct Params { + num_elements: u32, + C: u32, + HW: u32, + has_weight: u32, + has_bias: u32, + eps: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +// Inference batch_norm NCHW fp32: per-channel affine from running stats. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.num_elements) { + return; + } + let c = (i / params.HW) % params.C; + let inv = inverseSqrt(var_[c] + params.eps); + var g: f32 = 1.0; + if (params.has_weight == 1u) { + g = weight[c]; + } + var b: f32 = 0.0; + if (params.has_bias == 1u) { + b = bias[c]; + } + out[i] = (inp[i] - mean[c]) * inv * g + b; +} +)"; + +inline constexpr uint32_t kBatchNormWorkgroupSizeX = 256; +inline constexpr uint32_t kBatchNormWorkgroupSizeY = 1; +inline constexpr uint32_t kBatchNormWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu