From 0a73e83a5e60c238b33091ebfda8eaf2102806f8 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:26:47 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../runtime/ops/et_vk_conv2d/Conv2d.cpp | 115 ++++++++++---- .../runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl | 116 +++++++++++++++ .../ops/et_vk_conv2d/conv2d_gemm_wgsl.h | 140 ++++++++++++++++++ 3 files changed, 340 insertions(+), 31 deletions(-) create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp index dc8e8093400..96fcbafac02 100644 --- a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -313,15 +314,34 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector& 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& 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) { diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl new file mode 100644 index 00000000000..ecae242dbad --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl @@ -0,0 +1,116 @@ +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +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 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 a_sub: array, 32>; +var b_sub: array, 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, + @builtin(local_invocation_id) local_id: vec3) { + 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, 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; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm_wgsl.h new file mode 100644 index 00000000000..b86007f4d6a --- /dev/null +++ b/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 + +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 out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +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 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 a_sub: array, 32>; +var b_sub: array, 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, + @builtin(local_invocation_id) local_id: vec3) { + 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, 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