diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index fca12ffaf3f..1799f797651 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -61,6 +61,7 @@ set(WEBGPU_SRCS runtime/ops/addmm/Addmm.cpp runtime/ops/constant_pad_nd/ConstantPadNd.cpp runtime/ops/upsample_nearest2d/UpsampleNearest2d.cpp + runtime/ops/max_pool2d/MaxPool2d.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp new file mode 100644 index 00000000000..682eff61197 --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp @@ -0,0 +1,198 @@ +/* + * 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 PoolParams { + uint32_t N, C, IH, IW, OH, OW; + uint32_t kH, kW, sH, sW, pH, pW, dH, dW; + uint32_t _p0, _p1; +}; +static_assert(sizeof(PoolParams) == 64, "PoolParams must be 64 bytes"); + +// out_list = ValueList[values, indices] (mirrors Vulkan Pool.cpp). +void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 6) { + throw std::runtime_error("WebGPU max_pool2d: expected >=6 args"); + } + const int in_id = args.at(0); + const int kernel_id = args.at(1); + const int stride_id = args.at(2); + const int padding_id = args.at(3); + const int dilation_id = args.at(4); + const int out_list_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + if (graph.get_value_type(out_list_id) != WebGPUGraph::ValueType::ValueList) { + throw std::runtime_error("WebGPU max_pool2d: out is not a ValueList"); + } + const std::vector& outs = graph.get_value_list(out_list_id); + if (outs.empty()) { + throw std::runtime_error("WebGPU max_pool2d: empty out ValueList"); + } + const int values_id = outs.at(0); // [0]=values, [1]=indices + const bool has_indices = outs.size() > 1 && + graph.get_value_type(outs.at(1)) == WebGPUGraph::ValueType::Tensor; + if (has_indices) { + const auto& idx_t = graph.get_tensor(outs.at(1)); + if (idx_t.dims != graph.get_tensor(values_id).dims) { + throw std::runtime_error( + "WebGPU max_pool2d: indices output shape must match values output"); + } + // The WGSL kernel writes int32 indices (mirrors Vulkan's Pool.cpp, which + // also computes/stores 32-bit indices despite ATen's int64 schema) — throw + // rather than silently mis-stride the buffer if this graph's indices + // tensor ever turns out to be int64 (8 bytes/elem) instead. + const uint64_t idx_numel = utils::numel_of(idx_t.dims); + if (idx_t.nbytes != idx_numel * sizeof(int32_t)) { + throw std::runtime_error( + "WebGPU max_pool2d: indices output must be int32 (4 bytes/elem); " + "got a different byte width, would mis-stride the i32 kernel write"); + } + } + + const auto& in = graph.get_tensor(in_id); + const auto& out = graph.get_tensor(values_id); + if (in.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU max_pool2d: expected 4D in/out"); + } + + // kernel/stride/padding/dilation are int lists; PyTorch broadcasts a single + // value to both spatial dims. stride defaults to kernel_size when empty. + uint32_t kH, kW, sH, sW, pH, pW, dH, dW; + utils::parse_hw( + graph.get_int_list(kernel_id), kH, kW, "max_pool2d", "kernel_size"); + const std::vector stride_v = graph.get_int_list(stride_id); + if (stride_v.empty()) { + sH = kH; + sW = kW; + } else { + utils::parse_hw(stride_v, sH, sW, "max_pool2d", "stride"); + } + utils::parse_hw( + graph.get_int_list(padding_id), pH, pW, "max_pool2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "max_pool2d", "dilation"); + + 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]); + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU max_pool2d: zero stride"); + } + + const int64_t oh_num = + int64_t(IH) + 2 * int64_t(pH) - int64_t(dH) * (int64_t(kH) - 1) - 1; + const int64_t ow_num = + int64_t(IW) + 2 * int64_t(pW) - int64_t(dW) * (int64_t(kW) - 1) - 1; + if (oh_num < 0 || ow_num < 0) { + throw std::runtime_error("WebGPU max_pool2d: kernel larger than input"); + } + const uint32_t OH = static_cast(oh_num / int64_t(sH)) + 1; + const uint32_t OW = static_cast(ow_num / int64_t(sW)) + 1; + + // Validate against the serialized values output [N, C, OH, OW] (loud-fail if + // the arg interpretation is wrong, e.g. ceil_mode or a different layout). + if (static_cast(out.dims[0]) != N || + static_cast(out.dims[1]) != C || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU max_pool2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "max_pool2d", "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, "max_pool2d"), + kMaxPool2dWorkgroupSizeX, + "max_pool2d"); + + PoolParams params = {}; + params.N = N; + params.C = C; + params.IH = IH; + params.IW = IW; + params.OH = OH; + params.OW = OW; + params.kH = kH; + params.kW = kW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(PoolParams)); + graph.add_uniform_buffer_bytes(sizeof(PoolParams)); + + auto grid_constants = utils::make_grid_constants(grid); + WGPUConstantEntry constants[3] = {}; + constants[0] = grid_constants[0]; + constants[1] = grid_constants[1]; + constants[2].key = {"write_indices", WGPU_STRLEN}; + constants[2].value = has_indices ? 1.0 : 0.0; + + // write_indices==0 -> the shader never stores to out_idx, so a tiny dummy + // buffer is safe (mirrors NativeLayerNorm.cpp's dummy_affine pattern). + utils::OptionalBinding idx = utils::make_optional_binding( + device, + has_indices, + has_indices ? graph.get_tensor(outs.at(1)).buffer : nullptr, + has_indices ? graph.get_tensor(outs.at(1)).nbytes : 0); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kMaxPool2dWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(PoolParams)}, + {3, WGPUBufferBindingType_Storage, idx.buffer, idx.nbytes}, + }, + constants, + 3); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y); + + wgpuBufferRelease(uniform_buffer); + if (idx.owned_dummy != nullptr) { + wgpuBufferRelease(idx.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.max_pool2d_with_indices.default, max_pool2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl new file mode 100644 index 00000000000..3d7c67ddf9a --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d.wgsl @@ -0,0 +1,79 @@ +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill +override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) + +// max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per +// output element (n, c, oh, ow); gather the window, take the max. General +// stride/pad/dilation. Argmax is ALWAYS tracked (mirrors Vulkan Pool.cpp, which +// computes indices unconditionally and only gates the final write) so the +// values-only and with-indices paths share one kernel; write_indices==0 skips +// only the out_idx store, not the tracking, keeping both paths bit-identical on +// out_vals. Indices are the flat (ih*IW+iw) spatial-plane offset (matches +// torch.nn.functional.max_pool2d(return_indices=True)'s documented convention, +// NOT an absolute offset into the full NCHW tensor). Pad cells are skipped +// (init -inf); out_idx is left at its bound (dummy, when unused) value. +@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 = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; // * IW added per-row below + + var best: f32 = -3.4e38; // large finite negative max-init (Dawn/Tint rejects -3.40282347e38: > f32 max); matches q4gsw_dq8ca convention + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} diff --git a/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h new file mode 100644 index 00000000000..1dbe0571e40 --- /dev/null +++ b/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h @@ -0,0 +1,103 @@ +/* + * 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 max_pool2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 2215b7ba725face441baac5eca8e2133458b414377868354c3fda9046a88109b +inline constexpr const char* kMaxPool2dWGSL = R"( +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill +override write_indices: u32 = 0u; // 1 = also write out_idx (compile-time gate, mirrors Vulkan) + +// max_pool2d (values [+ optional indices]), NCHW row-major, fp32. One thread per +// output element (n, c, oh, ow); gather the window, take the max. General +// stride/pad/dilation. Argmax is ALWAYS tracked (mirrors Vulkan Pool.cpp, which +// computes indices unconditionally and only gates the final write) so the +// values-only and with-indices paths share one kernel; write_indices==0 skips +// only the out_idx store, not the tracking, keeping both paths bit-identical on +// out_vals. Indices are the flat (ih*IW+iw) spatial-plane offset (matches +// torch.nn.functional.max_pool2d(return_indices=True)'s documented convention, +// NOT an absolute offset into the full NCHW tensor). Pad cells are skipped +// (init -inf); out_idx is left at its bound (dummy, when unused) value. +@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 = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; // * IW added per-row below + + var best: f32 = -3.4e38; // large finite negative max-init (Dawn/Tint rejects -3.40282347e38: > f32 max); matches q4gsw_dq8ca convention + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} +)"; + +inline constexpr uint32_t kMaxPool2dWorkgroupSizeX = 256; +inline constexpr uint32_t kMaxPool2dWorkgroupSizeY = 1; +inline constexpr uint32_t kMaxPool2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu