[ExecuTorch][WebGPU] Run stride == kernel conv_transpose2d as a GEMM (#23880)

Summary:
conv_transpose2d uses a per-output gather kernel. When stride == kernel
(e.g. 2x
upsampling) each output takes exactly one tap, so the op is a GEMM plus
a
depth-to-space scatter.

Add conv_transpose2d_gemm.wgsl for that case (groups == 1, no padding,
dilation
or output_padding). Other shapes, GEMMs too narrow or small to beat the
gather,
and outputs whose GEMM grid would exceed the dispatch limit keep the
gather.
Updates the generated-output manifest snapshot for the new shader.

This change was authored with Claude Code.

Test Plan:
- Op tests transpose2x_tiled (multi-tile), transpose3x (GEMM),
transpose2x,
  transpose_overlap (gather) and transpose2x_grid_limit (heavy, gather
  fallback); all conv2d cases pass, as do the WebGPU Python tests.
- Negative controls: breaking the GEMM path fails its cases; dropping
the grid
  check fails transpose2x_grid_limit.
- Selection thresholds from GPU timings of GEMM vs gather across channel
counts.

Author: digantdesai

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -48,6 +48,7 @@
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_gemm_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h>
@@ -155,7 +156,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 138> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -429,6 +430,13 @@ constexpr std::array<WebGPUShaderInfo, 137> kShaderRegistry = {{
         kConvTranspose2dWorkgroupSizeY,
         kConvTranspose2dWorkgroupSizeZ,
     },
+    {
+        "conv_transpose2d_gemm",
+        kConvTranspose2dGemmWGSL,
+        kConvTranspose2dGemmWorkgroupSizeX,
+        kConvTranspose2dGemmWorkgroupSizeY,
+        kConvTranspose2dGemmWorkgroupSizeZ,
+    },
     {
         "conv_with_clamp",
         kConvWithClampWGSL,

backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp

@@ -13,6 +13,7 @@
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_gemm_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h>
 
 #include <webgpu/webgpu.h>
@@ -50,6 +51,13 @@ static_assert(
     sizeof(ConvParams) == 80,
     "ConvParams must be 80 bytes (16-mult)");
 
+// Below these a 32x32x32 GEMM tile is mostly padding, or the GEMM launches too
+// few tiles to fill the GPU, and the direct or gather kernel is as fast or
+// faster.
+constexpr uint64_t kGemmMinM = 8u;
+constexpr uint64_t kGemmMinK = 16u;
+constexpr uint64_t kGemmMinWorkgroups = 256u;
+
 // Transposed 2D convolution (gather form), folded into the convolution handler
 // (C2: a 2nd WEBGPU_REGISTER_OP(aten.convolution.default) would be silently
 // dropped). weight layout = torch convT [IC, OC/groups, KH, KW]. CPU-derisked
@@ -149,15 +157,44 @@ void conv_transpose2d_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const bool has_bias =
       graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor;
 
+  // Non-overlapping taps (stride == kernel, e.g. 2x upsampling) reduce to a
+  // GEMM (M = OC*KH*KW, N = input pixels, K = IC) plus a depth-to-space
+  // scatter. Other shapes, GEMMs too small to pay off, and GEMM tile grids past
+  // the dispatch limit use the per-output gather (which folds its grid).
+  constexpr uint32_t kConvTranspose2dGemmTile = 32u; // TILE in the shader
+  const uint64_t max_wgs = utils::queried_max_workgroups(device);
+  const uint64_t gemm_m = uint64_t(OC) * KH * KW;
+  const uint64_t gemm_cols =
+      utils::div_up<uint64_t>(uint64_t(B) * IH * IW, kConvTranspose2dGemmTile);
+  const uint64_t gemm_rows =
+      utils::div_up<uint64_t>(gemm_m, kConvTranspose2dGemmTile);
+  const bool use_gemm = groups == 1 && sH == KH && sW == KW && pH == 0 &&
+      pW == 0 && dH == 1 && dW == 1 && opH == 0 && opW == 0 &&
+      gemm_m >= kGemmMinM && IC >= kGemmMinK &&
+      gemm_cols * gemm_rows >= kGemmMinWorkgroups && gemm_cols <= max_wgs &&
+      gemm_rows <= max_wgs;
+
   // Up-front (throw before any buffer alloc -> no leak-on-throw).
-  // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535
-  // ceiling (FpnNeck @1008^2 out_numel exceeds it). stride_x lets the shader
-  // decode i = gid.y*stride_x + gid.x.
-  utils::DispatchGrid grid = utils::compute_dispatch_grid(
-      device,
-      utils::checked_u32(out_numel, "conv_transpose2d"),
-      kConvTranspose2dWorkgroupSizeX,
-      "conv_transpose2d");
+  utils::WgCount gemm_grid = {};
+  utils::DispatchGrid grid = {};
+  if (use_gemm) {
+    gemm_grid = utils::compute_tile_grid_2d(
+        device,
+        utils::checked_u32(
+            uint64_t(B) * uint64_t(IH) * uint64_t(IW), "conv_transpose2d_gemm"),
+        OC * KH * KW,
+        kConvTranspose2dGemmTile,
+        "conv_transpose2d_gemm");
+  } else {
+    // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535
+    // ceiling for large outputs. stride_x lets the shader decode
+    // i = gid.y*stride_x + gid.x.
+    grid = utils::compute_dispatch_grid(
+        device,
+        utils::checked_u32(out_numel, "conv_transpose2d"),
+        kConvTranspose2dWorkgroupSizeX,
+        "conv_transpose2d");
+  }
 
   ConvParams params = {};
   params.B = B;
@@ -187,29 +224,38 @@ void conv_transpose2d_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       has_bias,
       has_bias ? graph.get_tensor(bias_id).buffer : nullptr,
       has_bias ? graph.get_tensor(bias_id).nbytes : 0);
-  auto constants = utils::make_grid_constants(grid);
+  const std::vector<utils::BindingSpec> bindings = {
+      {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
+      {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes},
+      {2, WGPUBufferBindingType_ReadOnlyStorage, weight.buffer, weight.nbytes},
+      {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes},
+      {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(ConvParams)},
+  };
 
-  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-      device,
-      kConvTranspose2dWGSL,
-      {
-          {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes},
-          {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes},
-          {2,
-           WGPUBufferBindingType_ReadOnlyStorage,
-           weight.buffer,
-           weight.nbytes},
-          {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes},
-          {4,
-           WGPUBufferBindingType_Uniform,
-           uniform_buffer,
-           sizeof(ConvParams)},
-      },
-      constants.data(),
-      constants.size());
-
-  graph.add_dispatch_2d(
-      bundle.pipeline, bundle.bind_group, grid.count_x, grid.count_y);
+  if (use_gemm) {
+    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+        device, kConvTranspose2dGemmWGSL, bindings);
+    graph.add_dispatch(
+        {bundle.pipeline,
+         bundle.bind_group,
+         gemm_grid.x,
+         "conv_transpose2d_gemm",
+         gemm_grid.y});
+  } else {
+    auto constants = utils::make_grid_constants(grid);
+    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+        device,
+        kConvTranspose2dWGSL,
+        bindings,
+        constants.data(),
+        constants.size());
+    graph.add_dispatch(
+        {bundle.pipeline,
+         bundle.bind_group,
+         grid.count_x,
+         "conv_transpose2d",
+         grid.count_y});
+  }
 
   wgpuBufferRelease(uniform_buffer);
   if (bias.owned_dummy != nullptr) {

backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_gemm.wgsl

@@ -0,0 +1,111 @@
+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read> weight: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+
+struct Params {
+  B: u32, IC: u32, IH: u32, IW: u32,
+  OC: u32, OH: u32, OW: u32, KH: u32, KW: u32,
+  sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32,
+  groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+// Non-overlapping transposed conv (stride == kernel, no padding/dilation,
+// groups == 1): every output pixel takes exactly one kernel tap, so it is a
+// GEMM M=OC*KH*KW, N=B*IH*IW, K=IC followed by a depth-to-space scatter.
+// A = weight [IC, OC, KH, KW] read transposed (A[m,k] = weight[k*M + m]); B =
+// input [B, IC, IH*IW]. Same tiling as conv2d_gemm (TILE=32, RPT=4).
+const TILE: u32 = 32u;
+const RPT: u32 = 4u;
+
+var<workgroup> a_sub: array<array<f32, 32>, 32>;
+var<workgroup> b_sub: array<array<f32, 32>, 32>;
+
+fn gM() -> u32 { return params.OC * params.KH * params.KW; }
+fn gN() -> u32 { return params.B * params.IH * params.IW; }
+
+fn read_a(m: u32, k: u32) -> f32 {
+  if (m < gM() && k < params.IC) {
+    return weight[k * gM() + m];
+  }
+  return 0.0;
+}
+
+fn read_b(k: u32, n: u32) -> f32 {
+  if (k >= params.IC || n >= gN()) {
+    return 0.0;
+  }
+  let ihw = params.IH * params.IW;
+  return input[((n / ihw) * params.IC + k) * ihw + n % ihw];
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(
+    @builtin(workgroup_id) wg_id: vec3<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>) {
+  let M = gM();
+  let N = gN();
+  let K = params.IC;
+  let tile_row0 = wg_id.y * TILE;
+  let tile_col0 = wg_id.x * TILE;
+  let tile_row = local_id.y * RPT;
+  let tile_col = local_id.x * RPT;
+
+  var acc: array<array<f32, 4>, 4>;
+  for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+    for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) {
+      acc[ir][ic] = 0.0;
+    }
+  }
+
+  let num_tiles = (K + TILE - 1u) / TILE;
+  for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
+    let k_start = t * TILE;
+    for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+      let arow = local_id.y * RPT + ir;
+      for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) {
+        let col = local_id.x * RPT + kk;
+        a_sub[arow][col] = read_a(tile_row0 + arow, k_start + col);
+        b_sub[arow][col] = read_b(k_start + arow, tile_col0 + col);
+      }
+    }
+    workgroupBarrier();
+    for (var k: u32 = 0u; k < TILE; k = k + 1u) {
+      for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+        let aval = a_sub[tile_row + ir][k];
+        for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) {
+          acc[ir][ic] = acc[ir][ic] + aval * b_sub[k][tile_col + ic];
+        }
+      }
+    }
+    workgroupBarrier();
+  }
+
+  let khw = params.KH * params.KW;
+  let ihw = params.IH * params.IW;
+  for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+    let m = tile_row0 + tile_row + ir;
+    if (m >= M) {
+      continue;
+    }
+    let oc = m / khw;
+    let kh = (m % khw) / params.KW;
+    let kw = m % params.KW;
+    var bias_v = 0.0;
+    if (params.has_bias != 0u) {
+      bias_v = bias[oc];
+    }
+    for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) {
+      let n = tile_col0 + tile_col + ic;
+      if (n < N) {
+        let b = n / ihw;
+        let sp = n % ihw;
+        let oh = (sp / params.IW) * params.KH + kh;
+        let ow = (sp % params.IW) * params.KW + kw;
+        out[((b * params.OC + oc) * params.OH + oh) * params.OW + ow] =
+            acc[ir][ic] + bias_v;
+      }
+    }
+  }
+}