From 1c5c7b232c7bc65a69c1ca2964e7366787c3e027 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:24:47 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../runtime/ops/et_vk_conv2d/Conv2d.cpp | 391 ++++++++++++++++++ .../runtime/ops/et_vk_conv2d/conv2d.wgsl | 81 ++++ .../runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl | 98 +++++ .../ops/et_vk_conv2d/conv2d_vec4_wgsl.h | 122 ++++++ .../runtime/ops/et_vk_conv2d/conv2d_wgsl.h | 105 +++++ .../ops/et_vk_conv2d/conv_transpose2d.wgsl | 94 +++++ .../ops/et_vk_conv2d/conv_transpose2d_wgsl.h | 118 ++++++ 8 files changed, 1010 insertions(+) create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl create mode 100644 backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 7bcfb1d5e14..8efff290e3d 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -55,6 +55,7 @@ set(WEBGPU_SRCS runtime/ops/gelu/Gelu.cpp runtime/ops/layer_norm/NativeLayerNorm.cpp runtime/ops/linear_fp32/LinearFp32.cpp + runtime/ops/et_vk_conv2d/Conv2d.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp new file mode 100644 index 00000000000..dc8e8093400 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -0,0 +1,391 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct ConvParams { + uint32_t B; + uint32_t IC; + uint32_t IH; + uint32_t IW; + uint32_t OC; + uint32_t OH; + uint32_t OW; + uint32_t KH; + uint32_t KW; + uint32_t sH; + uint32_t sW; + uint32_t pH; + uint32_t pW; + uint32_t dH; + uint32_t dW; + uint32_t groups; + uint32_t has_bias; + uint32_t _p0; + uint32_t _p1; + uint32_t _p2; +}; +static_assert( + sizeof(ConvParams) == 80, + "ConvParams must be 80 bytes (16-mult)"); + +// 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 +// vs torch.conv_transpose2d to fp64 round-off incl. non-square spatial + +// kernel. +void conv_transpose2d_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int stride_id = args.at(3); + const int padding_id = args.at(4); + const int dilation_id = args.at(5); + const int output_padding_id = args.at(7); + const int groups_id = args.at(8); + const int out_id = args.at(9); + + WGPUDevice device = graph.device(); + + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + if (in.dims.size() != 4 || weight.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error( + "WebGPU conv_transpose2d: expected 4D in/weight/out"); + } + + uint32_t sH, sW, pH, pW, dH, dW, opH, opW; + utils::parse_hw( + graph.get_int_list(stride_id), sH, sW, "conv_transpose2d", "stride"); + utils::parse_hw( + graph.get_int_list(padding_id), pH, pW, "conv_transpose2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "conv_transpose2d", "dilation"); + utils::parse_hw( + graph.get_int_list(output_padding_id), + opH, + opW, + "conv_transpose2d", + "output_padding"); + + const uint32_t B = static_cast(in.dims[0]); + const uint32_t IC = static_cast(in.dims[1]); + const uint32_t IH = static_cast(in.dims[2]); + const uint32_t IW = static_cast(in.dims[3]); + const uint32_t groups = static_cast(graph.get_int(groups_id)); + // Transposed weight: [IC, OC/groups, KH, KW]. + const uint32_t KH = static_cast(weight.dims[2]); + const uint32_t KW = static_cast(weight.dims[3]); + const uint32_t OCpg = static_cast(weight.dims[1]); + const uint32_t OC = OCpg * groups; + + if (groups == 0 || IC % groups != 0 || OC % groups != 0) { + throw std::runtime_error( + "WebGPU conv_transpose2d: groups must divide IC/OC"); + } + if (static_cast(weight.dims[0]) != IC) { + throw std::runtime_error("WebGPU conv_transpose2d: weight dim0 != IC"); + } + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU conv_transpose2d: zero stride"); + } + if (opH >= sH || opW >= sW) { + throw std::runtime_error( + "WebGPU conv_transpose2d: output_padding >= stride"); + } + + // OH = (IH-1)*sH - 2*pH + dH*(KH-1) + output_padding + 1. + const int64_t oh_v = int64_t(IH - 1) * int64_t(sH) - 2 * int64_t(pH) + + int64_t(dH) * (int64_t(KH) - 1) + int64_t(opH) + 1; + const int64_t ow_v = int64_t(IW - 1) * int64_t(sW) - 2 * int64_t(pW) + + int64_t(dW) * (int64_t(KW) - 1) + int64_t(opW) + 1; + if (oh_v <= 0 || ow_v <= 0) { + throw std::runtime_error( + "WebGPU conv_transpose2d: invalid output geometry"); + } + const uint32_t OH = static_cast(oh_v); + const uint32_t OW = static_cast(ow_v); + + if (static_cast(out.dims[0]) != B || + static_cast(out.dims[1]) != OC || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU conv_transpose2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "conv_transpose2d", "output"); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + // 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"); + + ConvParams params = {}; + params.B = B; + params.IC = IC; + params.IH = IH; + params.IW = IW; + params.OC = OC; + params.OH = OH; + params.OW = OW; + params.KH = KH; + params.KW = KW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + params.groups = groups; + params.has_bias = has_bias ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvParams)); + + utils::OptionalBinding bias = utils::make_optional_binding( + device, + 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, + 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); + + wgpuBufferRelease(uniform_buffer); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +// aten.convolution.default args: [input, weight, bias, stride, padding, +// dilation, transposed, output_padding, groups, out] (Vulkan Convolution.cpp). +// Direct 2D conv (non-transposed), NCHW fp32, general +// stride/pad/dilation/groups. The fused variant et_vk.conv_with_clamp.default +// is NOT handled (no post-conv activation in the SigLIP patch-embed); it would +// error clearly at load if hit. +void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() != 10) { + throw std::runtime_error("WebGPU conv2d: expected 10 args (convolution)"); + } + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int stride_id = args.at(3); + const int padding_id = args.at(4); + const int dilation_id = args.at(5); + const int transposed_id = args.at(6); + const int groups_id = args.at(8); + const int out_id = args.at(9); + + WGPUDevice device = graph.device(); + + if (graph.get_bool(transposed_id)) { + // C2: fold — transposed conv is a distinct gather kernel, same + // registration. + conv_transpose2d_impl(graph, args); + return; + } + // output_padding (args[7]) is only meaningful for transposed conv; require it + // to be zero so a stale non-transposed export can't silently mis-size output. + for (int64_t v : graph.get_int_list(args.at(7))) { + if (v != 0) { + throw std::runtime_error( + "WebGPU conv2d: non-zero output_padding unsupported"); + } + } + + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + + if (in.dims.size() != 4 || weight.dims.size() != 4 || out.dims.size() != 4) { + throw std::runtime_error("WebGPU conv2d: expected 4D input/weight/out"); + } + + // stride/padding/dilation are int lists; PyTorch broadcasts a single value to + // both spatial dims (e.g. padding="valid" serializes as [0]). Accept size 1 + // (broadcast) or 2 (H, W). + uint32_t sH, sW, pH, pW, dH, dW; + utils::parse_hw(graph.get_int_list(stride_id), sH, sW, "conv2d", "stride"); + utils::parse_hw(graph.get_int_list(padding_id), pH, pW, "conv2d", "padding"); + utils::parse_hw( + graph.get_int_list(dilation_id), dH, dW, "conv2d", "dilation"); + + const uint32_t B = static_cast(in.dims[0]); + const uint32_t IC = static_cast(in.dims[1]); + const uint32_t IH = static_cast(in.dims[2]); + const uint32_t IW = static_cast(in.dims[3]); + const uint32_t OC = static_cast(weight.dims[0]); + const uint32_t KH = static_cast(weight.dims[2]); + const uint32_t KW = static_cast(weight.dims[3]); + const uint32_t groups = static_cast(graph.get_int(groups_id)); + + if (groups == 0 || OC % groups != 0 || IC % groups != 0) { + throw std::runtime_error("WebGPU conv2d: groups must divide IC and OC"); + } + // weight is [OC, IC/groups, KH, KW]; verify the channel dim. + if (static_cast(weight.dims[1]) != IC / groups) { + throw std::runtime_error("WebGPU conv2d: weight in-channels != IC/groups"); + } + + if (sH == 0 || sW == 0) { + throw std::runtime_error("WebGPU conv2d: zero stride"); + } + + // Compute in signed 64-bit: the numerator goes negative for invalid geometry + // (e.g. kernel larger than the padded input) and would wrap as unsigned. + const int64_t oh_num = + int64_t(IH) + 2 * int64_t(pH) - int64_t(dH) * (int64_t(KH) - 1) - 1; + const int64_t ow_num = + int64_t(IW) + 2 * int64_t(pW) - int64_t(dW) * (int64_t(KW) - 1) - 1; + if (oh_num < 0 || ow_num < 0) { + throw std::runtime_error( + "WebGPU conv2d: invalid geometry (kernel > input)"); + } + const uint32_t OH = static_cast(oh_num / int64_t(sH)) + 1; + const uint32_t OW = static_cast(ow_num / int64_t(sW)) + 1; + + // Validate against the serialized output tensor shape [B, OC, OH, OW]. + if (static_cast(out.dims[0]) != B || + static_cast(out.dims[1]) != OC || + static_cast(out.dims[2]) != OH || + static_cast(out.dims[3]) != OW) { + throw std::runtime_error("WebGPU conv2d: output shape mismatch"); + } + + uint64_t out_numel = utils::check_fp32(out, "conv2d", "output"); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + // NCHW's channel dim isn't memory-contiguous, so this is a register-packing + // vec4 (4 strided scalar loads gathered), not a coalesced load — still cuts + // the icg loop trip count 4x. Only valid when every group's input-channel + // count is a multiple of 4 (a real RGB stem's icpg=3 needs the scalar path). + 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"); + + ConvParams params = {}; + params.B = B; + params.IC = IC; + params.IH = IH; + params.IW = IW; + params.OC = OC; + params.OH = OH; + params.OW = OW; + params.KH = KH; + params.KW = KW; + params.sH = sH; + params.sW = sW; + params.pH = pH; + params.pW = pW; + params.dH = dH; + params.dW = dW; + params.groups = groups; + params.has_bias = has_bias ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvParams)); + + // Dummy 4-byte storage to satisfy the bias binding when None (gated in WGSL). + utils::OptionalBinding bias = utils::make_optional_binding( + device, + 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); + + wgpuBufferRelease(uniform_buffer); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.convolution.default, conv2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl new file mode 100644 index 00000000000..d87c4d70005 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl @@ -0,0 +1,81 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution (non-transposed), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. Supports general stride/padding/dilation and +// groups. input [B,IC,IH,IW]; weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; // *IW added per-row below + let w_c_base = (oc * icpg + icg) * params.KH; // *KW added per-row below + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl new file mode 100644 index 00000000000..48b38d5a3ad --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4.wgsl @@ -0,0 +1,98 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution, vec4-gathered over IC (host selects this only when +// icpg % 4 == 0 — NCHW's channel dim has stride IH*IW, not contiguous, so this +// is a register-packing vec4 (4 strided scalar loads combined), not a +// coalesced-memory vec4 load; it still cuts the icg loop trip count 4x. ONE +// thread per (b, oc, oh, ow) output element. input [B,IC,IH,IW]; +// weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let icpg4 = icpg / 4u; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let ic_stride = params.IH * params.IW; // channel stride in `input` + let w_ic_stride = params.KH * params.KW; // channel stride in `weight` + for (var icg4: u32 = 0u; icg4 < icpg4; icg4 = icg4 + 1u) { + let icg = icg4 * 4u; + let ic = ic0 + icg; + let in_c_base = b * params.IC * ic_stride + ic * ic_stride; + let w_c_base = oc * icpg * w_ic_stride + icg * w_ic_stride; + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let in_off = u32(ih) * params.IW + u32(iw); + let w_off = kh * params.KW + kw; + let in4 = vec4( + input[in_c_base + in_off], + input[in_c_base + ic_stride + in_off], + input[in_c_base + 2u * ic_stride + in_off], + input[in_c_base + 3u * ic_stride + in_off]); + let w4 = vec4( + weight[w_c_base + w_off], + weight[w_c_base + w_ic_stride + w_off], + weight[w_c_base + 2u * w_ic_stride + w_off], + weight[w_c_base + 3u * w_ic_stride + w_off]); + acc = acc + dot(in4, w4); + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h new file mode 100644 index 00000000000..19a124700bb --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_vec4_wgsl.h @@ -0,0 +1,122 @@ +/* + * 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_vec4.wgsl - DO NOT EDIT. +// wgsl-sha256: a60dbe62ac08849f9ed710700e6baaf3af9e618fb99852e6a82d9c5661440abe +inline constexpr const char* kConv2dVec4WGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution, vec4-gathered over IC (host selects this only when +// icpg % 4 == 0 — NCHW's channel dim has stride IH*IW, not contiguous, so this +// is a register-packing vec4 (4 strided scalar loads combined), not a +// coalesced-memory vec4 load; it still cuts the icg loop trip count 4x. ONE +// thread per (b, oc, oh, ow) output element. input [B,IC,IH,IW]; +// weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let icpg4 = icpg / 4u; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let ic_stride = params.IH * params.IW; // channel stride in `input` + let w_ic_stride = params.KH * params.KW; // channel stride in `weight` + for (var icg4: u32 = 0u; icg4 < icpg4; icg4 = icg4 + 1u) { + let icg = icg4 * 4u; + let ic = ic0 + icg; + let in_c_base = b * params.IC * ic_stride + ic * ic_stride; + let w_c_base = oc * icpg * w_ic_stride + icg * w_ic_stride; + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let in_off = u32(ih) * params.IW + u32(iw); + let w_off = kh * params.KW + kw; + let in4 = vec4( + input[in_c_base + in_off], + input[in_c_base + ic_stride + in_off], + input[in_c_base + 2u * ic_stride + in_off], + input[in_c_base + 3u * ic_stride + in_off]); + let w4 = vec4( + weight[w_c_base + w_off], + weight[w_c_base + w_ic_stride + w_off], + weight[w_c_base + 2u * w_ic_stride + w_off], + weight[w_c_base + 3u * w_ic_stride + w_off]); + acc = acc + dot(in4, w4); + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConv2dVec4WorkgroupSizeX = 256; +inline constexpr uint32_t kConv2dVec4WorkgroupSizeY = 1; +inline constexpr uint32_t kConv2dVec4WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h new file mode 100644 index 00000000000..42e67d8bbb0 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_wgsl.h @@ -0,0 +1,105 @@ +/* + * 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.wgsl - DO NOT EDIT. +// wgsl-sha256: 8f675cbc2d43986af1c4592ee312984d1e9b02dee39f67442ae5d79539e4fd11 +inline constexpr const char* kConv2dWGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution (non-transposed), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. Supports general stride/padding/dilation and +// groups. input [B,IC,IH,IW]; weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; // *IW added per-row below + let w_c_base = (oc * icpg + icg) * params.KH; // *KW added per-row below + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConv2dWorkgroupSizeX = 256; +inline constexpr uint32_t kConv2dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl new file mode 100644 index 00000000000..df2e64a15f9 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl @@ -0,0 +1,94 @@ +@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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Transposed 2D convolution (gather form), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. weight layout = torch convT [IC, OC/groups, +// KH, KW] (NOT flipped). For each kernel tap (kh,kw): an input row ih +// contributes iff (oh + pH - kh*dH) is divisible by sH and ih in range (the +// scatter-inversion). CPU-derisked vs torch.conv_transpose2d to fp64 round-off +// (/tmp/convtr_derisk.py), incl. non-square spatial + non-square kernel. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = + ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + // weight [IC, OC/groups, KH, KW]: index (ic, oc_in_g, kh, kw) + let w_idx = + ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + input[in_idx] * weight[w_idx]; + } + } + } + out[i] = acc; +} diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h new file mode 100644 index 00000000000..471c7494bf0 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d_wgsl.h @@ -0,0 +1,118 @@ +/* + * 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 conv_transpose2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 47bdb5dc2548e73c99e4f0f5b76c6fd996f9f3b0b5d6cfadbe675da806ab4a7b +inline constexpr const char* kConvTranspose2dWGSL = 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; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Transposed 2D convolution (gather form), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. weight layout = torch convT [IC, OC/groups, +// KH, KW] (NOT flipped). For each kernel tap (kh,kw): an input row ih +// contributes iff (oh + pH - kh*dH) is divisible by sH and ih in range (the +// scatter-inversion). CPU-derisked vs torch.conv_transpose2d to fp64 round-off +// (/tmp/convtr_derisk.py), incl. non-square spatial + non-square kernel. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * 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 oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = + ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + // weight [IC, OC/groups, KH, KW]: index (ic, oc_in_g, kh, kw) + let w_idx = + ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + input[in_idx] * weight[w_idx]; + } + } + } + out[i] = acc; +} +)"; + +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeX = 256; +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeY = 1; +inline constexpr uint32_t kConvTranspose2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu