webgpu: optimize Gemm and MatMul using subgroup feature

onnxruntime/core/providers/webgpu/math/gemm.cc

@@ -3,6 +3,7 @@
 
 #include "core/providers/webgpu/math/gemm.h"
 #include "core/providers/webgpu/math/gemm_packed.h"
+#include "core/providers/webgpu/vendor/intel/math/gemm.h"
 
 #include <vector>
 
@@ -147,6 +148,10 @@ Status Gemm::ComputeInternal(ComputeContext& context) const {
     return context.RunProgram(program);
   }
 
+  if (intel::CanApplyGemmIntel(context, M, N, K, transA_, transB_)) {
+    return intel::ApplyGemmIntel(A, B, C, transA_, transB_, alpha_, beta_, context);
+  }
+
   return ApplyGemmPacked(A, B, C, transA_, transB_, alpha_, beta_, context);
 }
 

onnxruntime/core/providers/webgpu/math/gemm_utils.cc

@@ -49,7 +49,7 @@ void HandleMaybeBiasForMatMul(ShaderHelper& shader,
     shader.AdditionalImplementation() << "    value = value + output_value_t(" << (is_channels_last ? bias->GetByOffset("colIn") : bias->GetByOffset("row")) << ");\n";
   }
   shader.AdditionalImplementation() << "    " << activation_snippet << "\n"
-                                    << output.SetByIndices("coords", "value") << "\n";
+                                    << "    " << output.SetByIndices("coords", "value") << "\n";
 }
 
 void HandleMatMulWithSplitK(
@@ -123,59 +123,63 @@ void MatMulReadFnSource(ShaderHelper& shader,
                         const ShaderIndicesHelper* batch_dims,
                         bool transA,
                         bool transB,
-                        bool is_vec4) {
-  int components = is_vec4 ? 4 : 1;
+                        bool is_vec4,
+                        bool use_subgroup) {
+  // Always read A with 1-component when using subgroup.
+  int components = use_subgroup ? 1 : (is_vec4 ? 4 : 1);
   const std::string data_type = "output_element_t";
-  const std::string type_string = MakeScalarOrVectorType(components, data_type);
+  std::string type_string = MakeScalarOrVectorType(components, data_type);
 
   shader.AdditionalImplementation()
       << "fn mm_readA(batch: i32, row: i32, colIn: i32 "
       << (batch_dims
               ? ", batch_indices: batch_dims_indices_t"
               : "")
-      << ") -> " << type_string << " {\n "
-      << "    var value = " << type_string << "(0);\n"
-      << "    let col = colIn * " << components << ";\n";
+      << ") -> " << type_string << " {\n"
+      << "  var value = " << type_string << "(0);\n"
+      << "  let col = colIn * " << components << ";\n";
   if (transA) {
-    shader.AdditionalImplementation() << "    if(row < i32(uniforms.dim_inner) && col < i32(uniforms.dim_a_outer)) {\n";
+    shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_inner) && col < i32(uniforms.dim_a_outer)) {\n";
   } else {
-    shader.AdditionalImplementation() << "    if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_inner)) {\n";
+    shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_inner)) {\n";
   }
-  shader.AdditionalImplementation() << "        var a_indices: a_indices_t;\n";
+  shader.AdditionalImplementation() << "    var a_indices: a_indices_t;\n";
 
   if (batch_dims) {
-    shader.AdditionalImplementation() << ConvertOutputBatchIndicesToInputBatchIndices("a", a, a.Rank() - 2, batch_dims ? batch_dims->Rank() : 0, " batch_indices ") << "\n";
+    shader.AdditionalImplementation() << ConvertOutputBatchIndicesToInputBatchIndices("a", a, a.Rank() - 2, batch_dims ? batch_dims->Rank() : 0, " batch_indices ");
   }
-  shader.AdditionalImplementation() << a.IndicesSet("a_indices", a.Rank() - 2, "u32(row)") << "\n"
-                                    << a.IndicesSet("a_indices", a.Rank() - 1, "u32(colIn)") << "\n"
-                                    << "        value = " << a.GetByIndices("a_indices") << ";\n"
-                                    << "    }\n"
-                                    << "    return value;\n"
+  shader.AdditionalImplementation() << "    " << a.IndicesSet("a_indices", a.Rank() - 2, "u32(row)") << "\n"
+                                    << "    " << a.IndicesSet("a_indices", a.Rank() - 1, "u32(colIn)") << "\n"
+                                    << "    value = " << a.GetByIndices("a_indices") << ";\n"
+                                    << "  }\n"
+                                    << "  return value;\n"
                                     << "}\n\n";
 
   // Add the mm_readB function
+  components = is_vec4 ? 4 : 1;
+  type_string = MakeScalarOrVectorType(components, data_type);
   shader.AdditionalImplementation()
       << "fn mm_readB(batch: i32, row: i32, colIn: i32 "
       << (batch_dims
               ? ", batch_indices: batch_dims_indices_t"
               : "")
-      << ") -> " << type_string << " {\n "
-      << "    var value = " << type_string << "(0);\n"
-      << "    let col = colIn * " << components << ";\n";
+      << ") -> " << type_string << " {\n"
+      << "  var value = " << type_string << "(0);\n"
+      << "  let col = colIn * " << components << ";\n";
 
   if (transB) {
-    shader.AdditionalImplementation() << "    if(row < i32(uniforms.dim_b_outer) && col < i32(uniforms.dim_inner)) {\n";
+    shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_b_outer) && col < i32(uniforms.dim_inner)) {\n";
   } else {
-    shader.AdditionalImplementation() << "    if(row < i32(uniforms.dim_inner) && col < i32(uniforms.dim_b_outer)) {\n";
+    shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_inner) && col < i32(uniforms.dim_b_outer)) {\n";
   }
 
-  shader.AdditionalImplementation() << "        var b_indices: b_indices_t;\n"
+  shader.AdditionalImplementation() << "    var b_indices: b_indices_t;\n"
                                     << ConvertOutputBatchIndicesToInputBatchIndices("b", b, b.Rank() - 2, batch_dims ? batch_dims->Rank() : 0, "batch_indices")
-                                    << b.IndicesSet("b_indices", b.Rank() - 2, "u32(row)") << "\n"
-                                    << b.IndicesSet("b_indices", b.Rank() - 1, "u32(colIn)") << "\n"
-                                    << "        value = " << b.GetByIndices("b_indices") << ";\n"
-                                    << "    }\n"
-                                    << "    return value;\n"
+                                    << "    " << b.IndicesSet("b_indices", b.Rank() - 2, "u32(row)") << "\n"
+                                    << "    " << b.IndicesSet("b_indices", b.Rank() - 1, "u32(colIn)") << "\n"
+                                    << "    value = " << b.GetByIndices("b_indices") << ";\n"
+                                    << "  }\n"
+                                    << "  return value;\n"
                                     << "}\n\n";
 }
 
@@ -191,12 +195,12 @@ void MatMulWriteFnSource(ShaderHelper& shader,
                          bool use_split_k,
                          ProgramVariableDataType output_variable_type) {
   shader.AdditionalImplementation()
-      << "fn mm_write(batch: i32, row: i32, colIn: i32, valueIn: output_value_t) { \n";
+      << "fn mm_write(batch: i32, row: i32, colIn: i32, valueIn: output_value_t) {\n";
 
   shader.AdditionalImplementation() << "  let col = colIn * " << output_components << ";\n";
 
-  shader.AdditionalImplementation() << "if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_b_outer)) { \n"
-                                    << "    var value = valueIn; \n";
+  shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_b_outer)) {\n"
+                                    << "    var value = valueIn;\n";
 
   if (use_split_k) {
     // Set output when MatMul is performed with Split-K.

onnxruntime/core/providers/webgpu/math/gemm_utils.h

@@ -14,7 +14,8 @@ void MatMulReadFnSource(ShaderHelper& shader,
                         const ShaderIndicesHelper* batch_dims,
                         bool transA,
                         bool transB,
-                        bool is_vec4);
+                        bool is_vec4,
+                        bool use_subgroup = false);
 
 void MatMulWriteFnSource(ShaderHelper& shader,
                          const ShaderVariableHelper& output,

onnxruntime/core/providers/webgpu/math/matmul.cc

@@ -8,6 +8,7 @@
 #include "core/providers/webgpu/webgpu_supported_types.h"
 #include "core/providers/webgpu/nn/fuse_utils.h"
 #include "core/providers/webgpu/data_transfer.h"
+#include "core/providers/webgpu/vendor/intel/math/matmul.h"
 
 namespace onnxruntime {
 namespace webgpu {
@@ -162,6 +163,10 @@ Status MatMul::ComputeInternal(ComputeContext& context) const {
     inputs.push_back(bias);
   }
 
+  if (intel::CanApplyMatMulIntel(context, helper.M(), helper.N(), helper.K())) {
+    return intel::ApplyMatMulIntel(context, Activation(), inputs, output_tensor);
+  }
+
   return ComputeMatMul(&context, Activation(), inputs, output_tensor, false);
 }
 

onnxruntime/core/providers/webgpu/vendor/intel/math/gemm.cc

@@ -0,0 +1,121 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/vendor/intel/math/gemm.h"
+#include "core/providers/webgpu/vendor/intel/math/gemm_subgroup.h"
+#include "core/providers/webgpu/math/gemm_utils.h"
+
+namespace onnxruntime {
+namespace webgpu {
+namespace intel {
+
+Status GemmSubgroupProgram::GenerateShaderCode(ShaderHelper& shader) const {
+  const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform |
+                                                                      ShaderUsage::UseValueTypeAlias |
+                                                                      ShaderUsage::UseElementTypeAlias);
+
+  if (need_handle_matmul_) {
+    const auto& a = shader.AddInput("a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias |
+                                             ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+    const auto& b = shader.AddInput("b", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias |
+                                             ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+
+    MatMulReadFnSource(shader, a, b, nullptr, transA_, transB_, is_vec4_, true);
+  }
+
+  ORT_RETURN_IF_ERROR(MakeMatMulSubgroupSource(shader, elements_per_thread_, nullptr, is_vec4_, transA_, transB_,
+                                               alpha_, need_handle_matmul_));
+  const ShaderVariableHelper* c = nullptr;
+  if (need_handle_bias_) {
+    c = &shader.AddInput("c", ShaderUsage::UseUniform);
+  }
+  MatMulWriteFnSource(shader, output, c, true, c_components_, output_components_, c_is_scalar_);
+
+  return Status::OK();
+}
+
+bool CanApplyGemmIntel(const ComputeContext& context, int64_t M, int64_t N, int64_t K, bool transA, bool transB) {
+  return CanApplySubgroup(context, M, N, K, transA, transB);
+}
+
+Status ApplyGemmIntel(const Tensor* a,
+                      const Tensor* b,
+                      const Tensor* c,
+                      bool transA,
+                      bool transB,
+                      float alpha,
+                      float beta,
+                      ComputeContext& context) {
+  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]);
+
+  std::vector<int64_t> output_dims{M, N};
+  auto* y = context.Output(0, output_dims);
+  int64_t output_size = y->Shape().Size();
+
+  if (output_size == 0) {
+    return Status::OK();
+  }
+
+  // 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;
+
+  const bool is_vec4 = b_shape[1] % 4 == 0;
+  // Components for A, B
+  int a_components = 1;
+  int b_components = is_vec4 ? 4 : 1;
+  // Components for Y
+  int output_components = (is_vec4 && N % 4 == 0) ? 4 : 1;
+  // Components for C.
+  int c_components = 1;
+
+  bool c_is_scalar = false;
+  if (need_handle_bias) {
+    const auto& c_shape = c->Shape();
+    int64_t c_last_dim = c_shape[c_shape.NumDimensions() - 1];
+    // `C` in GEMM might be broadcast to the output, and broadcasting requires the components to be consistent.
+    // So we use vec4 for C when its last dimension is N, and the output is also a vec4.
+    c_components = (c_last_dim == N && output_components == 4) ? 4 : 1;
+    c_is_scalar = c_shape.Size() == 1;
+  }
+
+  InlinedVector<int64_t> elements_per_thread = InlinedVector<int64_t>({4, intel::ElementsPerThreadY(is_vec4, M), 1});
+  const uint32_t dispatch_x = narrow<uint32_t>((N + kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0] - 1) /
+                                               (kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0]));
+  const uint32_t dispatch_y = narrow<uint32_t>((M + kSubgroupLogicalWorkGroupSizeY * elements_per_thread[1] - 1) /
+                                               (kSubgroupLogicalWorkGroupSizeY * elements_per_thread[1]));
+
+  GemmSubgroupProgram program{transA, transB, alpha, need_handle_bias, need_handle_matmul, c_components, c_is_scalar,
+                              output_components, is_vec4, elements_per_thread};
+
+  if (need_handle_matmul) {
+    program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, a_components},
+                       {b, ProgramTensorMetadataDependency::TypeAndRank, b_components}});
+  }
+
+  if (need_handle_bias) {
+    program.AddInput({c, ProgramTensorMetadataDependency::TypeAndRank, c_components});
+  }
+
+  program.CacheHint(alpha, transA, transB, c_is_scalar, is_vec4, absl::StrJoin(elements_per_thread, "-"))
+      .AddOutputs({{y, ProgramTensorMetadataDependency::TypeAndRank, output_components}})
+      .SetDispatchGroupSize(dispatch_x, dispatch_y, 1)
+      .SetWorkgroupSize(kSubgroupLogicalWorkGroupSizeX * kSubgroupLogicalWorkGroupSizeY, 1, 1)
+      .AddUniformVariables({{alpha},
+                            {beta},
+                            {M}, /* dim_a_outer */
+                            {N}, /* dim_b_outer */
+                            {K}} /*dim_inner */
+      );
+
+  return context.RunProgram(program);
+}
+
+}  // namespace intel
+}  // namespace webgpu
+}  // namespace onnxruntime

onnxruntime/core/providers/webgpu/vendor/intel/math/gemm.h

@@ -0,0 +1,66 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/providers/webgpu/webgpu_kernel.h"
+#include "core/providers/webgpu/shader_helper.h"
+#include "core/providers/webgpu/program.h"
+
+namespace onnxruntime {
+namespace webgpu {
+namespace intel {
+
+class GemmSubgroupProgram final : public Program<GemmSubgroupProgram> {
+ public:
+  GemmSubgroupProgram(bool transA, bool transB, float alpha, bool need_handle_bias, bool need_handle_matmul,
+                      int c_components, bool c_is_scalar, int output_components, bool is_vec4,
+                      const gsl::span<int64_t>& elements_per_thread)
+      : Program{"GemmSubgroup"},
+        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),
+        is_vec4_(is_vec4),
+        elements_per_thread_(elements_per_thread.begin(), elements_per_thread.end()) {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"alpha", ProgramUniformVariableDataType::Float32},
+      {"beta", ProgramUniformVariableDataType::Float32},
+      {"dim_a_outer", ProgramUniformVariableDataType::Uint32},
+      {"dim_b_outer", ProgramUniformVariableDataType::Uint32},
+      {"dim_inner", ProgramUniformVariableDataType::Uint32});
+
+ 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_;
+  bool is_vec4_ = false;
+  const InlinedVector<int64_t> elements_per_thread_;
+};
+
+bool CanApplyGemmIntel(const ComputeContext& context, int64_t M, int64_t N, int64_t K, bool transA, bool transB);
+
+Status ApplyGemmIntel(const Tensor* a,
+                      const Tensor* b,
+                      const Tensor* c,
+                      bool transA,
+                      bool transB,
+                      float alpha,
+                      float beta,
+                      ComputeContext& context);
+
+}  // namespace intel
+}  // namespace webgpu
+}  // namespace onnxruntime