From b8ce7e4d4dd413eb409353cb2d66c3f47493c3c1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:24:38 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../runtime/ops/linear_fp32/LinearFp32.cpp | 129 ++++++++++++++++++ .../ops/linear_fp32/linear_fp32_tiled.wgsl | 91 ++++++++++++ .../ops/linear_fp32/linear_fp32_tiled_wgsl.h | 115 ++++++++++++++++ .../ops/linear_fp32/linear_fp32_vec4.wgsl | 92 +++++++++++++ .../ops/linear_fp32/linear_fp32_vec4_wgsl.h | 116 ++++++++++++++++ 6 files changed, 544 insertions(+) create mode 100644 backends/webgpu/runtime/ops/linear_fp32/LinearFp32.cpp create mode 100644 backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled.wgsl create mode 100644 backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled_wgsl.h create mode 100644 backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4.wgsl create mode 100644 backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 86054bccec6..7bcfb1d5e14 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -54,6 +54,7 @@ set(WEBGPU_SRCS runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp runtime/ops/gelu/Gelu.cpp runtime/ops/layer_norm/NativeLayerNorm.cpp + runtime/ops/linear_fp32/LinearFp32.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/linear_fp32/LinearFp32.cpp b/backends/webgpu/runtime/ops/linear_fp32/LinearFp32.cpp new file mode 100644 index 00000000000..117c18f30a8 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_fp32/LinearFp32.cpp @@ -0,0 +1,129 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct LinearParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t has_bias; +}; +static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes"); + +constexpr uint32_t kTile = 32u; + +// aten.linear.default args: [in, weight, bias, out] (mirrors Vulkan Linear.cpp +// linear_packed_weight). Shared-memory-tiled GEMM (+bias), vec4-over-K when +// K%4==0 — mirrors the sibling `linear` op's linear_tiled.wgsl/linear_vec4.wgsl +// (Linear.cpp:95,98), with a bias epilogue added. weight is the prepacked +// [N, K] constant (et_vk.prepack copies it). bias may be None. +void linear_fp32_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 out_id = args.at(3); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& out_tensor = graph.get_tensor(out_id); + + if (in_tensor.dims.empty() || out_tensor.dims.empty()) { + throw std::runtime_error("WebGPU linear: empty input/output dims"); + } + + const uint64_t out_numel = utils::check_fp32(out_tensor, "linear", "output"); + + const uint32_t N = static_cast(out_tensor.dims.back()); + const uint32_t K = static_cast(in_tensor.dims.back()); + if (N == 0 || K == 0) { + throw std::runtime_error("WebGPU linear: zero N or K"); + } + const uint32_t M = static_cast(out_numel / N); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + // A genuinely-2D tile dispatch (mirrors Linear.cpp) doesn't need the 1D-flat + // stride_x recovery trick; throw before any allocation if it can't fit. + utils::WgCount tile_grid = + utils::compute_tile_grid_2d(device, N, M, kTile, "linear_fp32"); + const bool use_vec4 = (K % 4u == 0u); + + LinearParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.has_bias = has_bias ? 1u : 0u; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LinearParams)); + graph.add_uniform_buffer_bytes(sizeof(LinearParams)); + + // A 4-byte dummy storage buffer to satisfy the t_bias binding when bias is + // None; the shader never reads it (has_bias gate). + 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); + + // Tiled kernels have a fixed @workgroup_size(8, 8, 1) — no override constant. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kLinearFp32Vec4WGSL : kLinearFp32TiledWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LinearParams)}, + }); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, tile_grid.x, tile_grid.y); + + wgpuBufferRelease(uniform_buffer); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.linear.default, linear_fp32_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled.wgsl b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled.wgsl new file mode 100644 index 00000000000..afe2231d79b --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled.wgsl @@ -0,0 +1,91 @@ +struct Params { + M: u32, + N: u32, + K: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled linear (+bias), mirroring onnxruntime MatMulPacked +// (same skeleton as the sibling `linear` op's linear_tiled.wgsl). +// out[m,n] = sum_k in[m,k] * weight[n,k] (+ bias[n]); weight is [N,K]. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return t_in[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + // weight[n=col][k=krow] lives at weight[col*K + krow] (transposed [N,K]). + if (col < params.N && krow < params.K) { + return t_weight[col * params.K + krow]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + 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 = (params.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 r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + t_out[r * params.N + c] = v; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled_wgsl.h b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled_wgsl.h new file mode 100644 index 00000000000..1dd967a0f33 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_tiled_wgsl.h @@ -0,0 +1,115 @@ +/* + * 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 linear_fp32_tiled.wgsl - DO NOT EDIT. +// wgsl-sha256: 3cd543e827fc9fdb3323f573ff07448d4c74f41efccc581167b1ae8a49f8eaa3 +inline constexpr const char* kLinearFp32TiledWGSL = R"( +struct Params { + M: u32, + N: u32, + K: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled linear (+bias), mirroring onnxruntime MatMulPacked +// (same skeleton as the sibling `linear` op's linear_tiled.wgsl). +// out[m,n] = sum_k in[m,k] * weight[n,k] (+ bias[n]); weight is [N,K]. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return t_in[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + // weight[n=col][k=krow] lives at weight[col*K + krow] (transposed [N,K]). + if (col < params.N && krow < params.K) { + return t_weight[col * params.K + krow]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + 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 = (params.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 r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kLinearFp32TiledWorkgroupSizeX = 8; +inline constexpr uint32_t kLinearFp32TiledWorkgroupSizeY = 8; +inline constexpr uint32_t kLinearFp32TiledWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4.wgsl b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4.wgsl new file mode 100644 index 00000000000..ca0528ef37d --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4.wgsl @@ -0,0 +1,92 @@ +struct Params { + M: u32, + N: u32, + K: u32, + has_bias: u32, +}; + +// input/weight fp32 buffers viewed as vec4 over K (requires K%4==0). Output is +// scalar f32 (+bias). out[m,n] = sum_k input[m,k]*weight[n,k] (+bias[n]) +// = sum_k4 dot(in4, w4) (+bias[n]). Same skeleton as the sibling `linear` op's +// linear_vec4.wgsl, with a bias epilogue added. +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const TILE4: u32 = 8u; +const RPT: u32 = 4u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn in4(row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return t_in[row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +fn w4(row: u32, k4: u32) -> vec4 { + if (row < params.N && k4 * 4u < params.K) { + return t_weight[row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let row0 = wg_id.y * TILE; + let col0 = wg_id.x * TILE; + let tr = local_id.y * RPT; + let tc = 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 = (params.K + TILE - 1u) / TILE; + let flat0 = local_id.y * 8u + local_id.x; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4s = t * TILE4; + for (var i: u32 = 0u; i < 4u; i = i + 1u) { + let e = flat0 + i * 64u; + let rr = e / TILE4; + let kk = e % TILE4; + a_sub[rr][kk] = in4(row0 + rr, k4s + kk); + b_sub[rr][kk] = w4(col0 + rr, k4s + kk); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let av = a_sub[tr + ir][k4]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + dot(av, b_sub[tc + ic][k4]); + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = row0 + tr + ir; + let c = col0 + tc + ic; + if (r < params.M && c < params.N) { + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + t_out[r * params.N + c] = v; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4_wgsl.h b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4_wgsl.h new file mode 100644 index 00000000000..f9ce94f6226 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_fp32/linear_fp32_vec4_wgsl.h @@ -0,0 +1,116 @@ +/* + * 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 linear_fp32_vec4.wgsl - DO NOT EDIT. +// wgsl-sha256: 87bfab38ee3232bba769070b8f38765ed35ec015540ddff063742124a528f27c +inline constexpr const char* kLinearFp32Vec4WGSL = R"( +struct Params { + M: u32, + N: u32, + K: u32, + has_bias: u32, +}; + +// input/weight fp32 buffers viewed as vec4 over K (requires K%4==0). Output is +// scalar f32 (+bias). out[m,n] = sum_k input[m,k]*weight[n,k] (+bias[n]) +// = sum_k4 dot(in4, w4) (+bias[n]). Same skeleton as the sibling `linear` op's +// linear_vec4.wgsl, with a bias epilogue added. +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const TILE4: u32 = 8u; +const RPT: u32 = 4u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn in4(row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return t_in[row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +fn w4(row: u32, k4: u32) -> vec4 { + if (row < params.N && k4 * 4u < params.K) { + return t_weight[row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let row0 = wg_id.y * TILE; + let col0 = wg_id.x * TILE; + let tr = local_id.y * RPT; + let tc = 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 = (params.K + TILE - 1u) / TILE; + let flat0 = local_id.y * 8u + local_id.x; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4s = t * TILE4; + for (var i: u32 = 0u; i < 4u; i = i + 1u) { + let e = flat0 + i * 64u; + let rr = e / TILE4; + let kk = e % TILE4; + a_sub[rr][kk] = in4(row0 + rr, k4s + kk); + b_sub[rr][kk] = w4(col0 + rr, k4s + kk); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let av = a_sub[tr + ir][k4]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + dot(av, b_sub[tc + ic][k4]); + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = row0 + tr + ir; + let c = col0 + tc + ic; + if (r < params.M && c < params.N) { + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kLinearFp32Vec4WorkgroupSizeX = 8; +inline constexpr uint32_t kLinearFp32Vec4WorkgroupSizeY = 8; +inline constexpr uint32_t kLinearFp32Vec4WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu