[ExecuTorch][WebGPU] et_vk conv2d: route standard (groups==1) conv through an im2col tiled GEMM (1.1-2.4x)

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

@@ -9,6 +9,7 @@
 #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/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_wgsl.h>
@@ -313,15 +314,34 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint32_t icpg = IC / groups;
   const bool use_vec4 = (icpg % 4u == 0u);
 
-  // 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 (the documented SAM FpnNeck @1008^2 blocker). 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, "conv2d"),
-      kConv2dWorkgroupSizeX,
-      "et_vk_conv2d");
+  // groups==1 non-transposed -> im2col tiled GEMM (reuses the linear tiled-GEMM
+  // skeleton: M=OC, N=B*OH*OW, K=IC*KH*KW; input im2col-sampled on the fly,
+  // out-of-range -> 0.0 for padding). Canary M4 Pro: 1.1-2.4x over the direct
+  // kernel across stem/FPN shapes (biggest on the RGB stem, where the
+  // vec4-over-IC path is inert). Grouped/depthwise stay on the direct kernel
+  // (grouped GEMM is block-diagonal; mirrors ORT).
+  const bool use_gemm = (groups == 1u);
+
+  // Compute the dispatch grid UP-FRONT, before any buffer alloc, so a throw
+  // (grid exceeds the device dispatch/tile limit) can't leak the uniform/bias.
+  constexpr uint32_t kConv2dGemmTile =
+      32u; // fixed @workgroup_size(8,8), TILE=32
+  utils::WgCount gemm_grid = {};
+  utils::DispatchGrid direct_grid = {};
+  if (use_gemm) {
+    // 2D tile grid over (N cols, M rows); folds past the 65535 per-dim ceiling.
+    gemm_grid = utils::compute_tile_grid_2d(
+        device, B * OH * OW, OC, kConv2dGemmTile, "et_vk_conv2d_gemm");
+  } else {
+    // Adaptive 1D->2D dispatch: wg=clamp(256) + 2D-spill past the 65535 ceiling
+    // (the SAM FpnNeck @1008^2 blocker); stride_x decodes
+    // i=gid.y*stride_x+gid.x.
+    direct_grid = utils::compute_dispatch_grid(
+        device,
+        utils::checked_u32(out_numel, "conv2d"),
+        kConv2dWorkgroupSizeX,
+        "et_vk_conv2d");
+  }
 
   ConvParams params = {};
   params.B = B;
@@ -352,29 +372,62 @@ void conv2d_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);
-
-  utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
-      device,
-      use_vec4 ? kConv2dVec4WGSL : kConv2dWGSL,
-      {
-          {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) {
+    // vec4-over-IC is inert for NCHW (strided channel gather), so the GEMM is
+    // scalar — ORT skips vec4 for NCHW too.
+    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+        device,
+        kConv2dGemmWGSL,
+        {
+            {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)},
+        });
+    graph.add_dispatch_2d(
+        bundle.pipeline, bundle.bind_group, gemm_grid.x, gemm_grid.y);
+  } else {
+    // Direct conv (grouped/depthwise); vec4-over-IC when icpg%4==0 (register-
+    // packed, not coalesced — NCHW's channel stride isn't contiguous).
+    auto constants = utils::make_grid_constants(direct_grid);
+    utils::ComputePipelineBundle bundle = utils::make_compute_pipeline(
+        device,
+        use_vec4 ? kConv2dVec4WGSL : kConv2dWGSL,
+        {
+            {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,
+        direct_grid.count_x,
+        direct_grid.count_y);
+  }
 
   wgpuBufferRelease(uniform_buffer);
   if (bias.owned_dummy != nullptr) {

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

@@ -0,0 +1,116 @@
+@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;
+
+// Standard (groups==1) conv2d as an implicit-im2col tiled GEMM, reusing the
+// linear_fp32_tiled skeleton (TILE=32, RPT=4, shared-mem tiles). M=OC, N=B*OH*OW,
+// K=IC*KH*KW. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly (im2col),
+// out-of-range -> 0.0 implements padding. bias is per-ROW (OC). Output written NCHW.
+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 gK() -> u32 { return params.IC * params.KH * params.KW; }
+fn gN() -> u32 { return params.B * params.OH * params.OW; }
+
+fn read_a(m: u32, kk: u32) -> f32 { // weight[OC, IC*KH*KW] row-major
+  if (m < params.OC && kk < gK()) {
+    return weight[m * gK() + kk];
+  }
+  return 0.0;
+}
+
+fn read_b(kk: u32, n: u32) -> f32 { // im2col sample of input[B,IC,IH,IW]
+  if (kk >= gK() || n >= gN()) {
+    return 0.0;
+  }
+  let ohw = params.OH * params.OW;
+  let b = n / ohw;
+  let sp = n % ohw;
+  let oh = sp / params.OW;
+  let ow = sp % params.OW;
+  let khw = params.KH * params.KW;
+  let ic = kk / khw;
+  let r = kk % khw;
+  let kh = r / params.KW;
+  let kw = r % params.KW;
+  let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH);
+  let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW);
+  if (ih < 0 || ih >= i32(params.IH) || iw < 0 || iw >= i32(params.IW)) {
+    return 0.0; // padding
+  }
+  return input[((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw)];
+}
+
+@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 = params.OC;
+  let N = gN();
+  let K = gK();
+  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();
+  }
+
+  for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+    for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) {
+      let m = tile_row0 + tile_row + ir; // oc
+      let n = tile_col0 + tile_col + ic; // spatial index
+      if (m < M && n < N) {
+        var val = acc[ir][ic];
+        if (params.has_bias != 0u) {
+          val = val + bias[m];
+        }
+        let ohw = params.OH * params.OW;
+        let b = n / ohw;
+        let sp = n % ohw;
+        let oh = sp / params.OW;
+        let ow = sp % params.OW;
+        out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = val;
+      }
+    }
+  }
+}

backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h

@@ -0,0 +1,140 @@
+/*
+ * 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 conv2d_gemm.wgsl - DO NOT EDIT.
+// wgsl-sha256: 0cd4ee9c946a7d7029062fcd6cb137fe32db886ce7f07c3514ec58ede0fa695e
+inline constexpr const char* kConv2dGemmWGSL = R"(
+@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;
+
+// Standard (groups==1) conv2d as an implicit-im2col tiled GEMM, reusing the
+// linear_fp32_tiled skeleton (TILE=32, RPT=4, shared-mem tiles). M=OC, N=B*OH*OW,
+// K=IC*KH*KW. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly (im2col),
+// out-of-range -> 0.0 implements padding. bias is per-ROW (OC). Output written NCHW.
+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 gK() -> u32 { return params.IC * params.KH * params.KW; }
+fn gN() -> u32 { return params.B * params.OH * params.OW; }
+
+fn read_a(m: u32, kk: u32) -> f32 { // weight[OC, IC*KH*KW] row-major
+  if (m < params.OC && kk < gK()) {
+    return weight[m * gK() + kk];
+  }
+  return 0.0;
+}
+
+fn read_b(kk: u32, n: u32) -> f32 { // im2col sample of input[B,IC,IH,IW]
+  if (kk >= gK() || n >= gN()) {
+    return 0.0;
+  }
+  let ohw = params.OH * params.OW;
+  let b = n / ohw;
+  let sp = n % ohw;
+  let oh = sp / params.OW;
+  let ow = sp % params.OW;
+  let khw = params.KH * params.KW;
+  let ic = kk / khw;
+  let r = kk % khw;
+  let kh = r / params.KW;
+  let kw = r % params.KW;
+  let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH);
+  let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW);
+  if (ih < 0 || ih >= i32(params.IH) || iw < 0 || iw >= i32(params.IW)) {
+    return 0.0; // padding
+  }
+  return input[((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw)];
+}
+
+@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 = params.OC;
+  let N = gN();
+  let K = gK();
+  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();
+  }
+
+  for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) {
+    for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) {
+      let m = tile_row0 + tile_row + ir; // oc
+      let n = tile_col0 + tile_col + ic; // spatial index
+      if (m < M && n < N) {
+        var val = acc[ir][ic];
+        if (params.has_bias != 0u) {
+          val = val + bias[m];
+        }
+        let ohw = params.OH * params.OW;
+        let b = n / ohw;
+        let sp = n % ohw;
+        let oh = sp / params.OW;
+        let ow = sp % params.OW;
+        out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = val;
+      }
+    }
+  }
+}
+)";
+
+inline constexpr uint32_t kConv2dGemmWorkgroupSizeX = 8;
+inline constexpr uint32_t kConv2dGemmWorkgroupSizeY = 8;
+inline constexpr uint32_t kConv2dGemmWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu