diff --git a/onnxruntime/core/providers/webgpu/math/matmul.cc b/onnxruntime/core/providers/webgpu/math/matmul.cc index 6b0600fa03d1d..8e769186e56fc 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.cc +++ b/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 diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm.cc index 41376b07ce0d8..36d0216bc5fec 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm.cc +++ b/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 elements_per_thread = InlinedVector({4, intel::ElementsPerThreadY(is_vec4, M), 1}); + InlinedVector elements_per_thread = InlinedVector({4, intel::ElementsPerThreadY(context, M), 1}); const uint32_t dispatch_x = narrow((N + kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0] - 1) / (kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0])); const uint32_t dispatch_y = narrow((M + kSubgroupLogicalWorkGroupSizeY * elements_per_thread[1] - 1) / diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc index a6baf8dfb0239..43ebccddf60e2 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc +++ b/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, diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h index 89dca023d3e1b..884baa3ddf87d 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h +++ b/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& elements_per_thread, diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc index b6ec2e0c2b10b..dd94a6efb35b1 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc +++ b/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 elements_per_thread = InlinedVector({4, intel::ElementsPerThreadY(is_vec4, dim_a_outer), 1}); + InlinedVector elements_per_thread = InlinedVector({4, ElementsPerThreadY(context, dim_a_outer), 1}); const uint32_t dispatch_x = narrow((dim_b_outer + kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0] - 1) / (kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0])); diff --git a/onnxruntime/test/providers/webgpu/gemm_large_test.cc b/onnxruntime/test/providers/webgpu/gemm_large_test.cc new file mode 100644 index 0000000000000..520f55cc3707f --- /dev/null +++ b/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& 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& a_vals, const std::vector& b_vals, + const std::vector& c_vals, std::vector& 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 +void RunTestTyped(std::initializer_list a_dims, int64_t a_trans, std::initializer_list b_dims, + int64_t b_trans, std::initializer_list c_dims, float alpha = 1.0f, float beta = 1.0f) { + static_assert(std::is_same_v || std::is_same_v, "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 a_vals(random.Gaussian(AsSpan(a_dims), 0.0f, 0.25f)); + std::vector b_vals(random.Gaussian(AsSpan(b_dims), 0.0f, 0.25f)); + std::vector c_vals; + if (c_dims.size() != 0) { + c_vals = random.Gaussian(AsSpan(c_dims), 0.0f, 0.25f); + } + + std::vector 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) { + test.AddInput("A", a_dims, a_vals); + test.AddInput("B", b_dims, b_vals); + if (c_dims.size() != 0) + test.AddInput("C", c_dims, c_vals); + test.AddOutput("Y", {M, N}, expected_vals); + } else { + test.AddInput("A", a_dims, FloatsToMLFloat16s(a_vals)); + test.AddInput("B", b_dims, FloatsToMLFloat16s(b_vals)); + if (c_dims.size() != 0) + test.AddInput("C", c_dims, FloatsToMLFloat16s(c_vals)); + test.AddOutput("Y", {M, N}, FloatsToMLFloat16s(expected_vals)); + test.SetOutputAbsErr("Y", 0.055f); + test.SetOutputRelErr("Y", 0.02f); + } + + test.ConfigEp(std::move(webgpu_ep)).RunWithConfig(); +} + +template +void RunBothTypes(std::initializer_list a_dims, int64_t a_trans, + std::initializer_list b_dims, int64_t b_trans, + std::initializer_list c_dims, float alpha = 1.0f, float beta = 1.0f) { + RunTestTyped(a_dims, a_trans, b_dims, b_trans, c_dims, alpha, beta); + RunTestTyped(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({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 diff --git a/onnxruntime/test/providers/webgpu/matmul_large_test.cc b/onnxruntime/test/providers/webgpu/matmul_large_test.cc new file mode 100644 index 0000000000000..a1756a31824b3 --- /dev/null +++ b/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& a_vals, const std::vector& b_vals, + std::vector& 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 +void RunTestTyped(std::initializer_list a_dims, std::initializer_list b_dims) { + static_assert(std::is_same_v || std::is_same_v, "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 a_vals(random.Gaussian(AsSpan(a_dims), 0.0f, 0.25f)); + std::vector b_vals(random.Gaussian(AsSpan(b_dims), 0.0f, 0.25f)); + + std::vector expected_vals(output_shape.Size()); + ComputeExpectedResult(a_vals, b_vals, expected_vals, helper); + + std::vector output_dims(output_shape.NumDimensions()); + output_shape.CopyDims(output_dims.data(), output_shape.NumDimensions()); + + OpTester test("MatMul", version); + if constexpr (std::is_same_v) { + test.AddInput("A", a_dims, a_vals); + test.AddInput("B", b_dims, b_vals); + test.AddOutput("Y", output_dims, expected_vals); + } else { + test.AddInput("A", a_dims, FloatsToMLFloat16s(a_vals)); + test.AddInput("B", b_dims, FloatsToMLFloat16s(b_vals)); + test.AddOutput("Y", output_dims, FloatsToMLFloat16s(expected_vals)); + test.SetOutputAbsErr("Y", 0.055f); + test.SetOutputRelErr("Y", 0.02f); + } + + test.ConfigEp(std::move(webgpu_ep)).RunWithConfig(); +} + +template +void RunBothTypes(std::initializer_list a_dims, std::initializer_list b_dims) { + RunTestTyped(a_dims, b_dims); + RunTestTyped(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