[ExecuTorch][WebGPU] Add optimized addmm op

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})

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 <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/addmm/addmm_tiled_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <stdexcept>
+
+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<int>& 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<uint32_t>(out.dims[0]);
+  const uint32_t N = static_cast<uint32_t>(out.dims[1]);
+  const uint32_t K = static_cast<uint32_t>(mat1.dims[1]);
+  if (M == 0 || N == 0 || K == 0) {
+    throw std::runtime_error("WebGPU addmm: zero M/N/K");
+  }
+  if (static_cast<uint32_t>(mat2.dims[0]) != K ||
+      static_cast<uint32_t>(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, &params, 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

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<storage, read_write> t_out: array<f32>;
+@group(0) @binding(1) var<storage, read> t_self: array<f32>;
+@group(0) @binding(2) var<storage, read> t_mat1: array<f32>;
+@group(0) @binding(3) var<storage, read> t_mat2: array<f32>;
+@group(0) @binding(4) var<uniform> 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<workgroup> a_sub: array<array<f32, 32>, 32>;
+var<workgroup> b_sub: array<array<f32, 32>, 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<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>) {
+  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 = (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];
+      }
+    }
+  }
+}

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 <cstdint>
+
+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<storage, read_write> t_out: array<f32>;
+@group(0) @binding(1) var<storage, read> t_self: array<f32>;
+@group(0) @binding(2) var<storage, read> t_mat1: array<f32>;
+@group(0) @binding(3) var<storage, read> t_mat2: array<f32>;
+@group(0) @binding(4) var<uniform> 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<workgroup> a_sub: array<array<f32, 32>, 32>;
+var<workgroup> b_sub: array<array<f32, 32>, 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<u32>,
+    @builtin(local_invocation_id) local_id: vec3<u32>) {
+  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 = (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