Repository: microsoft/onnxruntime
Pull request #24478: [WebGPU] Optimize GEMM with vec4
Author: login: xiaofeihan1
Created: 2025-04-21T02:02:35Z | Merged: 2025-04-29T00:33:09Z

### Description
<!-- Describe your changes. -->
In this PR, we use vec4 to optimize GEMM when colums of A and B can be divided by 4, or use previous shader.
I will add u32/vec2 implementation in the future, and we will only keep one shader at that time.

###  Perf comparison
I run customized model only include GEMM(M = N = K = 1024) with nodejs on M2/M3 Max. Roughly 20% increase.
 || !transA&&!transB | transA     | transB     | transA&&transB |  
|------------------|------------|------------|----------------|------------|
| M2               | 9.36->7.41 | 9.45->7.54 | 11.21->8.19    | 9.66->8.37 |
| M3 max           | 8.07->6.99 | 7.54->6.53 | 8.42->5.89     | 5.47->5.29 |





Diff:
diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc
index 4057b63f0c65d..ac8d1a590c250 100644
--- a/onnxruntime/core/providers/webgpu/math/gemm.cc
+++ b/onnxruntime/core/providers/webgpu/math/gemm.cc
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "core/providers/webgpu/math/gemm.h"
+#include "core/providers/webgpu/math/gemm_vec4.h"
 
 #include <vector>
 
@@ -197,6 +198,11 @@ Status Gemm::ComputeInternal(ComputeContext& context) const {
     return Status::OK();
   }
 
+  // First try vec4 optimization if possible
+  if (CanApplyGemmVec4(A, B)) {
+    return ApplyGemmVec4(A, B, C, transA_, transB_, alpha_, beta_, context, Y);
+  }
+
   // WebGPU doesn't support binding a zero-sized buffer, so we need to check if A or B is empty.
   bool need_handle_matmul = A_shape.Size() > 0 && B_shape.Size() > 0;
   bool need_handle_bias = C && beta_;
@@ -220,14 +226,12 @@ Status Gemm::ComputeInternal(ComputeContext& context) const {
       .AddOutputs({{Y, ProgramTensorMetadataDependency::Type}})
       .SetDispatchGroupSize(num_tile_n * num_tile_m)
       .SetWorkgroupSize(TILE_SIZE, TILE_SIZE)
-      .AddUniformVariables({
-          {static_cast<uint32_t>(num_tile_n)},  // num_tile_n
-          {static_cast<uint32_t>(M)},           // M
-          {static_cast<uint32_t>(N)},           // N
-          {static_cast<uint32_t>(K)},           // K
-          {alpha_},                             // alpha
-          {beta_}                               // beta
-      });
+      .AddUniformVariables({{num_tile_n},
+                            {M},
+                            {N},
+                            {K},
+                            {alpha_},
+                            {beta_}});
 
   return context.RunProgram(program);
 }
diff --git a/onnxruntime/core/providers/webgpu/math/gemm_vec4.cc b/onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
new file mode 100644
index 0000000000000..6ba93df8247d2
--- /dev/null
+++ b/onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
@@ -0,0 +1,314 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+#include "core/providers/webgpu/webgpu_utils.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.AdditionalImplementation()
+      << "fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << " { \n"
+      << " if(col < total_cols && row < total_rows) {\n"
+      << "    return A[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+
+  shader.AdditionalImplementation()
+      << "fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << "{ \n"
+      << "  if(col < total_cols && row < total_rows) {\n"
+      << "    return B[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n";
+
+  if (output_components_ == 1) {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N; \n";
+  } else {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N4; \n";
+  }
+
+  shader.AdditionalImplementation() << "var values = valuesIn; \n"
+                                    << "if(col < total_cols && row < uniforms.M) { \n";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += output_element_t(uniforms.beta) * ";
+    // We can be allowed to use broadcasting only when both components are equal.
+    // There is only one case for c_components_ is not equal output_components_.
+    // I.g. the former is `1` and the latter is `4`.
+    // That means the shape of C is either {M,1} or {1,1}
+    if (c_components_ == output_components_) {
+      shader.AdditionalImplementation() << "output_value_t("
+                                        << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ");\n";
+    } else if (c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    }
+  }
+  shader.AdditionalImplementation() << "    output[row * total_cols + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.MainFunctionBody() << "  var values = " << type_string << "(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array< " << type_string << ", 8 >, 32 >;\n "
+                                      << "var<workgroup> tile_b: array<array< " << type_string << ", 8 >, 32 >;\n ";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u;
+        var col = k_start_a + (local_idx % 8u);
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u);
+        col = tile_col_start + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a1 = tile_a[i][local_idx / 32u];
+            let a2 = tile_a[i + 1u][local_idx / 32u];
+            let a3 = tile_a[i + 2u][local_idx / 32u];
+            let a4 = tile_a[i + 3u][local_idx / 32u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            let a = tile_a[i][local_idx / 32u];
+            let b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+         for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<output_element_t>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[i][local_idx % 8u];
+            let b2 = tile_b[i+1][local_idx % 8u];
+            let b3 = tile_b[i+2][local_idx % 8u];
+            let b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = output_element_t(uniforms.alpha) * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Write output
+  if (output_components_ == 1) {
+    shader.MainFunctionBody() << " for (var i = 0u; i < 4u; i = i + 1u) {\n"
+                              << "    mm_write(m, 4 * n + i, values[i]);\n"
+                              << "  }\n";
+  } else {
+    shader.MainFunctionBody() << "  mm_write(m, n, values);\n";
+  }
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  // When the number of columns in A and B is divisible by 4, we apply vec4 optimization to A and B.
+  // However, this doesn't necessarily mean that C and Y will use vec4.
+  // For example, C/output won't be vec4 if B is transposed and N is not divisible by 4.
+  // Also, C won't use vec4 when it's a scalar.
+  // The code would be simpler if we avoided vec4 optimization for C/output.
+  // But to maximize performance, we still apply vec4 when possible — even though it adds some complexity.
+  // I've added detailed comments explaining this logic.
+  // See MatMulReadFnSource and MatMulWriteFnSource, especially the parts related to broadcasting.
+  return a_shape[1] % 4 == 0 && b_shape[1] % 4 == 0;
+}
+
+Status ApplyGemmVec4(const Tensor* a,
+                     const Tensor* b,
+                     const Tensor* c,
+                     bool transA,
+                     bool transB,
+                     float alpha,
+                     float beta,
+                     ComputeContext& context,
+                     Tensor* y) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t M = onnxruntime::narrow<uint32_t>(transA ? a_shape[1] : a_shape[0]);
+  uint32_t K = onnxruntime::narrow<uint32_t>(transA ? a_shape[0] : a_shape[1]);
+  uint32_t N = onnxruntime::narrow<uint32_t>(transB ? b_shape[0] : b_shape[1]);
+
+  // WebGPU doesn't support binding a zero-sized buffer, so we need to check if A or B is empty.
+  bool need_handle_matmul = a_shape.Size() > 0 && b_shape.Size() > 0;
+  bool need_handle_bias = c && beta;
+
+  int c_components = 4;
+  bool c_is_scalar = false;
+
+  // We use vec4 for C when its last dimension equals N and N is divisible by 4.
+  if (need_handle_bias) {
+    const auto& c_shape = c->Shape();
+    int64_t c_last_dim = c_shape[c_shape.NumDimensions() - 1];
+    c_components = (c_last_dim == N && N % 4 == 0) ? 4 : 1;
+    c_is_scalar = c_shape.Size() == 1;
+  }
+
+  // We use vec4 for Y when N is divisible by 4.
+  const int output_components = N % 4 == 0 ? 4 : 1;
+
+  GemmVec4Program program{transA, transB, alpha, need_handle_bias, need_handle_matmul, c_components, c_is_scalar, output_components};
+
+  const int components = 4;
+
+  if (need_handle_matmul) {
+    program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, components},
+                       {b, ProgramTensorMetadataDependency::TypeAndRank, components}});
+  }
+
+  if (need_handle_bias) {
+    program.AddInput({c, ProgramTensorMetadataDependency::TypeAndRank, c_components});
+  }
+
+  const uint32_t TILE_SIZE = 32;
+  const uint32_t num_tile_n = (N + TILE_SIZE - 1) / TILE_SIZE;
+  const uint32_t num_tile_m = (M + TILE_SIZE - 1) / TILE_SIZE;
+
+  program.CacheHint(alpha, transA, transB, c_is_scalar)
+      .AddOutputs({{y, ProgramTensorMetadataDependency::TypeAndRank, output_components}})
+      .SetDispatchGroupSize(num_tile_n * num_tile_m)
+      .SetWorkgroupSize(256, 1, 1)
+      .AddUniformVariables({{num_tile_n},
+                            {M},
+                            {N},
+                            {K},
+                            {M / 4},
+                            {N / 4},
+                            {K / 4},
+                            {alpha},
+                            {beta}});
+
+  return context.RunProgram(program);
+}
+
+}  // namespace webgpu
+}  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/webgpu/math/gemm_vec4.h b/onnxruntime/core/providers/webgpu/math/gemm_vec4.h
new file mode 100644
index 0000000000000..ae7be49ce9218
--- /dev/null
+++ b/onnxruntime/core/providers/webgpu/math/gemm_vec4.h
@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/providers/webgpu/program.h"
+#include "core/providers/webgpu/webgpu_kernel.h"
+
+#include "core/providers/webgpu/shader_helper.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+class GemmVec4Program final : public Program<GemmVec4Program> {
+ public:
+  GemmVec4Program(bool transA, bool transB, float alpha, bool need_handle_bias, bool need_handle_matmul, int c_components, bool c_is_scalar, int output_components)
+      : Program{"GemmVec4"},
+        transA_{transA},
+        transB_{transB},
+        alpha_{alpha},
+        need_handle_bias_{need_handle_bias},
+        need_handle_matmul_{need_handle_matmul},
+        c_components_(c_components),
+        c_is_scalar_(c_is_scalar),
+        output_components_(output_components) {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  void MatMulReadFnSource(ShaderHelper& shader) const;
+  void MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"num_tile_n", ProgramUniformVariableDataType::Uint32},
+      {"M", ProgramUniformVariableDataType::Uint32},
+      {"N", ProgramUniformVariableDataType::Uint32},
+      {"K", ProgramUniformVariableDataType::Uint32},
+      {"M4", ProgramUniformVariableDataType::Uint32},
+      {"N4", ProgramUniformVariableDataType::Uint32},
+      {"K4", ProgramUniformVariableDataType::Uint32},
+      {"alpha", ProgramUniformVariableDataType::Float32},
+      {"beta", ProgramUniformVariableDataType::Float32});
+
+ private:
+  bool transA_;
+  bool transB_;
+  float alpha_;
+  bool need_handle_bias_;
+  bool need_handle_matmul_;
+  int c_components_;
+  bool c_is_scalar_ = false;
+  int output_components_;
+};
+
+Status ApplyGemmVec4(const Tensor* a,
+                     const Tensor* b,
+                     const Tensor* c,
+                     bool transA,
+                     bool transB,
+                     float alpha,
+                     float beta,
+                     ComputeContext& context,
+                     Tensor* y);
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b);
+
+}  // namespace webgpu
+}  // namespace onnxruntime
diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc
index d0069a0069646..400b5ab20930c 100644
--- a/onnxruntime/test/providers/cpu/math/gemm_test.cc
+++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc
@@ -433,6 +433,48 @@ TYPED_TEST(GemmOpTypedTests, TestGemm2DBroadcast_2) {
       .RunWithConfig();
 }
 
+TYPED_TEST(GemmOpTypedTests, TestGemm2DBroadcast_3) {
+  OpTester test("Gemm");
+
+  test.AddAttribute("transA", (int64_t)0);
+  test.AddAttribute("transB", (int64_t)0);
+  test.AddAttribute("alpha", 1.0f);
+  test.AddAttribute("beta", 1.0f);
+
+  // Same as GemmBroadcast, but adding the unnecessary first dimension.
+  test.AddInput<TypeParam>("A", {3, 4},
+                           std::vector<TypeParam>(12, static_cast<TypeParam>(1.0f)));
+  test.AddInput<TypeParam>("B", {4, 4}, std::vector<TypeParam>(16, static_cast<TypeParam>(1.0f)));
+  test.AddInput<TypeParam>("C", {1, 1}, std::vector<TypeParam>{static_cast<TypeParam>(1.0f)});
+  test.AddOutput<TypeParam>("Y", {3, 4},
+                            std::vector<TypeParam>{static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f),
+                                                   static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f),
+                                                   static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f)});
+  test.Config(run_with_tunable_op)
+      .RunWithConfig();
+}
+
+TYPED_TEST(GemmOpTypedTests, TestGemm2DBroadcast_4) {
+  OpTester test("Gemm");
+
+  test.AddAttribute("transA", (int64_t)0);
+  test.AddAttribute("transB", (int64_t)0);
+  test.AddAttribute("alpha", 1.0f);
+  test.AddAttribute("beta", 1.0f);
+
+  // Same as GemmBroadcast, but adding the unnecessary first dimension.
+  test.AddInput<TypeParam>("A", {3, 4},
+                           std::vector<TypeParam>(12, static_cast<TypeParam>(1.0f)));
+  test.AddInput<TypeParam>("B", {4, 4}, std::vector<TypeParam>(16, static_cast<TypeParam>(1.0f)));
+  test.AddInput<TypeParam>("C", {3, 1}, std::vector<TypeParam>{static_cast<TypeParam>(1.0f), static_cast<TypeParam>(2.0f), static_cast<TypeParam>(3.0f)});
+  test.AddOutput<TypeParam>("Y", {3, 4},
+                            std::vector<TypeParam>{static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f), static_cast<TypeParam>(5.0f),
+                                                   static_cast<TypeParam>(6.0f), static_cast<TypeParam>(6.0f), static_cast<TypeParam>(6.0f), static_cast<TypeParam>(6.0f),
+                                                   static_cast<TypeParam>(7.0f), static_cast<TypeParam>(7.0f), static_cast<TypeParam>(7.0f), static_cast<TypeParam>(7.0f)});
+  test.Config(run_with_tunable_op)
+      .RunWithConfig();
+}
+
 TYPED_TEST(GemmOpTypedTests, TestGemmFalseBroadcast) {
   OpTester test("Gemm");
 
@@ -914,5 +956,237 @@ TEST(GemmOpTest, SharedPrepackedWeights) {
 }
 #endif
 
+TEST(GemmOpTest, GemmOptimizeVec4) {
+  auto run_test = [](int64_t M, int64_t K, int64_t N) {
+    OpTester test("Gemm");
+
+    test.AddAttribute("transA", (int64_t)0);
+    test.AddAttribute("transB", (int64_t)0);
+    test.AddAttribute("alpha", 1.0f);
+    test.AddAttribute("beta", 1.0f);
+
+    // Matrix A: MxK filled with sequential numbers
+    std::vector<float> a_data;
+    a_data.reserve(M * K);
+    for (int64_t i = 0; i < M * K; ++i) {
+      a_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix B: KxN filled with sequential numbers
+    std::vector<float> b_data;
+    b_data.reserve(K * N);
+    for (int64_t i = 0; i < K * N; ++i) {
+      b_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix C: MxN filled with zeros
+    std::vector<float> c_data(M * N, 1.0f);
+
+    test.AddInput<float>("A", {M, K}, a_data);
+    test.AddInput<float>("B", {K, N}, b_data);
+    test.AddInput<float>("C", {M, N}, c_data);
+
+    // Calculate expected output
+    std::vector<float> expected_data(M * N, 0.0f);
+    for (int64_t i = 0; i < M; ++i) {
+      for (int64_t j = 0; j < N; ++j) {
+        float sum = 0.0f;
+        for (int64_t k = 0; k < K; ++k) {
+          sum += a_data[i * K + k] * b_data[k * N + j];
+        }
+        expected_data[i * N + j] = sum + c_data[i * N + j];
+      }
+    }
+
+    test.AddOutput<float>("Y", {M, N}, expected_data);
+    test.Config(run_with_tunable_op)
+        .RunWithConfig();
+  };
+
+  run_test(60, 16, 92);
+
+  run_test(8, 8, 8);
+  run_test(128, 128, 128);
+  run_test(128, 32, 64);
+  run_test(4, 8, 12);
+
+  run_test(96, 24, 48);
+  run_test(48, 48, 120);
+  run_test(72, 80, 84);
+}
+
+TEST(GemmOpTest, GemmOptimizeVec4TransA) {
+  auto run_test = [](int64_t M, int64_t K, int64_t N) {
+    OpTester test("Gemm");
+
+    test.AddAttribute("transA", (int64_t)1);  // A is transposed
+    test.AddAttribute("transB", (int64_t)0);
+    test.AddAttribute("alpha", 1.0f);
+    test.AddAttribute("beta", 1.0f);
+
+    // Matrix A: KxM (will be transposed to MxK) filled with sequential numbers
+    std::vector<float> a_data;
+    a_data.reserve(M * K);
+    for (int64_t i = 0; i < K * M; ++i) {
+      a_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix B: KxN filled with sequential numbers
+    std::vector<float> b_data;
+    b_data.reserve(K * N);
+    for (int64_t i = 0; i < K * N; ++i) {
+      b_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix C: MxN filled with zeros
+    std::vector<float> c_data(M * N, 1.0f);
+
+    test.AddInput<float>("A", {K, M}, a_data);  // Note dimensions are swapped
+    test.AddInput<float>("B", {K, N}, b_data);
+    test.AddInput<float>("C", {M, N}, c_data);
+
+    // Calculate expected output for transposed A
+    std::vector<float> expected_data(M * N, 0.0f);
+    for (int64_t i = 0; i < M; ++i) {
+      for (int64_t j = 0; j < N; ++j) {
+        float sum = 0.0f;
+        for (int64_t k = 0; k < K; ++k) {
+          sum += a_data[k * M + i] * b_data[k * N + j];  // Adjusted index for transposed A
+        }
+        expected_data[i * N + j] = sum + c_data[i * N + j];
+      }
+    }
+
+    test.AddOutput<float>("Y", {M, N}, expected_data);
+    test.Config(run_with_tunable_op)
+        .RunWithConfig();
+  };
+
+  run_test(60, 16, 92);
+  run_test(8, 8, 8);
+  run_test(128, 128, 128);
+  run_test(128, 32, 64);
+  run_test(4, 8, 12);
+  run_test(96, 24, 48);
+  run_test(48, 48, 120);
+  run_test(72, 80, 84);
+}
+
+TEST(GemmOpTest, GemmOptimizeVec4TransB) {
+  auto run_test = [](int64_t M, int64_t K, int64_t N) {
+    OpTester test("Gemm");
+
+    test.AddAttribute("transA", (int64_t)0);
+    test.AddAttribute("transB", (int64_t)1);  // B is transposed
+    test.AddAttribute("alpha", 1.0f);
+    test.AddAttribute("beta", 1.0f);
+
+    // Matrix A: MxK filled with sequential numbers
+    std::vector<float> a_data;
+    a_data.reserve(M * K);
+    for (int64_t i = 0; i < M * K; ++i) {
+      a_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix B: NxK (will be transposed to KxN) filled with sequential numbers
+    std::vector<float> b_data;
+    b_data.reserve(K * N);
+    for (int64_t i = 0; i < N * K; ++i) {
+      b_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix C: MxN filled with zeros
+    std::vector<float> c_data(M * N, 1.0f);
+
+    test.AddInput<float>("A", {M, K}, a_data);
+    test.AddInput<float>("B", {N, K}, b_data);  // Note dimensions are swapped
+    test.AddInput<float>("C", {M, N}, c_data);
+
+    // Calculate expected output for transposed B
+    std::vector<float> expected_data(M * N, 0.0f);
+    for (int64_t i = 0; i < M; ++i) {
+      for (int64_t j = 0; j < N; ++j) {
+        float sum = 0.0f;
+        for (int64_t k = 0; k < K; ++k) {
+          sum += a_data[i * K + k] * b_data[j * K + k];  // Adjusted index for transposed B
+        }
+        expected_data[i * N + j] = sum + c_data[i * N + j];
+      }
+    }
+
+    test.AddOutput<float>("Y", {M, N}, expected_data);
+    test.Config(run_with_tunable_op)
+        .RunWithConfig();
+  };
+
+  run_test(32, 32, 32);
+  run_test(60, 16, 92);
+  run_test(8, 8, 8);
+  run_test(64, 64, 64);
+  run_test(128, 128, 128);
+  run_test(128, 32, 64);
+  run_test(4, 8, 12);
+  run_test(96, 24, 48);
+  run_test(48, 48, 120);
+  run_test(72, 80, 84);
+}
+
+TEST(GemmOpTest, GemmOptimizeVec4TransAB) {
+  auto run_test = [](int64_t M, int64_t K, int64_t N) {
+    OpTester test("Gemm");
+
+    test.AddAttribute("transA", (int64_t)1);  // A is transposed
+    test.AddAttribute("transB", (int64_t)1);  // B is transposed
+    test.AddAttribute("alpha", 1.0f);
+    test.AddAttribute("beta", 1.0f);
+
+    // Matrix A: KxM (will be transposed to MxK) filled with sequential numbers
+    std::vector<float> a_data;
+    a_data.reserve(M * K);
+    for (int64_t i = 0; i < K * M; ++i) {
+      a_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix B: NxK (will be transposed to KxN) filled with sequential numbers
+    std::vector<float> b_data;
+    b_data.reserve(K * N);
+    for (int64_t i = 0; i < N * K; ++i) {
+      b_data.push_back(static_cast<float>((i % 7) + 1));
+    }
+
+    // Matrix C: MxN filled with zeros
+    std::vector<float> c_data(M * N, 1.0f);
+
+    test.AddInput<float>("A", {K, M}, a_data);  // Note dimensions are swapped
+    test.AddInput<float>("B", {N, K}, b_data);  // Note dimensions are swapped
+    test.AddInput<float>("C", {M, N}, c_data);
+
+    // Calculate expected output for both matrices transposed
+    std::vector<float> expected_data(M * N, 0.0f);
+    for (int64_t i = 0; i < M; ++i) {
+      for (int64_t j = 0; j < N; ++j) {
+        float sum = 0.0f;
+        for (int64_t k = 0; k < K; ++k) {
+          sum += a_data[k * M + i] * b_data[j * K + k];  // Adjusted indices for both transposed
+        }
+        expected_data[i * N + j] = sum + c_data[i * N + j];
+      }
+    }
+
+    test.AddOutput<float>("Y", {M, N}, expected_data);
+    test.Config(run_with_tunable_op)
+        .RunWithConfig();
+  };
+  run_test(32, 32, 32);
+  run_test(60, 16, 92);
+  run_test(8, 8, 8);
+  run_test(128, 128, 128);
+  run_test(128, 32, 64);
+  run_test(4, 8, 12);
+  run_test(96, 24, 48);
+  run_test(48, 48, 120);
+  run_test(72, 80, 84);
+}
+
 }  // namespace test
 }  // namespace onnxruntime


Reviews:
user: login: qjia7
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T02:45:24Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T00:23:00Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T03:21:39Z
user: login: qjia7
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T03:31:37Z
user: login: qjia7
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T04:57:01Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T09:02:40Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T10:33:00Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T08:48:34Z
user: login: qjia7
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T01:34:26Z
user: login: xiaofeihan1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T03:20:04Z
user: login: qjia7
body: LGTM, thanks.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T08:27:15Z
user: login: fs-eire
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T23:15:44Z

Review comments:
diff_hunk: @@ -0,0 +1,323 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) {"
+      << "var values = valuesIn;"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * "
+                                      << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Calculate bias
+  // if (need_handle_bias_) {
+  //   const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+  //   shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                             << "    values = values + uniforms.beta * "
+  //                             << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"
+  //                             << "  }\n";
+  // }
+
+  // Write output
+  // shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                           << "    " << output.SetByOffset("m * uniforms.N4 + n", "values") << "\n"
+  //                           << "  }\n";
+
+  shader.MainFunctionBody() << " mm_write(m, n, values);\n";
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b,
+                      const Tensor* c) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t A_rows = onnxruntime::narrow<uint32_t>(a_shape[0]);
+  uint32_t A_cols = onnxruntime::narrow<uint32_t>(a_shape[1]);
+  uint32_t B_rows = onnxruntime::narrow<uint32_t>(b_shape[0]);
+  uint32_t B_cols = onnxruntime::narrow<uint32_t>(b_shape[1]);
+
+  if (A_rows % 4 != 0 || A_cols % 4 != 0 || B_rows % 4 != 0 || B_cols % 4 != 0) {
+    return false;
+  }
+
+  // C is optional and is might scalar or 1D.
+  // We do vec4 for C so we need to check if C is 2D and its second dimension is divisible by 4.
+  if (c == nullptr) {
+    return true;
+  }
+
+  const auto& c_shape = c->Shape();
+  if (c_shape.NumDimensions() == 2) {
+    uint32_t C_cols = onnxruntime::narrow<uint32_t>(c_shape[1]);
+    return C_cols % 4 == 0;
+  }
+
+  uint32_t C_rows = onnxruntime::narrow<uint32_t>(c_shape[0]);
+  return C_rows % 4 == 0;
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: I think we can always support `C` no matter it's divisible by 4 or not. C's last dimension is either `N` or `1`. If the last dimension is `N`, it must be divisible by 4. If the last dimension is `1`, just read one scalar and make it to vec4<T>(scalar).
created_at: 2025-04-23T02:31:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,323 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) {"
+      << "var values = valuesIn;"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * "
+                                      << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Calculate bias
+  // if (need_handle_bias_) {
+  //   const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+  //   shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                             << "    values = values + uniforms.beta * "
+  //                             << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"
+  //                             << "  }\n";
+  // }
+
+  // Write output
+  // shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                           << "    " << output.SetByOffset("m * uniforms.N4 + n", "values") << "\n"
+  //                           << "  }\n";
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: Just remove above code if they are needed anymore. 
created_at: 2025-04-23T02:32:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,323 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) {"
+      << "var values = valuesIn;"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * "
+                                      << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Calculate bias
+  // if (need_handle_bias_) {
+  //   const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+  //   shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                             << "    values = values + uniforms.beta * "
+  //                             << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"
+  //                             << "  }\n";
+  // }
+
+  // Write output
+  // shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                           << "    " << output.SetByOffset("m * uniforms.N4 + n", "values") << "\n"
+  //                           << "  }\n";
+
+  shader.MainFunctionBody() << " mm_write(m, n, values);\n";
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b,
+                      const Tensor* c) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t A_rows = onnxruntime::narrow<uint32_t>(a_shape[0]);
+  uint32_t A_cols = onnxruntime::narrow<uint32_t>(a_shape[1]);
+  uint32_t B_rows = onnxruntime::narrow<uint32_t>(b_shape[0]);
+  uint32_t B_cols = onnxruntime::narrow<uint32_t>(b_shape[1]);
+
+  if (A_rows % 4 != 0 || A_cols % 4 != 0 || B_rows % 4 != 0 || B_cols % 4 != 0) {
+    return false;
+  }
+
+  // C is optional and is might scalar or 1D.
+  // We do vec4 for C so we need to check if C is 2D and its second dimension is divisible by 4.
+  if (c == nullptr) {
+    return true;
+  }
+
+  const auto& c_shape = c->Shape();
+  if (c_shape.NumDimensions() == 2) {
+    uint32_t C_cols = onnxruntime::narrow<uint32_t>(c_shape[1]);
+    return C_cols % 4 == 0;
+  }
+
+  uint32_t C_rows = onnxruntime::narrow<uint32_t>(c_shape[0]);
+  return C_rows % 4 == 0;
+}
+
+Status ApplyGemmVec4(const Tensor* a,
+                     const Tensor* b,
+                     const Tensor* c,
+                     bool transA_,
+                     bool transB_,
+                     float alpha,
+                     float beta,
+                     ComputeContext& context,
+                     Tensor* y) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t M = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[1] : a_shape[0]);
+  uint32_t K = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[0] : a_shape[1]);
+  uint32_t N = onnxruntime::narrow<uint32_t>(transB_ ? b_shape[0] : b_shape[1]);
+
+  // WebGPU doesn't support binding a zero-sized buffer, so we need to check if A or B is empty.
+  bool need_handle_matmul = a_shape.Size() > 0 && b_shape.Size() > 0;
+  bool need_handle_bias = c && beta;
+
+  GemmVec4Program program{transA_, transB_, alpha, need_handle_bias, need_handle_matmul};
+
+  const int components = 4;
+
+  if (need_handle_matmul) {
+    program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, components},
+                       {b, ProgramTensorMetadataDependency::TypeAndRank, components}});
+  }
+
+  if (need_handle_bias) {
+    program.AddInput({c, ProgramTensorMetadataDependency::TypeAndRank, components});
+  }
+
+  const uint32_t TILE_SIZE = 32;
+  const uint32_t num_tile_n = (N + TILE_SIZE - 1) / TILE_SIZE;
+  const uint32_t num_tile_m = (M + TILE_SIZE - 1) / TILE_SIZE;
+
+  program.CacheHint(alpha, transA_, transB_)
+      .AddOutputs({{y, ProgramTensorMetadataDependency::TypeAndRank, components}})
+      .SetDispatchGroupSize(num_tile_n * num_tile_m)
+      .SetWorkgroupSize(256, 1, 1)
+      .AddUniformVariables({
+          {static_cast<uint32_t>(num_tile_n)},  // num_tile_n
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: `static_cast<uint32_t>` is not needed since `num_tile_n` is type `uint32_t`. Similar for other places.
created_at: 2025-04-23T02:34:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,323 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) {"
+      << "var values = valuesIn;"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * "
+                                      << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Calculate bias
+  // if (need_handle_bias_) {
+  //   const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+  //   shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                             << "    values = values + uniforms.beta * "
+  //                             << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"
+  //                             << "  }\n";
+  // }
+
+  // Write output
+  // shader.MainFunctionBody() << "  if (m < uniforms.M && n < uniforms.N4) {\n"
+  //                           << "    " << output.SetByOffset("m * uniforms.N4 + n", "values") << "\n"
+  //                           << "  }\n";
+
+  shader.MainFunctionBody() << " mm_write(m, n, values);\n";
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b,
+                      const Tensor* c) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t A_rows = onnxruntime::narrow<uint32_t>(a_shape[0]);
+  uint32_t A_cols = onnxruntime::narrow<uint32_t>(a_shape[1]);
+  uint32_t B_rows = onnxruntime::narrow<uint32_t>(b_shape[0]);
+  uint32_t B_cols = onnxruntime::narrow<uint32_t>(b_shape[1]);
+
+  if (A_rows % 4 != 0 || A_cols % 4 != 0 || B_rows % 4 != 0 || B_cols % 4 != 0) {
+    return false;
+  }
+
+  // C is optional and is might scalar or 1D.
+  // We do vec4 for C so we need to check if C is 2D and its second dimension is divisible by 4.
+  if (c == nullptr) {
+    return true;
+  }
+
+  const auto& c_shape = c->Shape();
+  if (c_shape.NumDimensions() == 2) {
+    uint32_t C_cols = onnxruntime::narrow<uint32_t>(c_shape[1]);
+    return C_cols % 4 == 0;
+  }
+
+  uint32_t C_rows = onnxruntime::narrow<uint32_t>(c_shape[0]);
+  return C_rows % 4 == 0;
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: xiaofeihan1
body: Thanks for the suggestion! I have changed the code. For C's last dimension
- N.Use component = 4. Bias is `C.BroadcastedIndicesToOffset("vec2(row, col)", output)`
- 1.Use component = 1
  - shape size == 1. Bias is `vec4(C[0])` 
  - shape size == M. Bias is `vec4(C[row])`
Is this OK?

created_at: 2025-04-24T00:23:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -22,6 +22,10 @@ class GemmProgram final : public Program<GemmProgram> {
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
 
path: onnxruntime/core/providers/webgpu/math/gemm.h
user: login: xiaofeihan1
body: delete
created_at: 2025-04-24T03:21:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n"
+      << "var values = valuesIn; \n"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * ";
+    if (c_component_ == 1 && c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else if (c_component_ == 1) {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    } else {
+      shader.AdditionalImplementation() << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+    }
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  shader.MainFunctionBody() << " mm_write(m, n, values);\n";
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t A_rows = onnxruntime::narrow<uint32_t>(a_shape[0]);
+  uint32_t A_cols = onnxruntime::narrow<uint32_t>(a_shape[1]);
+  uint32_t B_rows = onnxruntime::narrow<uint32_t>(b_shape[0]);
+  uint32_t B_cols = onnxruntime::narrow<uint32_t>(b_shape[1]);
+
+  return A_rows % 4 == 0 && A_cols % 4 == 0 && B_rows % 4 == 0 && B_cols % 4 == 0;
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: Please loose the limitations. It's too strict so that only few cases can go to this path. Maybe `a_shape[1] % 4 == 0` and `b_shape[1]%4 == 0`  are enough.
created_at: 2025-04-24T03:31:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n"
+      << "var values = valuesIn; \n"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * ";
+    if (c_component_ == 1 && c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else if (c_component_ == 1) {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    } else {
+      shader.AdditionalImplementation() << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: Can we always get C's value like below no matter `c_component_` and  `c_is_scalar_`'s value?
```
shader.AdditionalImplementation() << " values += uniforms.beta * vec4<C_element_t>(";
    << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ");\n";` 
created_at: 2025-04-24T04:57:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n"
+      << "var values = valuesIn; \n"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * ";
+    if (c_component_ == 1 && c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else if (c_component_ == 1) {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    } else {
+      shader.AdditionalImplementation() << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: xiaofeihan1
body: The reason I wrote with this way is that`BroadcastedIndicesToOffset` requires that the component dimensions of `output` and `C` be identical. 
- When `c` components is 1 but `output` component is 4, broadcast will occur the error. e.g. A{3, 4}, B{4, 4}, C{3, 1}, Y{3*4}.
Given other comment you suggested that making `output` component is 1 or 4 depends on `N % 4 == 0`. I need to check all cases for C and output. That might be more complex than now.

created_at: 2025-04-24T09:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n"
+      << "var values = valuesIn; \n"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * ";
+    if (c_component_ == 1 && c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else if (c_component_ == 1) {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    } else {
+      shader.AdditionalImplementation() << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: xiaofeihan1
body: I take changeable output components into account and update the code.

| C's components | output's components | Note|
|---|---|---|
| 1 | 1 | Use broadcast|
| 1 | 4 | C is either {M,1} or {1,1}, so we can use C[i]/C[0] instead of broadcast|
| 4 | 1 | Such a case is impossible.|
| 4 | 4 | Use broadcast |

ps: 
`c_components = (c_last_dim == N && N % 4 == 0) ? 4 : 1;`
`output_components = N % 4 == 0 ? 4 : 1;`
created_at: 2025-04-24T10:33:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  shader.AdditionalImplementation()
+      << R"READA(fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+
+    if(col < total_cols&& row < total_rows) {
+        return A[row * total_cols+ col];
+    } else {
+        return output_value_t(0);
+    }
+}
+)READA";
+
+  shader.AdditionalImplementation()
+      << R"READB(fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> output_value_t {
+    if(col < total_cols && row < total_rows) {
+        return B[row * total_cols + col];
+    } else {
+        return output_value_t(0);
+    }
+    })READB";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n"
+      << "var values = valuesIn; \n"
+      << "if(col < uniforms.N4 && row < uniforms.M) {";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += uniforms.beta * ";
+    if (c_component_ == 1 && c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else if (c_component_ == 1) {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    } else {
+      shader.AdditionalImplementation() << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ";\n";
+    }
+  }
+  shader.AdditionalImplementation() << "    output[row * uniforms.N4 + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+
+  shader.MainFunctionBody() << "  var values = output_value_t(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array<output_value_t, 8>, 32>;\n"
+                                      << "var<workgroup> tile_b: array<array<output_value_t, 8>, 32>;\n\n";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u; // row is A index
+        var col = k_start_a + (local_idx % 8u); //
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u); // row is B index
+        col = tile_col_start + (local_idx % 8u); // col is B index
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a1 = output_value_t(0);
+        var a2 = output_value_t(0);
+        var a3 = output_value_t(0);
+        var a4 = output_value_t(0);
+
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a1 = tile_a[i][local_idx / 32u];
+            a2 = tile_a[i + 1u][local_idx / 32u];
+            a3 = tile_a[i + 2u][local_idx / 32u];
+            a4 = tile_a[i + 3u][local_idx / 32u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b = output_value_t(0);
+
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            a = tile_a[i][local_idx / 32u];
+            b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      // !transA_ && !transB_
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+
+         for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<f32>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        var a = output_value_t(0);
+        var b1 = output_value_t(0);
+        var b2 = output_value_t(0);
+        var b3 = output_value_t(0);
+        var b4 = output_value_t(0);
+        for (var i = 0u; i < 32; i = i + 4u) {
+            a = tile_a[local_idx / 8u][i/4u];
+            b1 = tile_b[i][local_idx % 8u];
+            b2 = tile_b[i+1][local_idx % 8u];
+            b3 = tile_b[i+2][local_idx % 8u];
+            b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = uniforms.alpha * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  shader.MainFunctionBody() << " mm_write(m, n, values);\n";
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t A_rows = onnxruntime::narrow<uint32_t>(a_shape[0]);
+  uint32_t A_cols = onnxruntime::narrow<uint32_t>(a_shape[1]);
+  uint32_t B_rows = onnxruntime::narrow<uint32_t>(b_shape[0]);
+  uint32_t B_cols = onnxruntime::narrow<uint32_t>(b_shape[1]);
+
+  return A_rows % 4 == 0 && A_cols % 4 == 0 && B_rows % 4 == 0 && B_cols % 4 == 0;
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: xiaofeihan1
body: Thanks for suggestion. Done. We added some code to support this, because which needs to consider the case where output may not be vec4 when B is transpose.
created_at: 2025-04-25T08:48:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,314 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+#include "core/providers/webgpu/webgpu_utils.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.AdditionalImplementation()
+      << "fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << " { \n"
+      << " if(col < total_cols && row < total_rows) {\n"
+      << "    return A[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+
+  shader.AdditionalImplementation()
+      << "fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << "{ \n"
+      << "  if(col < total_cols && row < total_rows) {\n"
+      << "    return B[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n";
+
+  if (output_components_ == 1) {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N; \n";
+  } else {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N4; \n";
+  }
+
+  shader.AdditionalImplementation() << "var values = valuesIn; \n"
+                                    << "if(col < total_cols && row < uniforms.M) { \n";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += output_element_t(uniforms.beta) * ";
+    // We can be allowed to use broadcasting only when both components are equal.
+    // There is only one case for c_components_ is not equal output_components_.
+    // I.g. the former is `1` and the latter is `4`.
+    // That means the shape of C is either {M,1} or {1,1}
+    if (c_components_ == output_components_) {
+      shader.AdditionalImplementation() << "output_value_t("
+                                        << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ");\n";
+    } else if (c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    }
+  }
+  shader.AdditionalImplementation() << "    output[row * total_cols + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.MainFunctionBody() << "  var values = " << type_string << "(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array< " << type_string << ", 8 >, 32 >;\n "
+                                      << "var<workgroup> tile_b: array<array< " << type_string << ", 8 >, 32 >;\n ";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u;
+        var col = k_start_a + (local_idx % 8u);
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u);
+        col = tile_col_start + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a1 = tile_a[i][local_idx / 32u];
+            let a2 = tile_a[i + 1u][local_idx / 32u];
+            let a3 = tile_a[i + 2u][local_idx / 32u];
+            let a4 = tile_a[i + 3u][local_idx / 32u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            let a = tile_a[i][local_idx / 32u];
+            let b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+         for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<output_element_t>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[i][local_idx % 8u];
+            let b2 = tile_b[i+1][local_idx % 8u];
+            let b3 = tile_b[i+2][local_idx % 8u];
+            let b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = output_element_t(uniforms.alpha) * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Write output
+  if (output_components_ == 1) {
+    shader.MainFunctionBody() << " for (var i = 0u; i < 4u; i = i + 1u) {\n"
+                              << "    mm_write(m, 4 * n + i, values[i]);\n"
+                              << "  }\n";
+  } else {
+    shader.MainFunctionBody() << "  mm_write(m, n, values);\n";
+  }
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  // When the number of columns in A and B is divisible by 4, we apply vec4 optimization to A and B.
+  // However, this doesn't necessarily mean that C and Y will use vec4.
+  // For example, C/output won't be vec4 if B is transposed and N is not divisible by 4.
+  // Also, C won't use vec4 when it's a scalar.
+  // The code would be simpler if we avoided vec4 optimization for C/output.
+  // But to maximize performance, we still apply vec4 when possible — even though it adds some complexity.
+  // I've added detailed comments explaining this logic.
+  // See MatMulReadFnSource and MatMulWriteFnSource, especially the parts related to broadcasting.
+  return a_shape[1] % 4 == 0 && b_shape[1] % 4 == 0;
+}
+
+Status ApplyGemmVec4(const Tensor* a,
+                     const Tensor* b,
+                     const Tensor* c,
+                     bool transA_,
+                     bool transB_,
+                     float alpha,
+                     float beta,
+                     ComputeContext& context,
+                     Tensor* y) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t M = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[1] : a_shape[0]);
+  uint32_t K = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[0] : a_shape[1]);
+  uint32_t N = onnxruntime::narrow<uint32_t>(transB_ ? b_shape[0] : b_shape[1]);
+
+  // WebGPU doesn't support binding a zero-sized buffer, so we need to check if A or B is empty.
+  bool need_handle_matmul = a_shape.Size() > 0 && b_shape.Size() > 0;
+  bool need_handle_bias = c && beta;
+
+  int c_components = 4;
+  bool c_is_scalar = false;
+
+  // We use vec4 for C when its last dimension equals N and N is divisible by 4.
+  if (need_handle_bias) {
+    const auto& c_shape = c->Shape();
+    int64_t c_last_dim = c_shape[c_shape.NumDimensions() - 1];
+    c_components = (c_last_dim == N && N % 4 == 0) ? 4 : 1;
+    c_is_scalar = c_shape.Size() == 1;
+  }
+
+  // We use vec4 for Y when N is divisible by 4.
+  const int output_components = N % 4 == 0 ? 4 : 1;
+
+  GemmVec4Program program{transA_, transB_, alpha, need_handle_bias, need_handle_matmul, c_components, c_is_scalar, output_components};
+
+  const int components = 4;
+
+  if (need_handle_matmul) {
+    program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, components},
+                       {b, ProgramTensorMetadataDependency::TypeAndRank, components}});
+  }
+
+  if (need_handle_bias) {
+    program.AddInput({c, ProgramTensorMetadataDependency::TypeAndRank, c_components});
+  }
+
+  const uint32_t TILE_SIZE = 32;
+  const uint32_t num_tile_n = (N + TILE_SIZE - 1) / TILE_SIZE;
+  const uint32_t num_tile_m = (M + TILE_SIZE - 1) / TILE_SIZE;
+
+  program.CacheHint(alpha, transA_, transB_)
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: qjia7
body: `c_is_scalar` also needs to be added into `CacheHint`?
created_at: 2025-04-28T01:33:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,314 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/gemm_vec4.h"
+
+#include "core/providers/webgpu/webgpu_utils.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+void GemmVec4Program::MatMulReadFnSource(ShaderHelper& shader) const {
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.AdditionalImplementation()
+      << "fn mm_readA(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << " { \n"
+      << " if(col < total_cols && row < total_rows) {\n"
+      << "    return A[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+
+  shader.AdditionalImplementation()
+      << "fn mm_readB(row: u32, col: u32, total_rows: u32, total_cols: u32) -> " << type_string << "{ \n"
+      << "  if(col < total_cols && row < total_rows) {\n"
+      << "    return B[row * total_cols + col];\n"
+      << "  } else {\n"
+      << "    return " << type_string << "(0);\n"
+      << "  }\n"
+      << "}\n\n";
+}
+
+void GemmVec4Program::MatMulWriteFnSource(ShaderHelper& shader, const ShaderVariableHelper& output) const {
+  shader.AdditionalImplementation()
+      << "fn mm_write(row: u32, col: u32, valuesIn: output_value_t) { \n";
+
+  if (output_components_ == 1) {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N; \n";
+  } else {
+    shader.AdditionalImplementation() << "  let total_cols = uniforms.N4; \n";
+  }
+
+  shader.AdditionalImplementation() << "var values = valuesIn; \n"
+                                    << "if(col < total_cols && row < uniforms.M) { \n";
+  if (need_handle_bias_) {
+    const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform);
+    shader.AdditionalImplementation() << "    values += output_element_t(uniforms.beta) * ";
+    // We can be allowed to use broadcasting only when both components are equal.
+    // There is only one case for c_components_ is not equal output_components_.
+    // I.g. the former is `1` and the latter is `4`.
+    // That means the shape of C is either {M,1} or {1,1}
+    if (c_components_ == output_components_) {
+      shader.AdditionalImplementation() << "output_value_t("
+                                        << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(row, col)", output)) << ");\n";
+    } else if (c_is_scalar_) {
+      shader.AdditionalImplementation() << "output_value_t(C[0]);\n";
+    } else {
+      shader.AdditionalImplementation() << "output_value_t(C[row]);\n";
+    }
+  }
+  shader.AdditionalImplementation() << "    output[row * total_cols + col] = values;\n"
+                                    << "  }\n"
+                                    << "}\n";
+}
+
+Status GemmVec4Program::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+
+  // We can’t treat `output_value_t` as the type of A and B, because output might not be a vec4, while A or B is.
+  const std::string data_type = "output_element_t";
+  const std::string type_string = MakeScalarOrVectorType(4 /*components */, data_type);
+
+  shader.MainFunctionBody() << "  var values = " << type_string << "(0);\n\n"
+                            << "  let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * 8u;\n"
+                            << "  let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * 32u;\n";
+
+  if (need_handle_matmul_) {
+    shader.AddInput("A", ShaderUsage::UseUniform);
+    shader.AddInput("B", ShaderUsage::UseUniform);
+
+    MatMulReadFnSource(shader);
+
+    // Add shared memory arrays for tiling
+    shader.AdditionalImplementation() << "var<workgroup> tile_a: array<array< " << type_string << ", 8 >, 32 >;\n "
+                                      << "var<workgroup> tile_b: array<array< " << type_string << ", 8 >, 32 >;\n ";
+
+    shader.MainFunctionBody()
+        << "  var k_start_a = 0u;\n"
+        << "  var k_start_b = 0u;\n\n"
+        << "  let num_tiles = (uniforms.K + 32 - 1) / 32;\n";
+
+    // Main loop for matrix multiplication
+    shader.MainFunctionBody()
+        << "  for (var t = 0u; t < num_tiles; t = t + 1u) {\n";
+    // Load TILE_A
+    if (transA_) {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = k_start_a + (local_idx / 8u);
+        var col =  tile_row_start/4 + local_idx % 8u;
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.K, uniforms.M4);
+        )TILE_A";
+    } else {
+      shader.MainFunctionBody() << R"TILE_A(
+        var row = tile_row_start + local_idx / 8u;
+        var col = k_start_a + (local_idx % 8u);
+        tile_a[local_idx / 8u][local_idx % 8u] = mm_readA(row, col, uniforms.M, uniforms.K4);
+        )TILE_A";
+    }
+    // Load TILE_B
+    if (transB_) {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = tile_col_start * 4 + (local_idx / 8u);
+        col = k_start_b + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.N, uniforms.K4);
+        )TILE_B";
+    } else {
+      shader.MainFunctionBody() << R"TILE_B(
+        row = k_start_b + (local_idx / 8u);
+        col = tile_col_start + (local_idx % 8u);
+        // load 1 vec4 into tile_b
+        tile_b[local_idx / 8u][local_idx % 8u] = mm_readB(row, col, uniforms.K, uniforms.N4);
+        )TILE_B";
+    }
+
+    shader.MainFunctionBody() << "    workgroupBarrier();\n\n";
+
+    if (transA_) {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 32u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_a = k_start_a + 8u; \n";
+    }
+
+    if (transB_) {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 8u; \n";
+    } else {
+      shader.MainFunctionBody() << "k_start_b = k_start_b + 32u; \n";
+    }
+
+    // Calculate output according to TILE_A and TILE_B
+    if (transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a1 = tile_a[i][local_idx / 32u];
+            let a2 = tile_a[i + 1u][local_idx / 32u];
+            let a3 = tile_a[i + 2u][local_idx / 32u];
+            let a4 = tile_a[i + 3u][local_idx / 32u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            var vec_idx = local_idx / 8u % 4;
+
+            values[0] += a1[vec_idx] * b1[0] + a2[vec_idx] * b1[1] + a3[vec_idx] * b1[2] + a4[vec_idx] * b1[3];
+            values[1] += a1[vec_idx] * b2[0] + a2[vec_idx] * b2[1] + a3[vec_idx] * b2[2] + a4[vec_idx] * b2[3];
+            values[2] += a1[vec_idx] * b3[0] + a2[vec_idx] * b3[1] + a3[vec_idx] * b3[2] + a4[vec_idx] * b3[3];
+            values[3] += a1[vec_idx] * b4[0] + a2[vec_idx] * b4[1] + a3[vec_idx] * b4[2] + a4[vec_idx] * b4[3];
+        }
+        )CALC";
+    } else if (transA_ && !transB_) {
+      shader.MainFunctionBody() << R"CALC(
+        // Calculate 4 output for each thread
+        // We read 32 vec4 from tile_a and 32 vec4 from tile_b in total.
+        for (var i = 0u; i < 32; i = i + 1u) {
+            let a = tile_a[i][local_idx / 32u];
+            let b = tile_b[i][local_idx % 8u];
+            values += a[(local_idx / 8u) % 4] * b;
+        })CALC";
+    } else if (!transA_ && transB_) {
+      shader.MainFunctionBody() << R"CALC(
+         for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[(local_idx % 8) * 4][i / 4u];
+            let b2 = tile_b[(local_idx % 8) * 4 + 1u][i / 4u];
+            let b3 = tile_b[(local_idx % 8) * 4 + 2u][i / 4u];
+            let b4 = tile_b[(local_idx % 8) * 4 + 3u][i / 4u];
+
+            values += vec4<output_element_t>(
+                dot(a, b1),
+                dot(a, b2),
+                dot(a, b3),
+                dot(a, b4)
+            );
+        }
+            )CALC";
+    } else {
+      shader.MainFunctionBody() << R"CALC(
+        for (var i = 0u; i < 32; i = i + 4u) {
+            let a = tile_a[local_idx / 8u][i/4u];
+            let b1 = tile_b[i][local_idx % 8u];
+            let b2 = tile_b[i+1][local_idx % 8u];
+            let b3 = tile_b[i+2][local_idx % 8u];
+            let b4 = tile_b[i+3][local_idx % 8u];
+
+            values += a.x * b1 + a.y * b2 + a.z * b3 + a.w * b4;
+        }
+        )CALC";
+    }
+    shader.MainFunctionBody() << "    workgroupBarrier();\n"
+                              << "  }\n\n";
+
+    // Calculate alpha
+    if (alpha_ != 1.0f) {
+      shader.MainFunctionBody() << "  values = output_element_t(uniforms.alpha) * values;\n";
+    }
+  }
+
+  MatMulWriteFnSource(shader, output);
+  shader.MainFunctionBody() << "  let m = tile_row_start + local_idx / 8u;\n"
+                            << "  let n = tile_col_start + local_idx % 8u;\n\n";
+
+  // Write output
+  if (output_components_ == 1) {
+    shader.MainFunctionBody() << " for (var i = 0u; i < 4u; i = i + 1u) {\n"
+                              << "    mm_write(m, 4 * n + i, values[i]);\n"
+                              << "  }\n";
+  } else {
+    shader.MainFunctionBody() << "  mm_write(m, n, values);\n";
+  }
+
+  return Status::OK();
+}
+
+bool CanApplyGemmVec4(const Tensor* a,
+                      const Tensor* b) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  // When the number of columns in A and B is divisible by 4, we apply vec4 optimization to A and B.
+  // However, this doesn't necessarily mean that C and Y will use vec4.
+  // For example, C/output won't be vec4 if B is transposed and N is not divisible by 4.
+  // Also, C won't use vec4 when it's a scalar.
+  // The code would be simpler if we avoided vec4 optimization for C/output.
+  // But to maximize performance, we still apply vec4 when possible — even though it adds some complexity.
+  // I've added detailed comments explaining this logic.
+  // See MatMulReadFnSource and MatMulWriteFnSource, especially the parts related to broadcasting.
+  return a_shape[1] % 4 == 0 && b_shape[1] % 4 == 0;
+}
+
+Status ApplyGemmVec4(const Tensor* a,
+                     const Tensor* b,
+                     const Tensor* c,
+                     bool transA_,
+                     bool transB_,
+                     float alpha,
+                     float beta,
+                     ComputeContext& context,
+                     Tensor* y) {
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+
+  uint32_t M = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[1] : a_shape[0]);
+  uint32_t K = onnxruntime::narrow<uint32_t>(transA_ ? a_shape[0] : a_shape[1]);
+  uint32_t N = onnxruntime::narrow<uint32_t>(transB_ ? b_shape[0] : b_shape[1]);
+
+  // WebGPU doesn't support binding a zero-sized buffer, so we need to check if A or B is empty.
+  bool need_handle_matmul = a_shape.Size() > 0 && b_shape.Size() > 0;
+  bool need_handle_bias = c && beta;
+
+  int c_components = 4;
+  bool c_is_scalar = false;
+
+  // We use vec4 for C when its last dimension equals N and N is divisible by 4.
+  if (need_handle_bias) {
+    const auto& c_shape = c->Shape();
+    int64_t c_last_dim = c_shape[c_shape.NumDimensions() - 1];
+    c_components = (c_last_dim == N && N % 4 == 0) ? 4 : 1;
+    c_is_scalar = c_shape.Size() == 1;
+  }
+
+  // We use vec4 for Y when N is divisible by 4.
+  const int output_components = N % 4 == 0 ? 4 : 1;
+
+  GemmVec4Program program{transA_, transB_, alpha, need_handle_bias, need_handle_matmul, c_components, c_is_scalar, output_components};
+
+  const int components = 4;
+
+  if (need_handle_matmul) {
+    program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, components},
+                       {b, ProgramTensorMetadataDependency::TypeAndRank, components}});
+  }
+
+  if (need_handle_bias) {
+    program.AddInput({c, ProgramTensorMetadataDependency::TypeAndRank, c_components});
+  }
+
+  const uint32_t TILE_SIZE = 32;
+  const uint32_t num_tile_n = (N + TILE_SIZE - 1) / TILE_SIZE;
+  const uint32_t num_tile_m = (M + TILE_SIZE - 1) / TILE_SIZE;
+
+  program.CacheHint(alpha, transA_, transB_)
path: onnxruntime/core/providers/webgpu/math/gemm_vec4.cc
user: login: xiaofeihan1
body: Thanks for catching this. Added it.
created_at: 2025-04-28T03:20:04Z
author_association: CONTRIBUTOR

PR comments:
user: login: fs-eire
created_at: 2025-04-22T00:00:38Z
body: Is there a way to reuse the implementation of `MatMul`? My understanding is that there are some kind of duplication between GEMM and MatMul, and it would be great if we can reuse the shared code
author_association: CONTRIBUTOR
user: login: xiaofeihan1
created_at: 2025-04-22T07:26:52Z
body: > Is there a way to reuse the implementation of `MatMul`? My understanding is that there are some kind of duplication between GEMM and MatMul, and it would be great if we can reuse the shared code

Thanks for the callout. That's what I'm gonna do next.
For current PR, I want to push forward to support vec4 for GEMM. I will take the refactor work in future PRs because it also require some effort to consider(e.g.There are also some differences between gemm and matmul, e.g. the latter supports batch size, the former supports transpose, etc). WDYT?
author_association: CONTRIBUTOR
user: login: fs-eire
created_at: 2025-04-22T18:44:03Z
body: /azp run Linux QNN CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI, Windows ARM64 QNN CI Pipeline, Windows GPU Doc Gen CI Pipeline, Windows x64 QNN CI Pipeline
author_association: CONTRIBUTOR
user: login: azure-pipelines[bot]
created_at: 2025-04-22T18:44:24Z
body: <samp>
Azure Pipelines successfully started running 5 pipeline(s).<br>

</samp>
author_association: NONE
user: login: fs-eire
created_at: 2025-04-24T21:10:02Z
body: /azp run Linux QNN CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI, Windows ARM64 QNN CI Pipeline, Windows GPU Doc Gen CI Pipeline, Windows x64 QNN CI Pipeline
author_association: CONTRIBUTOR
user: login: azure-pipelines[bot]
created_at: 2025-04-24T21:10:24Z
body: <samp>
Azure Pipelines successfully started running 5 pipeline(s).<br>

</samp>
author_association: NONE
user: login: fs-eire
created_at: 2025-04-28T23:12:43Z
body: /azp run Windows ARM64 QNN CI Pipeline,Windows x64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline,ONNX Runtime Web CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline
author_association: CONTRIBUTOR
user: login: fs-eire
created_at: 2025-04-28T23:12:44Z
body: /azp run Linux QNN CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,Big Models,Linux Android Emulator QNN CI Pipeline,Android CI Pipeline,iOS CI Pipeline,ONNX Runtime React Native CI Pipeline,Linux DNNL CI Pipeline,Linux MIGraphX CI Pipeline,Linux ROCm CI Pipeline
author_association: CONTRIBUTOR
user: login: azure-pipelines[bot]
created_at: 2025-04-28T23:13:04Z
body: <samp>
Azure Pipelines successfully started running 5 pipeline(s).<br>

</samp>
author_association: NONE
user: login: azure-pipelines[bot]
created_at: 2025-04-28T23:13:11Z
body: <samp>
Azure Pipelines successfully started running 6 pipeline(s), but failed to run 1 pipeline(s).<br>

</samp>
author_association: NONE