From d165dd4ed7c295cfff6cf9769cfc7e77badd7fde Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 16 Jul 2026 14:49:09 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../UpsampleBilinear2d.cpp | 132 ++++++++++++++++++ .../upsample_bilinear2d.wgsl | 68 +++++++++ .../upsample_bilinear2d_wgsl.h | 92 ++++++++++++ 4 files changed, 293 insertions(+) create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl create mode 100644 backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index e0d62e51179..61843a9d71a 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -65,6 +65,7 @@ set(WEBGPU_SRCS runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp runtime/ops/max_pool2d/MaxPool2d.cpp runtime/ops/leaky_relu/LeakyRelu.cpp + runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp b/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp new file mode 100644 index 00000000000..bb4b3d88e12 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/UpsampleBilinear2d.cpp @@ -0,0 +1,132 @@ +/* + * 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 align_corners; + uint32_t _p0; +}; +static_assert(sizeof(UpsampleParams) == 32, "UpsampleParams must be 32 bytes"); + +// aten.upsample_bilinear2d.vec: bilinear NCHW resize; align_corners src-index. +void upsample_bilinear2d_impl( + WebGPUGraph& graph, + const std::vector& args) { + if (args.size() < 3) { + throw std::runtime_error("WebGPU upsample_bilinear2d: expected >=3 args"); + } + const int in_id = args.at(0); + const int align_corners_id = args.at(2); + 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_bilinear2d: 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_bilinear2d: N/C mismatch"); + } + if (IH == 0 || IW == 0 || OH == 0 || OW == 0) { + throw std::runtime_error("WebGPU upsample_bilinear2d: zero spatial dim"); + } + + uint64_t out_numel = uint64_t(N) * C * OH * OW; + if (out.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU upsample_bilinear2d: fp32-only (byte mismatch)"); + } + if (out_numel > 0xFFFFFFFFull) { + throw std::runtime_error("WebGPU upsample_bilinear2d: output too large"); + } + + const uint32_t align_corners = graph.get_bool(align_corners_id) ? 1u : 0u; + + utils::DispatchGrid grid = utils::compute_dispatch_grid( + device, + static_cast(out_numel), + kUpsampleBilinear2dWorkgroupSizeX, + "upsample_bilinear2d"); + + UpsampleParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + params.align_corners = align_corners; + + 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, + kUpsampleBilinear2dWGSL, + { + {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()); + + 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); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.upsample_bilinear2d.vec, upsample_bilinear2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl new file mode 100644 index 00000000000..7656f3478a9 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl @@ -0,0 +1,68 @@ +@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, + align_corners: u32, + _p0: 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 + +// Bilinear NCHW fp32 upsample; src-index matches ATen upsample_bilinear2d. +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@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 sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} diff --git a/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h new file mode 100644 index 00000000000..bbdc7f77c11 --- /dev/null +++ b/backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d_wgsl.h @@ -0,0 +1,92 @@ +/* + * 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_bilinear2d.wgsl - DO NOT EDIT. +// wgsl-sha256: fa8d89f2ec0e316650315ce39a426efec9ed34daef1fb5b7bc19ae9cea64565a +inline constexpr const char* kUpsampleBilinear2dWGSL = 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, + align_corners: u32, + _p0: 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 + +// Bilinear NCHW fp32 upsample; src-index matches ATen upsample_bilinear2d. +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@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 sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} +)"; + +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeX = 256; +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeY = 1; +inline constexpr uint32_t kUpsampleBilinear2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu