[ExecuTorch][WebGPU] Add max_pool2d op

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})

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 <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/max_pool2d/max_pool2d_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <vector>
+
+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<int>& 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<int>& 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<int64_t> 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<uint32_t>(in.dims[0]);
+  const uint32_t C = static_cast<uint32_t>(in.dims[1]);
+  const uint32_t IH = static_cast<uint32_t>(in.dims[2]);
+  const uint32_t IW = static_cast<uint32_t>(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<uint32_t>(oh_num / int64_t(sH)) + 1;
+  const uint32_t OW = static_cast<uint32_t>(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<uint32_t>(out.dims[0]) != N ||
+      static_cast<uint32_t>(out.dims[1]) != C ||
+      static_cast<uint32_t>(out.dims[2]) != OH ||
+      static_cast<uint32_t>(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, &params, 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

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<storage, read_write> out_vals: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+@group(0) @binding(3) var<storage, read_write> out_idx: array<i32>;
+
+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<u32>) {
+  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;
+  }
+}

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 <cstdint>
+
+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<storage, read_write> out_vals: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+@group(0) @binding(3) var<storage, read_write> out_idx: array<i32>;
+
+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<u32>) {
+  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