From 9075bf25505b5aa22891fc9574f320d856d2beca Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:25:39 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../upsample_nearest2d/UpsampleNearest2d.cpp | 118 ++++++++++++++++++ .../upsample_nearest2d.wgsl | 40 ++++++ .../upsample_nearest2d_wgsl.h | 64 ++++++++++ 4 files changed, 223 insertions(+) create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl create mode 100644 backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 1fff7dc8fdc..fca12ffaf3f 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -60,6 +60,7 @@ set(WEBGPU_SRCS runtime/ops/embedding/Embedding.cpp runtime/ops/addmm/Addmm.cpp runtime/ops/constant_pad_nd/ConstantPadNd.cpp + runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp b/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp new file mode 100644 index 00000000000..0664e53a0ae --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp @@ -0,0 +1,118 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct UpsampleParams { + uint32_t N; + uint32_t C; + uint32_t IH; + uint32_t IW; + uint32_t OH; + uint32_t OW; + uint32_t _p0; + uint32_t _p1; +}; +static_assert(sizeof(UpsampleParams) == 32, "UpsampleParams must be 32 bytes"); + +// OH/OW come from the output tensor's own dims, not the size/scale args. +void upsample_nearest2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 2) { + throw std::runtime_error("WebGPU upsample_nearest2d: expected >=2 args"); + } + const int in_id = args.at(0); + 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); + + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU upsample_nearest2d: 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 IH = static_cast(in.dims[2]); + const uint32_t IW = static_cast(in.dims[3]); + const uint32_t OH = static_cast(out.dims[2]); + const uint32_t OW = static_cast(out.dims[3]); + + if (static_cast(out.dims[0]) != N || + static_cast(out.dims[1]) != C) { + throw std::runtime_error("WebGPU upsample_nearest2d: N/C mismatch"); + } + if (IH == 0 || IW == 0 || OH == 0 || OW == 0) { + throw std::runtime_error("WebGPU upsample_nearest2d: zero spatial dim"); + } + + uint64_t out_numel = utils::check_fp32(out, "upsample_nearest2d", "output"); + + // Up-front (throw before any buffer alloc -> no leak-on-throw). + // 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, "upsample_nearest2d"), + kUpsampleNearest2dWorkgroupSizeX, + "upsample_nearest2d"); + + UpsampleParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(UpsampleParams)); + graph.add_uniform_buffer_bytes(sizeof(UpsampleParams)); + + auto constants = utils::make_grid_constants(grid); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kUpsampleNearest2dWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UpsampleParams)}, + }, + 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.upsample_nearest2d.vec, upsample_nearest2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl new file mode 100644 index 00000000000..7f5310597ca --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d.wgsl @@ -0,0 +1,40 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _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 + +// Nearest-neighbour 2D upsample, NCHW row-major, fp32. One thread per output +// element (n, c, oh, ow). oh -> floor(oh*IH/OH), matching real ATen's +// nearest_neighbor_compute_source_index (UpSample.h) — intentionally NOT +// mirroring Vulkan's upsample_2d.glsl, which uses the half-pixel-center +// formula correct for bilinear but not for non-exact nearest (see diff summary). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} diff --git a/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h new file mode 100644 index 00000000000..89da4f3474c --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_nearest2d/upsample_nearest2d_wgsl.h @@ -0,0 +1,64 @@ +/* + * 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 upsample_nearest2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 3eb5bb5604a34371e493ecc150f02995374ab183fed7e7c4a4e036b5b45563c0 +inline constexpr const char* kUpsampleNearest2dWGSL = R"( +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _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 + +// Nearest-neighbour 2D upsample, NCHW row-major, fp32. One thread per output +// element (n, c, oh, ow). oh -> floor(oh*IH/OH), matching real ATen's +// nearest_neighbor_compute_source_index (UpSample.h) — intentionally NOT +// mirroring Vulkan's upsample_2d.glsl, which uses the half-pixel-center +// formula correct for bilinear but not for non-exact nearest (see diff summary). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} +)"; + +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeX = 256; +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeY = 1; +inline constexpr uint32_t kUpsampleNearest2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu