webgpu: adjust the parms for gemm-subgroup kernel

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

@@ -194,15 +194,15 @@ Status ComputeMatMul(ComputeContext* context,
   // When B is a matrix (batch is 1), we fold batchA into the M dimension for better
   // performance (e.g., [2,3,5] → [1,6,5]).
   if (batchA != 1 && batchB == 1) {
-    // dimensions of A: [1,`batchA`, M, K]
+    // dimensions of A: [`batchA` * M, K]
     int64_t batchAndM = a_shape.SizeToDimension(a_shape.NumDimensions() - 1);
-    TensorShapeVector dims_a = {1, batchAndM, helper.K()};
-    // dimensions of B: [1,K,N]
-    TensorShapeVector dims_b = {1, helper.K(), helper.N()};
+    TensorShapeVector dims_a = {batchAndM, helper.K()};
+    // dimensions of B: [K, N]
+    TensorShapeVector dims_b = {helper.K(), helper.N()};
 
     a_shape = TensorShape(dims_a);
     b_shape = TensorShape(dims_b);
-    output_shape = {1, batchAndM, helper.N()};
+    output_shape = {batchAndM, helper.N()};
   }
 
   // helpful dimension variables

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

@@ -84,7 +84,7 @@ Status ApplyGemmIntel(const Tensor* a,
     c_is_scalar = c_shape.Size() == 1;
   }
 
-  InlinedVector<int64_t> elements_per_thread = InlinedVector<int64_t>({4, intel::ElementsPerThreadY(is_vec4, M), 1});
+  InlinedVector<int64_t> elements_per_thread = InlinedVector<int64_t>({4, intel::ElementsPerThreadY(context, 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) /

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

@@ -94,8 +94,12 @@ bool CanApplySubgroup(const ComputeContext& context, int64_t M, int64_t N, int64
   return false;
 }
 
-int64_t ElementsPerThreadY(bool is_vec4, uint32_t M) {
-  return is_vec4 ? (M <= 8 ? 1 : (M <= 16 ? 2 : (M <= 32 ? 4 : 8))) : 4;
+int64_t ElementsPerThreadY(ComputeContext& context, uint32_t M) {
+  // For Xe-LPG and Xe-3LPG, we have observed that 4 elements per thread is optimal when M is large.
+  const auto& arch = context.AdapterInfo().architecture;
+  const bool is_xe_lpg_or_xe_3lpg = arch == std::string_view("xe-lpg") ||
+                                    arch == std::string_view("xe-3lpg");
+  return M <= 8 ? 1 : (M <= 16 ? 2 : (M <= 32 ? 4 : (is_xe_lpg_or_xe_3lpg ? 4 : 8)));
 }
 
 Status MakeMatMulSubgroupSource(ShaderHelper& shader,

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

@@ -15,7 +15,7 @@ const uint32_t kSubgroupLogicalWorkGroupSizeZ = 1;
 
 bool CanApplySubgroup(const ComputeContext& context, int64_t M, int64_t N, int64_t K, bool transA = false, bool transB = false);
 
-int64_t ElementsPerThreadY(bool is_vec4, uint32_t M);
+int64_t ElementsPerThreadY(ComputeContext& context, uint32_t M);
 
 Status MakeMatMulSubgroupSource(ShaderHelper& shader,
                                 const InlinedVector<int64_t>& elements_per_thread,

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

@@ -52,21 +52,31 @@ Status ApplyMatMulIntel(ComputeContext& context,
   ORT_THROW_IF_ERROR(helper.Compute(a_shape, b_shape));
   int64_t batchA = a_shape.SizeToDimension(a_shape.NumDimensions() - 2);
   int64_t batchB = b_shape.SizeToDimension(b_shape.NumDimensions() - 2);
-
   TensorShape output_shape = helper.OutputShape();
 
   // When B is a matrix (batch is 1), we fold batchA into the M dimension for better
   // performance (e.g., [2,3,5] → [1,6,5]).
-  if (batchA != 1 && batchB == 1) {
-    // dimensions of A: [1,`batchA`, M, K]
+  // On Xe-LPG/3LPG, folding a batched matmul into a single large M loses Z-dispatch
+  // parallelism. Only fold when each per-batch M leaves a large fraction of invalid
+  // threads in its trailing workgroup (m_mod_32 ∈ [1, 24] = 25%–97% wasted), so the
+  // fold actually claws that waste back. Otherwise the Z-dispatch path wins.
+  const int64_t M = output_shape[output_shape.NumDimensions() - 2];
+  const auto& arch = context.AdapterInfo().architecture;
+  const bool is_xe_lpg_or_xe_3lpg = arch == std::string_view("xe-lpg") ||
+                                    arch == std::string_view("xe-3lpg");
+  // 32 = kSubgroupLogicalWorkGroupSizeY * ElementsPerThreadY(M > 32) on Xe-LPG/3LPG
+  const int64_t m_mod_32 = M % 32;
+  const bool xe_lpg_or_xe_3lpg_fold_ok = (m_mod_32 > 0 && m_mod_32 <= 24);
+  if (batchA != 1 && batchB == 1 && (!is_xe_lpg_or_xe_3lpg || xe_lpg_or_xe_3lpg_fold_ok)) {
+    // dimensions of A: [`batchA` * M, K]
     int64_t batchAndM = a_shape.SizeToDimension(a_shape.NumDimensions() - 1);
-    TensorShapeVector dims_a = {1, batchAndM, helper.K()};
-    // dimensions of B: [1,K,N]
-    TensorShapeVector dims_b = {1, helper.K(), helper.N()};
+    TensorShapeVector dims_a = {batchAndM, helper.K()};
+    // dimensions of B: [K, N]
+    TensorShapeVector dims_b = {helper.K(), helper.N()};
 
     a_shape = TensorShape(dims_a);
     b_shape = TensorShape(dims_b);
-    output_shape = {1, batchAndM, helper.N()};
+    output_shape = {batchAndM, helper.N()};
   }
 
   // helpful dimension variables
@@ -91,7 +101,7 @@ Status ApplyMatMulIntel(ComputeContext& context,
 
   // Always access A with 1-component when using subgroup.
   const bool is_vec4 = dim_b_outer % 4 == 0;
-  InlinedVector<int64_t> elements_per_thread = InlinedVector<int64_t>({4, intel::ElementsPerThreadY(is_vec4, dim_a_outer), 1});
+  InlinedVector<int64_t> elements_per_thread = InlinedVector<int64_t>({4, ElementsPerThreadY(context, dim_a_outer), 1});
 
   const uint32_t dispatch_x = narrow<uint32_t>((dim_b_outer + kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0] - 1) /
                                                (kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0]));

onnxruntime/test/providers/webgpu/gemm_large_test.cc

@@ -0,0 +1,158 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+
+#include "core/providers/cpu/math/gemm_helper.h"
+#include "test/providers/provider_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
+#include "default_providers.h"
+
+namespace onnxruntime {
+namespace test {
+
+static float GetBiasValue(const std::vector<float>& c_vals, const TensorShape& c_shape,
+                          int64_t M, int64_t N, int64_t m, int64_t n) {
+  if (c_vals.empty())
+    return 0.0f;
+  if (c_shape.Size() == 1)
+    return c_vals[0];
+  if (c_shape.NumDimensions() == 1 && c_shape[0] == N)
+    return c_vals[n];
+  if (c_shape.NumDimensions() == 2) {
+    if (c_shape[0] == M && c_shape[1] == 1)
+      return c_vals[m];
+    if (c_shape[0] == M && c_shape[1] == N)
+      return c_vals[m * N + n];
+    if (c_shape[0] == 1 && c_shape[1] == N)
+      return c_vals[n];
+  }
+  ADD_FAILURE() << "Unsupported bias shape in test reference";
+  return 0.0f;
+}
+
+static void ComputeExpectedResult(int64_t M, int64_t K, int64_t N,
+                                  const std::vector<float>& a_vals, const std::vector<float>& b_vals,
+                                  const std::vector<float>& c_vals, std::vector<float>& expected_vals,
+                                  int64_t a_trans, int64_t b_trans, const TensorShape& c_shape,
+                                  float alpha, float beta) {
+  for (int64_t m = 0; m < M; ++m) {
+    for (int64_t n = 0; n < N; ++n) {
+      float sum = 0.0f;
+      for (int64_t k = 0; k < K; ++k) {
+        float a_val = a_trans ? a_vals[k * M + m] : a_vals[m * K + k];
+        float b_val = b_trans ? b_vals[n * K + k] : b_vals[k * N + n];
+        sum += a_val * b_val;
+      }
+      expected_vals[m * N + n] = sum * alpha + GetBiasValue(c_vals, c_shape, M, N, m, n) * beta;
+    }
+  }
+}
+
+template <typename T, int version = 13>
+void RunTestTyped(std::initializer_list<int64_t> a_dims, int64_t a_trans, std::initializer_list<int64_t> b_dims,
+                  int64_t b_trans, std::initializer_list<int64_t> c_dims, float alpha = 1.0f, float beta = 1.0f) {
+  static_assert(std::is_same_v<T, float> || std::is_same_v<T, MLFloat16>, "unexpected type for T");
+
+  auto webgpu_ep = DefaultWebGpuExecutionProvider();
+  if (!webgpu_ep) {
+    GTEST_SKIP() << "WebGPU execution provider is not available.";
+  }
+
+  TensorShape a_shape(a_dims);
+  TensorShape b_shape(b_dims);
+  TensorShape c_shape(c_dims);
+  GemmHelper helper(a_shape, a_trans != 0, b_shape, b_trans != 0, c_shape);
+  ASSERT_STATUS_OK(helper.State());
+  const auto M = helper.M();
+  const auto K = helper.K();
+  const auto N = helper.N();
+
+  RandomValueGenerator random{1234};
+  std::vector<float> a_vals(random.Gaussian<float>(AsSpan(a_dims), 0.0f, 0.25f));
+  std::vector<float> b_vals(random.Gaussian<float>(AsSpan(b_dims), 0.0f, 0.25f));
+  std::vector<float> c_vals;
+  if (c_dims.size() != 0) {
+    c_vals = random.Gaussian<float>(AsSpan(c_dims), 0.0f, 0.25f);
+  }
+
+  std::vector<float> expected_vals(M * N);
+  ComputeExpectedResult(M, K, N, a_vals, b_vals, c_vals, expected_vals, a_trans, b_trans, c_shape, alpha, beta);
+
+  OpTester test("Gemm", version);
+  test.AddAttribute("transA", a_trans);
+  test.AddAttribute("transB", b_trans);
+  test.AddAttribute("alpha", alpha);
+  test.AddAttribute("beta", beta);
+  if constexpr (std::is_same_v<T, float>) {
+    test.AddInput<T>("A", a_dims, a_vals);
+    test.AddInput<T>("B", b_dims, b_vals);
+    if (c_dims.size() != 0)
+      test.AddInput<T>("C", c_dims, c_vals);
+    test.AddOutput<T>("Y", {M, N}, expected_vals);
+  } else {
+    test.AddInput<T>("A", a_dims, FloatsToMLFloat16s(a_vals));
+    test.AddInput<T>("B", b_dims, FloatsToMLFloat16s(b_vals));
+    if (c_dims.size() != 0)
+      test.AddInput<T>("C", c_dims, FloatsToMLFloat16s(c_vals));
+    test.AddOutput<T>("Y", {M, N}, FloatsToMLFloat16s(expected_vals));
+    test.SetOutputAbsErr("Y", 0.055f);
+    test.SetOutputRelErr("Y", 0.02f);
+  }
+
+  test.ConfigEp(std::move(webgpu_ep)).RunWithConfig();
+}
+
+template <int version = 13>
+void RunBothTypes(std::initializer_list<int64_t> a_dims, int64_t a_trans,
+                  std::initializer_list<int64_t> b_dims, int64_t b_trans,
+                  std::initializer_list<int64_t> c_dims, float alpha = 1.0f, float beta = 1.0f) {
+  RunTestTyped<float, version>(a_dims, a_trans, b_dims, b_trans, c_dims, alpha, beta);
+  RunTestTyped<MLFloat16, version>(a_dims, a_trans, b_dims, b_trans, c_dims, alpha, beta);
+}
+
+// Aligned dimensions and baseline large shapes.
+TEST(Gemm_Large, DISABLED_Aligned) {
+  RunBothTypes({512, 1024}, 0, {1024, 1024}, 0, {512, 1024});
+  RunBothTypes({127, 1024}, 0, {1024, 1024}, 0, {1024});
+}
+
+// Unaligned dimensions and edge shapes.
+TEST(Gemm_Large, DISABLED_Unaligned) {
+  RunBothTypes({127, 1023}, 0, {1023, 1023}, 0, {1023});
+  RunBothTypes({511, 1024}, 0, {1024, 1023}, 0, {511, 1});
+  RunBothTypes({6, 1024}, 0, {1024, 192}, 0, {6, 1});
+  RunBothTypes({49, 1024}, 0, {1024, 600}, 0, {49, 600});
+}
+
+// Transpose combinations and transposed bias cases.
+TEST(Gemm_Large, DISABLED_TransAB) {
+  RunBothTypes({1024, 512}, 1, {1024, 1024}, 0, {512, 1});
+  RunBothTypes({512, 1024}, 0, {1024, 1024}, 1, {512, 1});
+  RunBothTypes({1024, 512}, 1, {1024, 1024}, 1, {512, 1});
+  RunBothTypes({1024, 512}, 1, {1024, 1024}, 1, {512, 1}, 1.5f, 1.3f);
+  RunBothTypes({1024, 16}, 1, {1024, 192}, 0, {192});
+  RunBothTypes({16, 1024}, 0, {192, 1024}, 1, {192});
+  RunBothTypes({1024, 16}, 1, {192, 1024}, 1, {192});
+}
+
+// Bias broadcast coverage, including no-bias and alpha/beta scaling.
+TEST(Gemm_Large, DISABLED_BiasBroadcast) {
+  RunBothTypes({16, 1024}, 0, {1024, 191}, 0, {1, 191});
+  RunBothTypes({15, 1024}, 0, {1024, 191}, 0, {15, 191});
+  RunBothTypes({15, 1024}, 0, {1024, 192}, 0, {15, 1});
+  RunTestTyped<float>({16, 1024}, 0, {1024, 192}, 0, {});  // no bias
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {16, 1});
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {1, 192});
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {192});
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {16, 192});
+  RunBothTypes({16, 1024}, 0, {1024, 600}, 0, {1, 600});
+
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {16, 1}, 1.5f, 1.3f);
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {1, 192}, 1.5f, 1.3f);
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {192}, 1.5f, 1.3f);
+  RunBothTypes({16, 1024}, 0, {1024, 192}, 0, {16, 192}, 1.5f, 1.3f);
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/providers/webgpu/matmul_large_test.cc

@@ -0,0 +1,116 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+
+#include "core/providers/cpu/math/matmul_helper.h"
+#include "test/providers/provider_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
+#include "default_providers.h"
+
+namespace onnxruntime {
+namespace test {
+
+// Reference matmul using MatMulComputeHelper for shape/offset computation.
+// Supports arbitrary-rank batched matmul with broadcasting.
+static void ComputeExpectedResult(const std::vector<float>& a_vals, const std::vector<float>& b_vals,
+                                  std::vector<float>& out_vals,
+                                  const MatMulComputeHelper& helper) {
+  const auto M = helper.M();
+  const auto K = helper.K();
+  const auto N = helper.N();
+  const auto& left_offsets = helper.LeftOffsets();
+  const auto& right_offsets = helper.RightOffsets();
+  const auto& output_offsets = helper.OutputOffsets();
+  const size_t num_batches = output_offsets.size();
+
+  for (size_t batch = 0; batch < num_batches; ++batch) {
+    const float* a = a_vals.data() + left_offsets[batch];
+    const float* b = b_vals.data() + right_offsets[batch];
+    float* out = out_vals.data() + output_offsets[batch];
+    for (ptrdiff_t m = 0; m < M; ++m) {
+      for (ptrdiff_t n = 0; n < N; ++n) {
+        float sum = 0.0f;
+        for (ptrdiff_t k = 0; k < K; ++k) {
+          sum += a[m * K + k] * b[k * N + n];
+        }
+        out[m * N + n] = sum;
+      }
+    }
+  }
+}
+
+template <typename T, int version = 13>
+void RunTestTyped(std::initializer_list<int64_t> a_dims, std::initializer_list<int64_t> b_dims) {
+  static_assert(std::is_same_v<T, float> || std::is_same_v<T, MLFloat16>, "unexpected type for T");
+
+  auto webgpu_ep = DefaultWebGpuExecutionProvider();
+  if (!webgpu_ep) {
+    GTEST_SKIP() << "WebGPU execution provider is not available.";
+  }
+
+  TensorShape a_shape(a_dims);
+  TensorShape b_shape(b_dims);
+  MatMulComputeHelper helper;
+  ASSERT_STATUS_OK(helper.Compute(a_shape, b_shape));
+  const TensorShape& output_shape = helper.OutputShape();
+
+  RandomValueGenerator random{1234};
+  std::vector<float> a_vals(random.Gaussian<float>(AsSpan(a_dims), 0.0f, 0.25f));
+  std::vector<float> b_vals(random.Gaussian<float>(AsSpan(b_dims), 0.0f, 0.25f));
+
+  std::vector<float> expected_vals(output_shape.Size());
+  ComputeExpectedResult(a_vals, b_vals, expected_vals, helper);
+
+  std::vector<int64_t> output_dims(output_shape.NumDimensions());
+  output_shape.CopyDims(output_dims.data(), output_shape.NumDimensions());
+
+  OpTester test("MatMul", version);
+  if constexpr (std::is_same_v<T, float>) {
+    test.AddInput<T>("A", a_dims, a_vals);
+    test.AddInput<T>("B", b_dims, b_vals);
+    test.AddOutput<T>("Y", output_dims, expected_vals);
+  } else {
+    test.AddInput<T>("A", a_dims, FloatsToMLFloat16s(a_vals));
+    test.AddInput<T>("B", b_dims, FloatsToMLFloat16s(b_vals));
+    test.AddOutput<T>("Y", output_dims, FloatsToMLFloat16s(expected_vals));
+    test.SetOutputAbsErr("Y", 0.055f);
+    test.SetOutputRelErr("Y", 0.02f);
+  }
+
+  test.ConfigEp(std::move(webgpu_ep)).RunWithConfig();
+}
+
+template <int version = 13>
+void RunBothTypes(std::initializer_list<int64_t> a_dims, std::initializer_list<int64_t> b_dims) {
+  RunTestTyped<float, version>(a_dims, b_dims);
+  RunTestTyped<MLFloat16, version>(a_dims, b_dims);
+}
+
+// 2D aligned baseline shapes.
+TEST(MatMul_Large, DISABLED_Aligned) {
+  RunBothTypes({128, 64}, {64, 1024});
+}
+
+// 2D unaligned edge shapes.
+TEST(MatMul_Large, DISABLED_Unaligned) {
+  RunBothTypes({127, 64}, {64, 1024});
+  RunBothTypes({127, 63}, {63, 1023});
+}
+
+// 3D broadcast and non-broadcast cases.
+TEST(MatMul_Large, DISABLED_Broadcast3D) {
+  RunBothTypes({2, 128, 64}, {64, 1024});
+  RunBothTypes({2, 128, 64}, {2, 64, 1024});
+  RunBothTypes({2, 128, 64}, {64, 1023});
+  RunBothTypes({2, 128, 64}, {2, 64, 1023});
+}
+
+// 4D broadcast cases.
+TEST(MatMul_Large, DISABLED_Broadcast4D) {
+  RunBothTypes({2, 2, 128, 64}, {2, 64, 1024});
+  RunBothTypes({2, 2, 128, 64}, {2, 64, 1023});
+}
+
+}  // namespace test
+}  // namespace onnxruntime