From a2a7d5f249bd40709fc5a1339b291cb194b4d86b Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Fri, 14 Mar 2025 14:32:36 +0800 Subject: [PATCH 1/8] implement --- .../core/providers/webgpu/math/gemm.cc | 144 ++++++++++++++++++ onnxruntime/core/providers/webgpu/math/gemm.h | 66 ++++++++ .../webgpu/webgpu_execution_provider.cc | 8 +- .../test/providers/cpu/math/gemm_test.cc | 6 +- 4 files changed, 218 insertions(+), 6 deletions(-) create mode 100644 onnxruntime/core/providers/webgpu/math/gemm.cc create mode 100644 onnxruntime/core/providers/webgpu/math/gemm.h diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc new file mode 100644 index 0000000000000..5a29d568dcaf1 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -0,0 +1,144 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/math/gemm.h" + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" + +namespace onnxruntime { +namespace webgpu { + +#define WEBGPU_GEMM_VERSIONED_KERNEL(start, end) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + Gemm, \ + kOnnxDomain, \ + start, \ + end, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", WebGpuSupportedNumberTypes()), \ + Gemm); + +#define WEBGPU_GEMM_KERNEL(version) \ + ONNX_OPERATOR_KERNEL_EX( \ + Gemm, \ + kOnnxDomain, \ + version, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", WebGpuSupportedNumberTypes()), \ + Gemm); + +WEBGPU_GEMM_VERSIONED_KERNEL(7, 8) +WEBGPU_GEMM_VERSIONED_KERNEL(9, 10) +WEBGPU_GEMM_VERSIONED_KERNEL(11, 12) +WEBGPU_GEMM_KERNEL(13) + +Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { + const ShaderVariableHelper& A = shader.AddInput("A", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + const ShaderVariableHelper& B = shader.AddInput("B", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + + const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") + << " let m = global_idx / uniforms.N;\n" + << " let n = global_idx % uniforms.N;\n" + << " var value = A_value_t(0);\n" + << "\n" + << " for (var k = 0u; k < uniforms.K; k = k + 1u) {\n"; + + if (transA_ && transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") + << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + } else if (transA_ && !transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") + << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + } else if (!transA_ && transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") + << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + } else { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") + << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + } + shader.MainFunctionBody() << " }\n" + << "\n"; + // Calculate Alpha + if (alpha_) { + shader.MainFunctionBody() << " value = value * A_value_t(uniforms.alpha);\n"; + } + + // Calculate Bias + if (need_handle_bias_) { + const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + shader.MainFunctionBody() << " value = value + A_value_t(uniforms.beta) * " + << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"; + } + + shader.MainFunctionBody() << output.SetByOffset("global_idx", "value") << "\n"; + + return Status::OK(); +} + +Status Gemm::ComputeInternal(ComputeContext& context) const { + const auto* A = context.Input(0); + const auto* B = context.Input(1); + const auto* C = context.Input(2); + + if (A == nullptr || B == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Gemm requires input tensors A and B."); + } + + const auto& A_shape = A->Shape(); + const auto& B_shape = B->Shape(); + + if (A_shape.NumDimensions() != 2 || B_shape.NumDimensions() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input tensors A and B must be 2 dimensional."); + } + + int64_t M = transA_ ? A_shape[1] : A_shape[0]; + int64_t K = transA_ ? A_shape[0] : A_shape[1]; + int64_t N = transB_ ? B_shape[0] : B_shape[1]; + + if ((transA_ ? A_shape[0] : A_shape[1]) != (transB_ ? B_shape[1] : B_shape[0])) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inner dimensions of A and B must match."); + } + + std::vector 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(); + } + + constexpr size_t TILE_SIZE = 16; + int64_t num_tiles_m = (M + TILE_SIZE - 1) / TILE_SIZE; + int64_t num_tiles_n = (N + TILE_SIZE - 1) / TILE_SIZE; + int64_t dispatch_size = num_tiles_m * num_tiles_n; + + GemmProgram program{transA_, transB_, alpha_, beta_, C && beta_}; + program.AddInputs({{A, ProgramTensorMetadataDependency::Type}, + {B, ProgramTensorMetadataDependency::Type}}); + + if (C && beta_) { + program.AddInput({C, ProgramTensorMetadataDependency::Rank}); + } + + program.AddOutputs({Y}) + .SetDispatchGroupSize(dispatch_size) + .SetWorkgroupSize(TILE_SIZE, TILE_SIZE, 1) + .AddUniformVariables({ + {static_cast(output_size)}, // output_size + {static_cast(M)}, // M + {static_cast(N)}, // N + {static_cast(K)}, // K + {alpha_}, // alpha + {beta_} // beta + }); + + return context.RunProgram(program); +} + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/gemm.h b/onnxruntime/core/providers/webgpu/math/gemm.h new file mode 100644 index 0000000000000..2f20da4c2a7e1 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/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 { + +class GemmProgram final : public Program { + public: + GemmProgram(bool transA, bool transB, float alpha, float beta, bool need_handle_bias) + : Program{"Gemm"}, + transA_{transA}, + transB_{transB}, + alpha_{alpha}, + beta_{beta}, + need_handle_bias_{need_handle_bias} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"output_size", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"alpha", ProgramUniformVariableDataType::Float32}, + {"beta", ProgramUniformVariableDataType::Float32}); + + private: + bool transA_; + bool transB_; + float alpha_; + float beta_; + bool need_handle_bias_; +}; + +class Gemm final : public WebGpuKernel { + public: + Gemm(const OpKernelInfo& info) : WebGpuKernel(info) { + int64_t transA_temp; + info.GetAttrOrDefault("transA", &transA_temp, static_cast(0)); + transA_ = transA_temp != 0; + + int64_t transB_temp; + info.GetAttrOrDefault("transB", &transB_temp, static_cast(0)); + transB_ = transB_temp != 0; + + info.GetAttrOrDefault("alpha", &alpha_, 1.0f); + info.GetAttrOrDefault("beta", &beta_, 1.0f); + } + + Status ComputeInternal(ComputeContext& context) const override; + + private: + bool transA_; + bool transB_; + float alpha_; + float beta_; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index d44cf4674d8a3..81cc13c866fed 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -610,10 +610,10 @@ std::unique_ptr RegisterKernels() { // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index d0069a0069646..1ddb6cdde4ba0 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -761,9 +761,10 @@ TYPED_TEST(GemmOpTypedTests, ZeroKWithBias) { test.AddInput("C", {4}, std::vector(4, static_cast(1.0f))); test.AddOutput("Y", {4, 4}, std::vector(16, static_cast(1.0f))); + // WebGPU EP doesn't support zero buffer. test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider, kDmlExecutionProvider, kDnnlExecutionProvider, kQnnExecutionProvider, - kOpenVINOExecutionProvider}) + kOpenVINOExecutionProvider, kWebGpuExecutionProvider}) .Config(run_with_tunable_op) .RunWithConfig(); } @@ -780,9 +781,10 @@ TYPED_TEST(GemmOpTypedTests, ZeroKWithNoBias) { test.AddInput("B", {0, 4}, {}); test.AddOutput("Y", {4, 4}, std::vector(16, static_cast(0.0f))); + // WebGPU EP doesn't support zero buffer. test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider, kDmlExecutionProvider, kDnnlExecutionProvider, kQnnExecutionProvider, - kOpenVINOExecutionProvider}) + kOpenVINOExecutionProvider, kWebGpuExecutionProvider}) .Config(run_with_tunable_op) .RunWithConfig(); } From b1a43f3eece1fbbf4b7429598ce9184db798ffa3 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Mon, 17 Mar 2025 17:53:19 +0800 Subject: [PATCH 2/8] handle A/B --- .../core/providers/webgpu/math/gemm.cc | 81 ++++++++++--------- onnxruntime/core/providers/webgpu/math/gemm.h | 6 +- .../test/providers/cpu/math/gemm_test.cc | 6 +- 3 files changed, 51 insertions(+), 42 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index 5a29d568dcaf1..55bac433a7fa7 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -36,42 +36,47 @@ WEBGPU_GEMM_VERSIONED_KERNEL(11, 12) WEBGPU_GEMM_KERNEL(13) Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { - const ShaderVariableHelper& A = shader.AddInput("A", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - const ShaderVariableHelper& B = shader.AddInput("B", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - - const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") << " let m = global_idx / uniforms.N;\n" << " let n = global_idx % uniforms.N;\n" - << " var value = A_value_t(0);\n" - << "\n" - << " for (var k = 0u; k < uniforms.K; k = k + 1u) {\n"; - - if (transA_ && transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") - << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; - } else if (transA_ && !transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") - << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; - } else if (!transA_ && transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") - << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; - } else { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") - << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; - } - shader.MainFunctionBody() << " }\n" + << " var value = output_value_t(0);\n" << "\n"; + + // When K == 0, we don't bind A and B. Because WebGPU doesn't support binding a zero-sized buffer, + if (need_handle_matmul_) { + const ShaderVariableHelper& A = shader.AddInput("A", ShaderUsage::UseUniform); + const ShaderVariableHelper& B = shader.AddInput("B", ShaderUsage::UseUniform); + + shader.MainFunctionBody() << " for (var k = 0u; k < uniforms.K; k = k + 1u) {\n"; + + if (transA_ && transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") + << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + } else if (transA_ && !transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") + << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + } else if (!transA_ && transB_) { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") + << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + } else { + shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") + << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + } + shader.MainFunctionBody() << " }\n" + << "\n"; + } + // Calculate Alpha if (alpha_) { - shader.MainFunctionBody() << " value = value * A_value_t(uniforms.alpha);\n"; + shader.MainFunctionBody() << " value = value * output_value_t(uniforms.alpha);\n"; } // Calculate Bias if (need_handle_bias_) { - const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); - shader.MainFunctionBody() << " value = value + A_value_t(uniforms.beta) * " + const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform); + shader.MainFunctionBody() << " value = value + output_value_t(uniforms.beta) * " << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"; } @@ -112,22 +117,26 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { return Status::OK(); } - constexpr size_t TILE_SIZE = 16; - int64_t num_tiles_m = (M + TILE_SIZE - 1) / TILE_SIZE; - int64_t num_tiles_n = (N + TILE_SIZE - 1) / TILE_SIZE; - int64_t dispatch_size = num_tiles_m * num_tiles_n; + // WebGPU doesn't support binding a zero-sized buffer, so we need to check if K == 0 + bool need_handle_matmul = K != 0; + bool need_handle_bias = C && beta_; + + GemmProgram program{transA_, transB_, alpha_, beta_, need_handle_bias, need_handle_matmul}; - GemmProgram program{transA_, transB_, alpha_, beta_, C && beta_}; - program.AddInputs({{A, ProgramTensorMetadataDependency::Type}, - {B, ProgramTensorMetadataDependency::Type}}); + if (need_handle_matmul) { + program.AddInputs({{A, ProgramTensorMetadataDependency::Type}, + {B, ProgramTensorMetadataDependency::Type}}); + } - if (C && beta_) { + if (need_handle_bias) { program.AddInput({C, ProgramTensorMetadataDependency::Rank}); } - program.AddOutputs({Y}) - .SetDispatchGroupSize(dispatch_size) - .SetWorkgroupSize(TILE_SIZE, TILE_SIZE, 1) + constexpr size_t WORKGROUP_SIZE = 64; + + program.AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) + .SetDispatchGroupSize(output_size + WORKGROUP_SIZE - 1 / WORKGROUP_SIZE) + .SetWorkgroupSize(WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(output_size)}, // output_size {static_cast(M)}, // M diff --git a/onnxruntime/core/providers/webgpu/math/gemm.h b/onnxruntime/core/providers/webgpu/math/gemm.h index 2f20da4c2a7e1..1e888495b018c 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.h +++ b/onnxruntime/core/providers/webgpu/math/gemm.h @@ -12,13 +12,14 @@ namespace webgpu { class GemmProgram final : public Program { public: - GemmProgram(bool transA, bool transB, float alpha, float beta, bool need_handle_bias) + GemmProgram(bool transA, bool transB, float alpha, float beta, bool need_handle_bias, bool need_handle_matmul) : Program{"Gemm"}, transA_{transA}, transB_{transB}, alpha_{alpha}, beta_{beta}, - need_handle_bias_{need_handle_bias} {} + need_handle_bias_{need_handle_bias}, + need_handle_matmul_{need_handle_matmul} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -36,6 +37,7 @@ class GemmProgram final : public Program { float alpha_; float beta_; bool need_handle_bias_; + bool need_handle_matmul_; }; class Gemm final : public WebGpuKernel { diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index 1ddb6cdde4ba0..d0069a0069646 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -761,10 +761,9 @@ TYPED_TEST(GemmOpTypedTests, ZeroKWithBias) { test.AddInput("C", {4}, std::vector(4, static_cast(1.0f))); test.AddOutput("Y", {4, 4}, std::vector(16, static_cast(1.0f))); - // WebGPU EP doesn't support zero buffer. test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider, kDmlExecutionProvider, kDnnlExecutionProvider, kQnnExecutionProvider, - kOpenVINOExecutionProvider, kWebGpuExecutionProvider}) + kOpenVINOExecutionProvider}) .Config(run_with_tunable_op) .RunWithConfig(); } @@ -781,10 +780,9 @@ TYPED_TEST(GemmOpTypedTests, ZeroKWithNoBias) { test.AddInput("B", {0, 4}, {}); test.AddOutput("Y", {4, 4}, std::vector(16, static_cast(0.0f))); - // WebGPU EP doesn't support zero buffer. test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider, kDmlExecutionProvider, kDnnlExecutionProvider, kQnnExecutionProvider, - kOpenVINOExecutionProvider, kWebGpuExecutionProvider}) + kOpenVINOExecutionProvider}) .Config(run_with_tunable_op) .RunWithConfig(); } From 2198ff928393f7b3c7ab680972e6ad0d04000dbe Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Tue, 18 Mar 2025 10:04:46 +0800 Subject: [PATCH 3/8] fix pipeline error --- onnxruntime/core/providers/webgpu/math/gemm.cc | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index 55bac433a7fa7..9582f97fbd03d 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -132,10 +132,8 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { program.AddInput({C, ProgramTensorMetadataDependency::Rank}); } - constexpr size_t WORKGROUP_SIZE = 64; - program.AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) - .SetDispatchGroupSize(output_size + WORKGROUP_SIZE - 1 / WORKGROUP_SIZE) + .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .SetWorkgroupSize(WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(output_size)}, // output_size From 9438a2b934541fe2d4ff4f792aaccede789de1a5 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Tue, 18 Mar 2025 13:48:47 +0800 Subject: [PATCH 4/8] remove unneccessary variable --- onnxruntime/core/providers/webgpu/math/gemm.cc | 2 +- onnxruntime/core/providers/webgpu/math/gemm.h | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index 9582f97fbd03d..d2f8aebc00561 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -121,7 +121,7 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { bool need_handle_matmul = K != 0; bool need_handle_bias = C && beta_; - GemmProgram program{transA_, transB_, alpha_, beta_, need_handle_bias, need_handle_matmul}; + GemmProgram program{transA_, transB_, alpha_, need_handle_bias, need_handle_matmul}; if (need_handle_matmul) { program.AddInputs({{A, ProgramTensorMetadataDependency::Type}, diff --git a/onnxruntime/core/providers/webgpu/math/gemm.h b/onnxruntime/core/providers/webgpu/math/gemm.h index 1e888495b018c..75d63d3eba5ec 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.h +++ b/onnxruntime/core/providers/webgpu/math/gemm.h @@ -12,12 +12,11 @@ namespace webgpu { class GemmProgram final : public Program { public: - GemmProgram(bool transA, bool transB, float alpha, float beta, bool need_handle_bias, bool need_handle_matmul) + GemmProgram(bool transA, bool transB, float alpha, bool need_handle_bias, bool need_handle_matmul) : Program{"Gemm"}, transA_{transA}, transB_{transB}, alpha_{alpha}, - beta_{beta}, need_handle_bias_{need_handle_bias}, need_handle_matmul_{need_handle_matmul} {} @@ -35,7 +34,6 @@ class GemmProgram final : public Program { bool transA_; bool transB_; float alpha_; - float beta_; bool need_handle_bias_; bool need_handle_matmul_; }; From 4f3159253dc3c8a26a9329fe36fce3a20a267752 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Tue, 18 Mar 2025 15:35:57 +0800 Subject: [PATCH 5/8] optimize --- .../core/providers/webgpu/math/gemm.cc | 140 ++++++++++++++---- onnxruntime/core/providers/webgpu/math/gemm.h | 2 +- 2 files changed, 112 insertions(+), 30 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index d2f8aebc00561..2b725eb4ac410 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -36,36 +36,108 @@ WEBGPU_GEMM_VERSIONED_KERNEL(11, 12) WEBGPU_GEMM_KERNEL(13) Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { + const uint32_t TILE_SIZE = 16; + + // Add shared memory arrays + shader.AdditionalImplementation() << "var tile_a: array, " << TILE_SIZE << ">;\n" + << "var tile_b: array, " << TILE_SIZE << ">;\n\n"; + const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") - << " let m = global_idx / uniforms.N;\n" - << " let n = global_idx % uniforms.N;\n" - << " var value = output_value_t(0);\n" - << "\n"; + shader.MainFunctionBody() << " var value = output_value_t(0);\n\n" + << " let tile_col_start = (workgroup_id.x % uniforms.num_tile_n) * " << TILE_SIZE << "u;\n" + << " let tile_row_start = (workgroup_id.x / uniforms.num_tile_n) * " << TILE_SIZE << "u;\n"; - // When K == 0, we don't bind A and B. Because WebGPU doesn't support binding a zero-sized buffer, + // When K == 0, we don't bind A and B. Because WebGPU doesn't support binding a zero-sized buffer. if (need_handle_matmul_) { const ShaderVariableHelper& A = shader.AddInput("A", ShaderUsage::UseUniform); const ShaderVariableHelper& B = shader.AddInput("B", ShaderUsage::UseUniform); - shader.MainFunctionBody() << " for (var k = 0u; k < uniforms.K; k = k + 1u) {\n"; + shader.MainFunctionBody() + << " let num_tiles = (uniforms.K - 1u) / " << TILE_SIZE << "u + 1u;\n" + << " var k_start = 0u;\n" + << " for (var t = 0u; t < num_tiles; t = t + 1u) {\n"; + + // Fill workgroup shared memory + if (transA_ && transB_) { + shader.MainFunctionBody() << " var col = tile_row_start + local_id.x;\n" + << " var row = k_start + local_id.y;\n" + << " if (col < uniforms.M && row < uniforms.K) {\n" + << " tile_a[local_id.y][local_id.x] = " << A.GetByOffset("row * uniforms.M + col") << ";\n" + << " } else {\n" + << " tile_a[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n\n" + << " col = k_start + local_id.x;\n" + << " row = tile_col_start + local_id.y;\n" + << " if (col < uniforms.K && row < uniforms.N) {\n" + << " tile_b[local_id.y][local_id.x] = " << B.GetByOffset("row * uniforms.K + col") << ";\n" + << " } else {\n" + << " tile_b[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n"; + } else if (transA_ && !transB_) { + shader.MainFunctionBody() << " var col = tile_row_start + local_id.x;\n" + << " var row = k_start + local_id.y;\n" + << " if (col < uniforms.M && row < uniforms.K) {\n" + << " tile_a[local_id.y][local_id.x] = " << A.GetByOffset("row * uniforms.M + col") << ";\n" + << " } else {\n" + << " tile_a[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n\n" + << " col = tile_col_start + local_id.x;\n" + << " row = k_start + local_id.y;\n" + << " if (col < uniforms.N && row < uniforms.K) {\n" + << " tile_b[local_id.y][local_id.x] = " << B.GetByOffset("row * uniforms.N + col") << ";\n" + << " } else {\n" + << " tile_b[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n"; + } else if (!transA_ && transB_) { + shader.MainFunctionBody() << " var col = k_start + local_id.x;\n" + << " var row = tile_row_start + local_id.y;\n" + << " if (col < uniforms.K && row < uniforms.M) {\n" + << " tile_a[local_id.y][local_id.x] = " << A.GetByOffset("row * uniforms.K + col") << ";\n" + << " } else {\n" + << " tile_a[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n\n" + << " col = k_start + local_id.x;\n" + << " row = tile_col_start + local_id.y;\n" + << " if (col < uniforms.K && row < uniforms.N) {\n" + << " tile_b[local_id.y][local_id.x] = " << B.GetByOffset("row * uniforms.K + col") << ";\n" + << " } else {\n" + << " tile_b[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n"; + } else { + shader.MainFunctionBody() << " var col = k_start + local_id.x;\n" + << " var row = tile_row_start + local_id.y;\n" + << " if (col < uniforms.K && row < uniforms.M) {\n" + << " tile_a[local_id.y][local_id.x] = " << A.GetByOffset("row * uniforms.K + col") << ";\n" + << " } else {\n" + << " tile_a[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n\n" + << " col = tile_col_start + local_id.x;\n" + << " row = k_start + local_id.y;\n" + << " if (col < uniforms.N && row < uniforms.K) {\n" + << " tile_b[local_id.y][local_id.x] = " << B.GetByOffset("row * uniforms.N + col") << ";\n" + << " } else {\n" + << " tile_b[local_id.y][local_id.x] = output_value_t(0);\n" + << " }\n"; + } + + shader.MainFunctionBody() << " k_start = k_start + " << TILE_SIZE << "u;\n" + << " workgroupBarrier();\n\n" + << " for (var k = 0u; k < " << TILE_SIZE << "u; k = k + 1u) {\n"; if (transA_ && transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") - << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + shader.MainFunctionBody() << " value = value + tile_a[k][local_id.y] * tile_b[local_id.x][k];\n"; } else if (transA_ && !transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("k * uniforms.M + m") - << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + shader.MainFunctionBody() << " value = value + tile_a[k][local_id.y] * tile_b[k][local_id.x];\n"; } else if (!transA_ && transB_) { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") - << " * " << B.GetByOffset("n * uniforms.K + k") << ";\n"; + shader.MainFunctionBody() << " value = value + tile_a[local_id.y][k] * tile_b[local_id.x][k];\n"; } else { - shader.MainFunctionBody() << " value = value + " << A.GetByOffset("m * uniforms.K + k") - << " * " << B.GetByOffset("k * uniforms.N + n") << ";\n"; + shader.MainFunctionBody() << " value = value + tile_a[local_id.y][k] * tile_b[k][local_id.x];\n"; } - shader.MainFunctionBody() << " }\n" - << "\n"; + + shader.MainFunctionBody() << " }\n" + << " workgroupBarrier();\n" + << " }\n\n"; } // Calculate Alpha @@ -73,6 +145,9 @@ Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.MainFunctionBody() << " value = value * output_value_t(uniforms.alpha);\n"; } + shader.MainFunctionBody() << " let m = tile_row_start + local_id.y;\n" + << " let n = tile_col_start + local_id.x;\n"; + // Calculate Bias if (need_handle_bias_) { const ShaderVariableHelper& C = shader.AddInput("C", ShaderUsage::UseUniform); @@ -80,7 +155,10 @@ Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { << C.GetByOffset(C.BroadcastedIndicesToOffset("vec2(m, n)", output)) << ";\n"; } - shader.MainFunctionBody() << output.SetByOffset("global_idx", "value") << "\n"; + // Write output + shader.MainFunctionBody() << " if (m < uniforms.M && n < uniforms.N) {\n" + << " " << output.SetByOffset("m * uniforms.N + n", "value") << "\n" + << " }\n"; return Status::OK(); } @@ -101,9 +179,9 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input tensors A and B must be 2 dimensional."); } - int64_t M = transA_ ? A_shape[1] : A_shape[0]; - int64_t K = transA_ ? A_shape[0] : A_shape[1]; - int64_t N = transB_ ? B_shape[0] : B_shape[1]; + uint32_t M = gsl::narrow(transA_ ? A_shape[1] : A_shape[0]); + uint32_t K = gsl::narrow(transA_ ? A_shape[0] : A_shape[1]); + uint32_t N = gsl::narrow(transB_ ? B_shape[0] : B_shape[1]); if ((transA_ ? A_shape[0] : A_shape[1]) != (transB_ ? B_shape[1] : B_shape[0])) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inner dimensions of A and B must match."); @@ -132,16 +210,20 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { program.AddInput({C, ProgramTensorMetadataDependency::Rank}); } + const uint32_t TILE_SIZE = 16; + 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.AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) - .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .SetWorkgroupSize(WORKGROUP_SIZE) + .SetDispatchGroupSize(num_tile_n * num_tile_m) + .SetWorkgroupSize(TILE_SIZE, TILE_SIZE) .AddUniformVariables({ - {static_cast(output_size)}, // output_size - {static_cast(M)}, // M - {static_cast(N)}, // N - {static_cast(K)}, // K - {alpha_}, // alpha - {beta_} // beta + {static_cast(num_tile_n)}, // num_tile_n + {static_cast(M)}, // M + {static_cast(N)}, // N + {static_cast(K)}, // K + {alpha_}, // alpha + {beta_} // beta }); return context.RunProgram(program); diff --git a/onnxruntime/core/providers/webgpu/math/gemm.h b/onnxruntime/core/providers/webgpu/math/gemm.h index 75d63d3eba5ec..7fee1091de5f8 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.h +++ b/onnxruntime/core/providers/webgpu/math/gemm.h @@ -23,7 +23,7 @@ class GemmProgram final : public Program { Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"output_size", ProgramUniformVariableDataType::Uint32}, + {"num_tile_n", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, From bc54156f2b56827633b48194ca12306a270452ef Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Wed, 19 Mar 2025 15:29:45 +0800 Subject: [PATCH 6/8] resolve comments --- onnxruntime/core/providers/webgpu/math/gemm.cc | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index 2b725eb4ac410..c39317bd0ac55 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -3,6 +3,8 @@ #include "core/providers/webgpu/math/gemm.h" +#include + #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" @@ -45,10 +47,10 @@ Status GemmProgram::GenerateShaderCode(ShaderHelper& shader) const { const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); shader.MainFunctionBody() << " var value = output_value_t(0);\n\n" - << " let tile_col_start = (workgroup_id.x % uniforms.num_tile_n) * " << TILE_SIZE << "u;\n" - << " let tile_row_start = (workgroup_id.x / uniforms.num_tile_n) * " << TILE_SIZE << "u;\n"; + << " let tile_col_start = (workgroup_idx % uniforms.num_tile_n) * " << TILE_SIZE << "u;\n" + << " let tile_row_start = (workgroup_idx / uniforms.num_tile_n) * " << TILE_SIZE << "u;\n"; - // When K == 0, we don't bind A and B. Because WebGPU doesn't support binding a zero-sized buffer. + // When A or B is empty, we don't bind A and B. Because WebGPU doesn't support binding a zero-sized buffer. if (need_handle_matmul_) { const ShaderVariableHelper& A = shader.AddInput("A", ShaderUsage::UseUniform); const ShaderVariableHelper& B = shader.AddInput("B", ShaderUsage::UseUniform); @@ -179,9 +181,9 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input tensors A and B must be 2 dimensional."); } - uint32_t M = gsl::narrow(transA_ ? A_shape[1] : A_shape[0]); - uint32_t K = gsl::narrow(transA_ ? A_shape[0] : A_shape[1]); - uint32_t N = gsl::narrow(transB_ ? B_shape[0] : B_shape[1]); + uint32_t M = onnxruntime::narrow(transA_ ? A_shape[1] : A_shape[0]); + uint32_t K = onnxruntime::narrow(transA_ ? A_shape[0] : A_shape[1]); + uint32_t N = onnxruntime::narrow(transB_ ? B_shape[0] : B_shape[1]); if ((transA_ ? A_shape[0] : A_shape[1]) != (transB_ ? B_shape[1] : B_shape[0])) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inner dimensions of A and B must match."); @@ -195,8 +197,8 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { return Status::OK(); } - // WebGPU doesn't support binding a zero-sized buffer, so we need to check if K == 0 - bool need_handle_matmul = K != 0; + // 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_; GemmProgram program{transA_, transB_, alpha_, need_handle_bias, need_handle_matmul}; From 996f5143cd128832f0a5a5b9993fddc74c19d9e7 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Thu, 20 Mar 2025 21:10:43 +0800 Subject: [PATCH 7/8] use cachehint --- onnxruntime/core/providers/webgpu/math/gemm.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index c39317bd0ac55..575d22616b7b2 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -4,6 +4,7 @@ #include "core/providers/webgpu/math/gemm.h" #include +#include #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" @@ -216,7 +217,8 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { 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.AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) + program.CacheHint(absl::StrJoin(std::make_tuple(alpha_, transA_, transB_), ",")) + .AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) .SetDispatchGroupSize(num_tile_n * num_tile_m) .SetWorkgroupSize(TILE_SIZE, TILE_SIZE) .AddUniformVariables({ From ff864c6f1dbb42715db44b3acc7f615c0c93c038 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Fri, 21 Mar 2025 16:49:54 +0800 Subject: [PATCH 8/8] remove strjoin --- onnxruntime/core/providers/webgpu/math/gemm.cc | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/gemm.cc b/onnxruntime/core/providers/webgpu/math/gemm.cc index 575d22616b7b2..4057b63f0c65d 100644 --- a/onnxruntime/core/providers/webgpu/math/gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/gemm.cc @@ -4,7 +4,6 @@ #include "core/providers/webgpu/math/gemm.h" #include -#include #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" @@ -217,7 +216,7 @@ Status Gemm::ComputeInternal(ComputeContext& context) const { 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(absl::StrJoin(std::make_tuple(alpha_, transA_, transB_), ",")) + program.CacheHint(alpha_, transA_, transB_) .AddOutputs({{Y, ProgramTensorMetadataDependency::Type}}) .SetDispatchGroupSize(num_tile_n * num_tile_m) .SetWorkgroupSize(TILE_SIZE, TILE_SIZE)