[ExecuTorch][WebGPU] Add constant_pad_nd op (#20857)

JCNTH · web-flow · commit bd00ba502202 · 2026-07-23T08:04:17.000-07:00
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20868 * #20996 * #20995 * #20994 * #20993 * #20992 * #20991 * #20990 * #20989 * #20988 * #20987 * #20986 * #20876 * #20875 * #20874 * #20873 * #20872 * #20871 * #20866 * #20865 * #20864 * #20863 * #20862 * #20861 * #20860 * #20859 * #20858 * __->__ #20857 * #20856 * #20855 * #20854 * #20852 * #20851 * #20850 * #20849 * #20848 * #20846 * #20845 * #20844 * #20843 * #20842 **Adds `aten.constant_pad_nd.default` to the WebGPU backend, unblocking the DaViT window-padding path in vision models.** **Problem** — the backend had no `constant_pad_nd` handler, so any graph that pads a tensor (e.g. DaViT's window partitioning) could not fully delegate to WebGPU and threw at runtime. **Solution** — a single gather/fill compute kernel: Before — no handler; `aten.constant_pad_nd.default` unsupported at runtime. After — one thread per output element gathers the source element when its coordinates land inside the input, otherwise writes the constant fill `value`. **Implementation**: - The handler right-aligns the (rank 1..4) dims into fixed `vec4<u32>` params (`out_dims`, `in_dims`, `left`); leading slots get extent 1 / pad 0, so the WGSL is rank-agnostic and always iterates 4 dims. - The `pad` `IntList` is reversed-dim (innermost-first `(left, right)` pairs); the handler expands it to per-dim `left`/`right`, then validates `out.dims[d] == in.dims[d] + left[d] + right[d]` before any buffer allocation (loud-fail, no leak-on-throw). - The kernel decodes each output element's 4D coords (last dim fastest), subtracts each dim's `left` pad as an unsigned wrapping subtract (a negative coord wraps to a huge value and is rejected by the `< in_dims` bound check); if all four coords are in-bounds it copies `inp[flat_in]`, else it writes `value` — a pure copy/fill, so bit-exact. - The fill `value` is read via `utils::scalar_or` (a `Scalar` may serialize as `Int` or `Double`), defaulting to `0`. - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid` (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 workgroup-count ceiling; the `stride_x` override lets the shader decode `i = gid.y*stride_x + gid.x`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Pad.cpp` (same reversed-dim `(before, after)` pad convention and `constant_pad_nd` resize logic). **Constraints** — fp32 only (`nbytes == numel*4` guard); rank 1..4; `pad` must be even-length and no longer than the rank; the output element count must fit `u32` (`<= 2^32`). Co-authored-with: Claude Code. @exported-using-ghexport Differential Revision: [D110836672](https://our.internmc.facebook.com/intern/diff/D110836672/) Differential Revision: [D110836672](https://our.internmc.facebook.com/intern/diff/D110836672)
diff --git a/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp b/backends/webgpu/runtime/ops/constant_pad_nd/ConstantPadNd.cpp
@@ -0,0 +1,156 @@
+/*
+ * 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/constant_pad_nd/constant_pad_nd_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <array>
+#include <cstdint>
+#include <stdexcept>
+
+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<int>& 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<int64_t>& 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<int64_t, 4> left = {0, 0, 0, 0};
+  std::array<int64_t, 4> 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++) {
+    // The kernel's pad params are u32; a negative pad (cropping) would wrap
+    // into a huge offset, so reject it loudly rather than gather out of bounds.
+    if (left[d] < 0 || right[d] < 0) {
+      throw std::runtime_error(
+          "WebGPU constant_pad_nd: negative pad (cropping) not supported");
+    }
+    const int64_t expect = in.dims[d] + left[d] + right[d];
+    if (expect < 0 || static_cast<int64_t>(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");
+  utils::check_fp32(in, "constant_pad_nd", "input");
+
+  // 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<uint32_t>(out.dims[d]);
+    params.in_dims[off + d] = static_cast<uint32_t>(in.dims[d]);
+    params.left[off + d] = static_cast<uint32_t>(left[d]);
+  }
+  params.out_numel = static_cast<uint32_t>(out_numel);
+  params.value = value;
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, 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
@@ -0,0 +1,63 @@
+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+
+// 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<u32>,
+  in_dims: vec4<u32>,
+  left: vec4<u32>,
+  out_numel: u32,
+  value: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(2) var<uniform> 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<u32>) {
+  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
@@ -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 <cstdint>
+
+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<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+
+// 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<u32>,
+  in_dims: vec4<u32>,
+  left: vec4<u32>,
+  out_numel: u32,
+  value: f32,
+  _p0: u32,
+  _p1: u32,
+}
+@group(0) @binding(2) var<uniform> 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<u32>) {
+  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
