diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 75f89fe17cb..9eece554fe6 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -58,6 +58,7 @@ set(WEBGPU_SRCS runtime/ops/et_vk_conv2d/Conv2d.cpp runtime/ops/et_vk_sdpa/EtVkSdpa.cpp runtime/ops/embedding/Embedding.cpp + runtime/ops/addmm/Addmm.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/addmm/Addmm.cpp b/backends/webgpu/runtime/ops/addmm/Addmm.cpp new file mode 100644 index 00000000000..f3bc2b696b7 --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/Addmm.cpp @@ -0,0 +1,134 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +struct AddmmParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t self_2d; + float beta; + float alpha; + uint32_t _pad[2]; +}; +static_assert(sizeof(AddmmParams) == 32, "AddmmParams must be 32 bytes"); + +constexpr uint32_t kTile = 32u; + +// aten.addmm.default args: [self, mat1, mat2, beta, alpha, out]. +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N], self [N] or +// [M,N] (HF Linear lowers to addmm with a [N] bias). Shared-memory-tiled GEMM, +// re-derived for mat2's [K,N] layout (mirrors the sibling `linear` op's +// linear_tiled.wgsl skeleton, NOT its [N,K]-shaped read_b). No vec4 variant: +// unlike linear's weight [N,K] (K contiguous, vec4-over-K natural on both +// sides), mat2 [K,N] has N contiguous — vec4-over-K would need a strided +// gather on the mat2 side, eroding the benefit; the tiled kernel alone already +// solves the coalescing gap via the cooperative shared-memory tile load. +void addmm_impl(WebGPUGraph& graph, const std::vector& args) { + const int self_id = args.at(0); + const int mat1_id = args.at(1); + const int mat2_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& self_t = graph.get_tensor(self_id); + const auto& mat1 = graph.get_tensor(mat1_id); + const auto& mat2 = graph.get_tensor(mat2_id); + const auto& out = graph.get_tensor(out_id); + + if (self_t.buffer == nullptr || mat1.buffer == nullptr || + mat2.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error("WebGPU addmm: null buffer binding"); + } + if (mat1.dims.size() != 2 || mat2.dims.size() != 2 || out.dims.size() != 2) { + throw std::runtime_error("WebGPU addmm: expected 2D self-mm/mat1/mat2/out"); + } + + const uint32_t M = static_cast(out.dims[0]); + const uint32_t N = static_cast(out.dims[1]); + const uint32_t K = static_cast(mat1.dims[1]); + if (M == 0 || N == 0 || K == 0) { + throw std::runtime_error("WebGPU addmm: zero M/N/K"); + } + if (static_cast(mat2.dims[0]) != K || + static_cast(mat2.dims[1]) != N) { + throw std::runtime_error("WebGPU addmm: mat2 shape != [K, N]"); + } + const uint64_t out_numel = utils::check_fp32(out, "addmm", "output"); + + const uint64_t self_numel = utils::numel(self_t.dims); + const bool self_2d = self_numel == out_numel; + if (!self_2d && self_numel != N) { + throw std::runtime_error( + "WebGPU addmm: self must broadcast from [N] or [M,N]"); + } + + const float beta = utils::scalar_or(graph, args.at(3), 1.0f); + const float alpha = utils::scalar_or(graph, args.at(4), 1.0f); + + // A genuinely-2D tile dispatch 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, "addmm"); + + AddmmParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.self_2d = self_2d ? 1u : 0u; + params.beta = beta; + params.alpha = alpha; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AddmmParams)); + graph.add_uniform_buffer_bytes(sizeof(AddmmParams)); + + // Tiled kernel has a fixed @workgroup_size(8, 8, 1) — no override constant. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAddmmTiledWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + self_t.buffer, + self_t.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, mat1.buffer, mat1.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, mat2.buffer, mat2.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AddmmParams)}, + }); + + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, tile_grid.x, tile_grid.y); + + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.addmm.default, addmm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl b/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl new file mode 100644 index 00000000000..a2e9a24ecdc --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/addmm_tiled.wgsl @@ -0,0 +1,99 @@ +struct Params { + M: u32, + N: u32, + K: u32, + self_2d: u32, + beta: f32, + alpha: f32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_self: array; +@group(0) @binding(2) var t_mat1: array; +@group(0) @binding(3) var t_mat2: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled addmm (same skeleton as the sibling `linear` op's +// linear_tiled.wgsl), re-derived for mat2's actual [K,N] layout (NOT [N,K] +// like linear's weight — mat2[k,n] is contiguous over n, so the tile-load's +// per-thread-row reads are cacheline-coalesced, unlike linear's weight; still +// RPT=4-float-strided across threads within a row, not 1-float-strided). +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N]. +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_mat1[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + if (krow < params.K && col < params.N) { + return t_mat2[krow * params.N + col]; + } + 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 self_val: f32; + if (params.self_2d == 1u) { + self_val = t_self[r * params.N + c]; + } else { + self_val = t_self[c]; + } + t_out[r * params.N + c] = params.beta * self_val + params.alpha * acc[ir][ic]; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h b/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h new file mode 100644 index 00000000000..47e5ac60527 --- /dev/null +++ b/backends/webgpu/runtime/ops/addmm/addmm_tiled_wgsl.h @@ -0,0 +1,123 @@ +/* + * 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 addmm_tiled.wgsl - DO NOT EDIT. +// wgsl-sha256: 70d1e4fe13aea0d042fa657c8a9d6da81f3ea7fa8e35de98209a5a140308d364 +inline constexpr const char* kAddmmTiledWGSL = R"( +struct Params { + M: u32, + N: u32, + K: u32, + self_2d: u32, + beta: f32, + alpha: f32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_self: array; +@group(0) @binding(2) var t_mat1: array; +@group(0) @binding(3) var t_mat2: array; +@group(0) @binding(4) var params: Params; + +// Shared-memory tiled addmm (same skeleton as the sibling `linear` op's +// linear_tiled.wgsl), re-derived for mat2's actual [K,N] layout (NOT [N,K] +// like linear's weight — mat2[k,n] is contiguous over n, so the tile-load's +// per-thread-row reads are cacheline-coalesced, unlike linear's weight; still +// RPT=4-float-strided across threads within a row, not 1-float-strided). +// out = beta*self + alpha*(mat1 @ mat2); mat1 [M,K], mat2 [K,N]. +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_mat1[row * params.K + col]; + } + return 0.0; +} + +fn read_b(krow: u32, col: u32) -> f32 { + if (krow < params.K && col < params.N) { + return t_mat2[krow * params.N + col]; + } + 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 self_val: f32; + if (params.self_2d == 1u) { + self_val = t_self[r * params.N + c]; + } else { + self_val = t_self[c]; + } + t_out[r * params.N + c] = params.beta * self_val + params.alpha * acc[ir][ic]; + } + } + } +} +)"; + +inline constexpr uint32_t kAddmmTiledWorkgroupSizeX = 8; +inline constexpr uint32_t kAddmmTiledWorkgroupSizeY = 8; +inline constexpr uint32_t kAddmmTiledWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu