[ExecuTorch][WebGPU] Add upsample_bilinear2d op

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

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 <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/upsample_bilinear2d/upsample_bilinear2d_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+
+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<int>& 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<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]);
+  const uint32_t OH = static_cast<uint32_t>(out.dims[2]);
+  const uint32_t OW = static_cast<uint32_t>(out.dims[3]);
+
+  if (static_cast<uint32_t>(out.dims[0]) != N ||
+      static_cast<uint32_t>(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<uint32_t>(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, &params, 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

backends/webgpu/runtime/ops/upsample_bilinear2d/upsample_bilinear2d.wgsl

@@ -0,0 +1,68 @@
+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+
+struct Params {
+  N: u32,
+  C: u32,
+  IH: u32,
+  IW: u32,
+  OH: u32,
+  OW: u32,
+  align_corners: u32,
+  _p0: 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
+
+// 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<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 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;
+}

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 <cstdint>
+
+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<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> inp: array<f32>;
+
+struct Params {
+  N: u32,
+  C: u32,
+  IH: u32,
+  IW: u32,
+  OH: u32,
+  OW: u32,
+  align_corners: u32,
+  _p0: 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
+
+// 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<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 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