From 78c8d94122683f175dc444abaee440c7dfd6928b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:25:26 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../ops/constant_pad_nd/ConstantPadNd.cpp | 149 ++++++++++++++++++ .../ops/constant_pad_nd/constant_pad_nd.wgsl | 63 ++++++++ .../constant_pad_nd/constant_pad_nd_wgsl.h | 87 ++++++++++ 4 files changed, 300 insertions(+) create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl create mode 100644 backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 9eece554fe6..1fff7dc8fdc 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -59,6 +59,7 @@ set(WEBGPU_SRCS runtime/ops/et_vk_sdpa/EtVkSdpa.cpp runtime/ops/embedding/Embedding.cpp runtime/ops/addmm/Addmm.cpp + runtime/ops/constant_pad_nd/ConstantPadNd.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp b/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp new file mode 100644 index 00000000000..1b7f1690f37 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp @@ -0,0 +1,149 @@ +/* + * 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 PadParams { + uint32_t out_dims[4]; + uint32_t in_dims[4]; + uint32_t left[4]; + uint32_t out_numel; + float value; + uint32_t _p0; + uint32_t _p1; +}; +static_assert(sizeof(PadParams) == 64, "PadParams must be 64 bytes"); + +// `pad` is reversed-dim (last-dim-first pairs); output right-aligned into 4D. +void constant_pad_nd_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 3) { + throw std::runtime_error("WebGPU constant_pad_nd: expected >=3 args"); + } + const int in_id = args.at(0); + const int pad_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(out_id); + + const size_t nd = in.dims.size(); + if (nd == 0 || nd > 4) { + throw std::runtime_error("WebGPU constant_pad_nd: rank must be 1..4"); + } + if (out.dims.size() != nd) { + throw std::runtime_error("WebGPU constant_pad_nd: in/out rank mismatch"); + } + + if (graph.get_value_type(pad_id) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error("WebGPU constant_pad_nd: pad is not an IntList"); + } + const std::vector& pad = graph.get_int_list(pad_id); + if (pad.size() % 2 != 0) { + throw std::runtime_error("WebGPU constant_pad_nd: pad must be even-length"); + } + const size_t npad = pad.size() / 2; + if (npad > nd) { + throw std::runtime_error("WebGPU constant_pad_nd: pad longer than rank"); + } + + // value scalar (default 0). Vulkan serializes a Scalar as Int or Double. + float value = 0.0f; + if (args.size() >= 4) { + value = utils::scalar_or(graph, args.at(2), 0.0f); + } + + // Per-dim left/right pad (pad list is reversed-dim, from the LAST dim). + std::array left = {0, 0, 0, 0}; + std::array right = {0, 0, 0, 0}; + for (size_t k = 0; k < npad; k++) { + const size_t d = nd - 1 - k; // k-th pad entry -> dim (nd-1-k) + left[d] = pad[2 * k]; + right[d] = pad[2 * k + 1]; + } + + // Validate output dims == in + left + right per dim (loud-fail on a wrong + // pad-list interpretation), before any buffer alloc -> no leak-on-throw. + for (size_t d = 0; d < nd; d++) { + const int64_t expect = in.dims[d] + left[d] + right[d]; + if (expect < 0 || static_cast(out.dims[d]) != expect) { + throw std::runtime_error("WebGPU constant_pad_nd: output shape mismatch"); + } + } + + const uint64_t out_numel = + utils::check_fp32(out, "constant_pad_nd", "output"); + + // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 + // ceiling. stride_x lets the shader decode idx = gid.y*stride_x + gid.x. + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + utils::checked_u32(out_numel, "constant_pad_nd"), + kConstantPadNdWorkgroupSizeX, + "constant_pad_nd"); + + // Right-align dims into [4]: leading (4-nd) entries get extent 1, pad 0. + PadParams params = {}; + for (int s = 0; s < 4; s++) { + params.out_dims[s] = 1; + params.in_dims[s] = 1; + params.left[s] = 0; + } + const size_t off = 4 - nd; + for (size_t d = 0; d < nd; d++) { + params.out_dims[off + d] = static_cast(out.dims[d]); + params.in_dims[off + d] = static_cast(in.dims[d]); + params.left[off + d] = static_cast(left[d]); + } + params.out_numel = static_cast(out_numel); + params.value = value; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(PadParams)); + graph.add_uniform_buffer_bytes(sizeof(PadParams)); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConstantPadNdWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(PadParams)}, + }, + constants.data(), + constants.size()); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.constant_pad_nd.default, constant_pad_nd_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl new file mode 100644 index 00000000000..9d7d4b76142 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd.wgsl @@ -0,0 +1,63 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +// Up to 4D. The handler right-aligns dims into [4] (leading entries = 1, left/ +// right pad = 0 for unpadded/leading dims), so the shader is rank-agnostic and +// always iterates 4 dims. in_dims[d] = input extent, left[d] = that dim's +// left-pad, out_dims[d] = in_dims[d] + left[d] + right[d]. +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// constant_pad_nd, gather form, NCHW row-major fp32. One thread per OUTPUT +// element: decode its 4D coords, subtract each dim's left-pad to get the input +// coord; if ALL input coords are in-bounds -> copy inp[flat_in], else write +// `value`. Pure copy/fill -> bit-exact. (CPU-derisked == torch at 0.) +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.out_numel) { + return; + } + + // decode out coords (last dim fastest) + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + // subtract left pad -> input coord (wrapping subtract; check via < in_dim on + // the unsigned result catches negatives because they wrap to huge values) + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = params.value; + } +} diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h new file mode 100644 index 00000000000..a8380304890 --- /dev/null +++ b/backends/webgpu/runtime/ops/constant_pad_nd/constant_pad_nd_wgsl.h @@ -0,0 +1,87 @@ +/* + * 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 constant_pad_nd.wgsl - DO NOT EDIT. +// wgsl-sha256: 67496c3b851bbe69c64d26bcca415cd4c0de20d555b0ed5cccbe66499a757592 +inline constexpr const char* kConstantPadNdWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +// Up to 4D. The handler right-aligns dims into [4] (leading entries = 1, left/ +// right pad = 0 for unpadded/leading dims), so the shader is rank-agnostic and +// always iterates 4 dims. in_dims[d] = input extent, left[d] = that dim's +// left-pad, out_dims[d] = in_dims[d] + left[d] + right[d]. +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// constant_pad_nd, gather form, NCHW row-major fp32. One thread per OUTPUT +// element: decode its 4D coords, subtract each dim's left-pad to get the input +// coord; if ALL input coords are in-bounds -> copy inp[flat_in], else write +// `value`. Pure copy/fill -> bit-exact. (CPU-derisked == torch at 0.) +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.y * stride_x + gid.x; + if (i >= params.out_numel) { + return; + } + + // decode out coords (last dim fastest) + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + // subtract left pad -> input coord (wrapping subtract; check via < in_dim on + // the unsigned result catches negatives because they wrap to huge values) + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = params.value; + } +} +)"; + +inline constexpr uint32_t kConstantPadNdWorkgroupSizeX = 256; +inline constexpr uint32_t kConstantPadNdWorkgroupSizeY = 1; +inline constexpr uint32_t kConstantPadNdWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu