From 8241759aa16dfdcf94a2efa52a63e50adfcf0682 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 10 Mar 2025 15:40:17 -0700 Subject: [PATCH 1/5] gelu contrib ops impl --- .../webgpu/bert/bias_split_gelu.cc | 80 +++++++++++++++++++ .../contrib_ops/webgpu/bert/bias_split_gelu.h | 32 ++++++++ onnxruntime/contrib_ops/webgpu/bert/gelu.cc | 27 +++++++ .../contrib_ops/webgpu/bert/quick_gelu.cc | 56 +++++++++++++ .../contrib_ops/webgpu/bert/quick_gelu.h | 37 +++++++++ .../webgpu/webgpu_contrib_kernels.cc | 6 +- .../webgpu/math/unary_elementwise_ops.cc | 19 ++--- .../webgpu/math/unary_elementwise_ops.h | 26 ++++++ .../core/providers/webgpu/webgpu_utils.h | 20 +++++ 9 files changed, 289 insertions(+), 14 deletions(-) create mode 100644 onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc create mode 100644 onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.h create mode 100644 onnxruntime/contrib_ops/webgpu/bert/gelu.cc create mode 100644 onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc create mode 100644 onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h create mode 100644 onnxruntime/core/providers/webgpu/webgpu_utils.h diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc new file mode 100644 index 0000000000000..99cd643423400 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc @@ -0,0 +1,80 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "contrib_ops/webgpu/bert/bias_split_gelu.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/providers/webgpu/webgpu_utils.h" +#include "core/providers/webgpu/math/unary_elementwise_ops.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +ONNX_OPERATOR_KERNEL_EX( + BiasSplitGelu, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()), + BiasSplitGelu); + +Status BiasSplitGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { + const ShaderVariableHelper& input = shader.AddInput("input"); + const ShaderVariableHelper& bias = shader.AddInput("bias"); + const ShaderVariableHelper& output = shader.AddOutput("output"); + + shader.AdditionalImplementation() << ErfImpl; + + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") + << "const M_SQRT2: f32 = sqrt(2.0);\n" + << "const halfChannels = uniforms.channels / 2u;\n" + << "let biasIdx = global_idx % halfChannels;\n" + << "let batchIndex = global_idx / halfChannels;\n" + << "let inputOffset = biasIdx + batchIndex * halfChannels * 2;\n" + << "let valueLeft = " << input.GetByOffset("inputOffset") << " + " << bias.GetByOffset("biasIdx") << ";\n" + << "let valueRight = " << input.GetByOffset("inputOffset + halfChannels") << " + " << bias.GetByOffset("biasIdx + halfChannels") << ";\n" + << "let geluRight = valueRight * 0.5 * (erf_v(valueRight / M_SQRT2) + 1);\n" + << output.SetByOffset("global_idx", "valueLeft * geluRight"); + + return Status::OK(); +} + +Status BiasSplitGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + const auto* input = context.Input(0); + const auto* bias = context.Input(1); + + TensorShape input_shape = input->Shape(); + + if (input_shape.NumDimensions() != 3) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "BiasSplitGelu input should have 3 dimensions."); + } + + int64_t channels = input_shape[2]; + int64_t components = GetMaxComponents(channels); + channels /= components; + input_shape[2] = channels / 2; // for output shape calculation (N,S,D) -> (N,S,D/2) + + TensorShape bias_shape = bias->Shape(); + if (bias_shape.NumDimensions() != 1 || bias_shape[0] != channels) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "BiasSplitGelu bias should have 1 dimension with size equal to the number of channels."); + } + + auto* output = context.Output(0, input_shape); + int64_t output_size = output->Shape().Size() / components; + + BiasSplitGeluProgram program{}; + program.AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank}, + {bias}}) + .AddOutput({output}) + .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({{static_cast(output_size)}, + {static_cast(channels)}}); + return context.RunProgram(program); +} + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.h b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.h new file mode 100644 index 0000000000000..ccc3dd8c89b7b --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.h @@ -0,0 +1,32 @@ +// 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" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class BiasSplitGeluProgram final : public Program { + public: + BiasSplitGeluProgram() : Program{"BiasSplitGelu"} {} + Status GenerateShaderCode(ShaderHelper& sh) const override; + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32}, + {"channels", ProgramUniformVariableDataType::Uint32}); +}; + +class BiasSplitGelu final : public WebGpuKernel { + public: + BiasSplitGelu(const OpKernelInfo& info) : WebGpuKernel(info) {} + Status ComputeInternal(ComputeContext& context) const override; +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/gelu.cc new file mode 100644 index 0000000000000..8dafecfae83e5 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/gelu.cc @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/math/unary_elementwise_ops.h" // contains Gelu definition +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +ONNX_OPERATOR_KERNEL_EX( + Gelu, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()), + Gelu); + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc new file mode 100644 index 0000000000000..bcc1da4c7d37a --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/math/unary_elementwise_ops.h" +#include "contrib_ops/webgpu/bert/quick_gelu.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +ONNX_OPERATOR_KERNEL_EX( + QuickGelu, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()), + QuickGelu); + +Status QuickGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + const auto& y = shader.AddOutput("y", ShaderUsage::UseUniform); + + shader.AdditionalImplementation() << QuickGeluImpl; + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") + << " var a = " << x.GetByOffset("global_idx") << ";\n" + << y.SetByOffset("global_idx", onnxruntime::webgpu::QuickGeluExpr); + + return Status::OK(); +} + +Status QuickGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + const auto* input = context.Input(0); + auto* output = context.Output(0, input->Shape()); + + uint32_t data_size = gsl::narrow(output->Shape().Size()); + if (data_size == 0) { + return Status::OK(); + } + + const auto vec_size = (data_size + 3) / 4; + + QuickGeluProgram program{}; + program.AddInput({input, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) + .AddOutput({output, ProgramTensorMetadataDependency::None, {vec_size}, 4}) + .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({{vec_size}, {alpha_}}); + return context.RunProgram(program); +} + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h new file mode 100644 index 0000000000000..df65eccca86b0 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h @@ -0,0 +1,37 @@ +// 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" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class QuickGeluProgram final : public Program { + public: + QuickGeluProgram() : Program{"QuickGelu"} {} + Status GenerateShaderCode(ShaderHelper& sh) const override; + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"vec_size", ProgramUniformVariableDataType::Uint32}, + {"alpha", ProgramUniformVariableDataType::Float32}); +}; + +class QuickGelu final : public WebGpuKernel { + public: + QuickGelu(const OpKernelInfo& info) : WebGpuKernel(info) { + alpha_ = info.GetAttr("alpha"); + } + Status ComputeInternal(ComputeContext& context) const override; + + private: + float alpha_; +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 068a94c7390e2..6e63ba3a0caa4 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -38,14 +38,14 @@ Status RegisterWebGpuContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, // default entry to avoid the list become empty after ops-reducing // BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, // LayerNormalization used to be a contrib op that (incorrectly) used kOnnxDomain so we need to version it diff --git a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc index 189d7baafce6a..18db4b847f24c 100644 --- a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc +++ b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc @@ -285,17 +285,14 @@ class LinearUnit : public UnaryElementwise { WEBGPU_LU_IMPL(Elu, "elu_v(a)", EluImpl, 1.0) WEBGPU_ELEMENTWISE_KERNEL(Elu, 6, WebGpuSupportedFloatTypes()) -class Gelu : public UnaryElementwise { - public: - Gelu(const OpKernelInfo& info) - : UnaryElementwise{info, - "Gelu", - info.GetAttrOrDefault("approximate", "none") == "tanh" ? FastGeluExpr : GeluExpr, - info.GetAttrOrDefault("approximate", "none") == "tanh" ? TanhImpl : ErfImpl, - ShaderUsage::UseValueTypeAlias} { - cache_hint = info.GetAttrOrDefault("approximate", "none"); - } -}; +Gelu::Gelu(const OpKernelInfo& info) + : UnaryElementwise{info, + "Gelu", + info.GetAttrOrDefault("approximate", "none") == "tanh" ? FastGeluExpr : GeluExpr, + info.GetAttrOrDefault("approximate", "none") == "tanh" ? TanhImpl : ErfImpl, + ShaderUsage::UseValueTypeAlias} { + cache_hint = info.GetAttrOrDefault("approximate", "none"); +} WEBGPU_ELEMENTWISE_KERNEL(Gelu, 20, WebGpuSupportedFloatTypes()) diff --git a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h index 70fa81d21f95d..bdc9b7cc835dc 100644 --- a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h +++ b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h @@ -60,6 +60,11 @@ class UnaryElementwise : public WebGpuKernel { ShaderUsage additional_usage_; }; +class Gelu : public UnaryElementwise { + public: + Gelu(const OpKernelInfo& info); +} + constexpr const char ErfImpl[] = R"( const r0 = 0.3275911; const r1 = 0.254829592; @@ -104,11 +109,32 @@ fn elu_v(v: vec4) -> vec4 { } )"; +constexpr const char QuickGeluImpl[] = R"( +fn quick_gelu_v(a: x_value_t) -> x_value_t { + let one = 1.0; + let zero = 0.0; + let alpha_vec = x_value_t(uniforms.alpha); + let v = a * alpha_vec; + var x1 : x_value_t; + for (var i = 0; i < 4; i = i + 1) { + if (v[i] >= zero) { + x1[i] = one / (one + exp(-v[i])); + } else { + x1[i] = one - one / (one + exp(v[i])); + } + } + return a * x1; +} +)"; + // default GELU expression, depending on ErfImpl constexpr const char GeluExpr[] = "0.5 * a * (1.0 + erf_v(a * 0.7071067811865475))"; // fast GELU expression, depending on TanhImpl constexpr const char FastGeluExpr[] = "a * (0.5 + 0.5 * tanh_v(a * (0.035677408136300125 * a * a + 0.7978845608028654)))"; +// quick GELU expression, depending on QuickGeluImpl +constexpr const char QuickGeluExpr[] = "quick_gelu_v(a)"; + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/webgpu_utils.h b/onnxruntime/core/providers/webgpu/webgpu_utils.h new file mode 100644 index 0000000000000..4f9018646905d --- /dev/null +++ b/onnxruntime/core/providers/webgpu/webgpu_utils.h @@ -0,0 +1,20 @@ +// Licensed under the MIT License. + +#pragma once + +#include + +namespace onnxruntime { +namespace webgpu { + +inline int64_t GetMaxComponents(int64_t size) { + if (size % 4 == 0) { + return 4; + } else if (size % 2 == 0) { + return 2; + } + return 1; +} + +} // namespace webgpu +} // namespace onnxruntime \ No newline at end of file From 82a8e03edbfb8729cb5e71c700c4f56af6e07eb1 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 10 Mar 2025 16:26:38 -0700 Subject: [PATCH 2/5] bug fix --- onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h index bdc9b7cc835dc..86aa729efb375 100644 --- a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h +++ b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.h @@ -63,7 +63,7 @@ class UnaryElementwise : public WebGpuKernel { class Gelu : public UnaryElementwise { public: Gelu(const OpKernelInfo& info); -} +}; constexpr const char ErfImpl[] = R"( const r0 = 0.3275911; From ab47a5ac5372d11a746fed4d83805adfcc567fd6 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 10 Mar 2025 17:45:41 -0700 Subject: [PATCH 3/5] gsl::narrow to static_cast --- onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc index bcc1da4c7d37a..49763193f2d6a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc @@ -36,7 +36,7 @@ Status QuickGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const auto* input = context.Input(0); auto* output = context.Output(0, input->Shape()); - uint32_t data_size = gsl::narrow(output->Shape().Size()); + uint32_t data_size = static_cast(output->Shape().Size()); if (data_size == 0) { return Status::OK(); } From 06378067f3ef16d9458d18577b096a70cf8f0959 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Tue, 11 Mar 2025 16:41:19 -0700 Subject: [PATCH 4/5] quickgelu simple --- .../contrib_ops/webgpu/bert/quick_gelu.cc | 35 ++---------------- .../contrib_ops/webgpu/bert/quick_gelu.h | 37 ------------------- .../webgpu/math/unary_elementwise_ops.cc | 25 ++----------- .../webgpu/math/unary_elementwise_ops.h | 36 ++++++++++++++---- 4 files changed, 36 insertions(+), 97 deletions(-) delete mode 100644 onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc index 49763193f2d6a..bd2a05f58d615 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc @@ -4,13 +4,15 @@ #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/math/unary_elementwise_ops.h" -#include "contrib_ops/webgpu/bert/quick_gelu.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" namespace onnxruntime { namespace contrib { namespace webgpu { +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + ONNX_OPERATOR_KERNEL_EX( QuickGelu, kMSDomain, @@ -20,37 +22,6 @@ ONNX_OPERATOR_KERNEL_EX( .TypeConstraint("T", WebGpuSupportedFloatTypes()), QuickGelu); -Status QuickGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - const auto& y = shader.AddOutput("y", ShaderUsage::UseUniform); - - shader.AdditionalImplementation() << QuickGeluImpl; - shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") - << " var a = " << x.GetByOffset("global_idx") << ";\n" - << y.SetByOffset("global_idx", onnxruntime::webgpu::QuickGeluExpr); - - return Status::OK(); -} - -Status QuickGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { - const auto* input = context.Input(0); - auto* output = context.Output(0, input->Shape()); - - uint32_t data_size = static_cast(output->Shape().Size()); - if (data_size == 0) { - return Status::OK(); - } - - const auto vec_size = (data_size + 3) / 4; - - QuickGeluProgram program{}; - program.AddInput({input, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) - .AddOutput({output, ProgramTensorMetadataDependency::None, {vec_size}, 4}) - .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .AddUniformVariables({{vec_size}, {alpha_}}); - return context.RunProgram(program); -} - } // namespace webgpu } // namespace contrib } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h deleted file mode 100644 index df65eccca86b0..0000000000000 --- a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.h +++ /dev/null @@ -1,37 +0,0 @@ -// 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" - -namespace onnxruntime { -namespace contrib { -namespace webgpu { - -using namespace onnxruntime::webgpu; -using onnxruntime::webgpu::ComputeContext; - -class QuickGeluProgram final : public Program { - public: - QuickGeluProgram() : Program{"QuickGelu"} {} - Status GenerateShaderCode(ShaderHelper& sh) const override; - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"vec_size", ProgramUniformVariableDataType::Uint32}, - {"alpha", ProgramUniformVariableDataType::Float32}); -}; - -class QuickGelu final : public WebGpuKernel { - public: - QuickGelu(const OpKernelInfo& info) : WebGpuKernel(info) { - alpha_ = info.GetAttr("alpha"); - } - Status ComputeInternal(ComputeContext& context) const override; - - private: - float alpha_; -}; - -} // namespace webgpu -} // namespace contrib -} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc index 18db4b847f24c..e16327b9facad 100644 --- a/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc +++ b/onnxruntime/core/providers/webgpu/math/unary_elementwise_ops.cc @@ -256,26 +256,6 @@ WEBGPU_CLIP_KERNEL(MLFloat16) // activation // -class LinearUnit : public UnaryElementwise { - public: - LinearUnit(const OpKernelInfo& info, - const std::string& kernel_name, - const std::string& expression, - const std::string& additional_impl, - float default_alpha) - : UnaryElementwise{info, kernel_name, expression, additional_impl, ShaderUsage::UseElementTypeAlias} { - info.GetAttrOrDefault("alpha", &alpha_, default_alpha); - } - - Status ConfigureProgram(const ComputeContext& /*context*/, UnaryElementwiseProgram& program) const override { - program.AddUniformVariables({alpha_}); - return Status::OK(); - } - - protected: - float alpha_; -}; - #define WEBGPU_LU_IMPL(OP_TYPE, ...) \ class OP_TYPE final : public LinearUnit { \ public: \ @@ -294,6 +274,9 @@ Gelu::Gelu(const OpKernelInfo& info) cache_hint = info.GetAttrOrDefault("approximate", "none"); } +QuickGelu::QuickGelu(const OpKernelInfo& info) + : LinearUnit{info, "QuickGelu", "quick_gelu_v(a)", QuickGeluImpl, 1.702f} {} + WEBGPU_ELEMENTWISE_KERNEL(Gelu, 20, WebGpuSupportedFloatTypes()) WEBGPU_ELEMENTWISE_IMPL(Relu, "select(x_value_t(0), a, a > x_value_t(0))", "", ShaderUsage::UseValueTypeAlias) @@ -309,4 +292,4 @@ WEBGPU_LU_IMPL(ThresholdedRelu, "select(vec4(0), a, a > vec4) -> vec4 { )"; constexpr const char QuickGeluImpl[] = R"( -fn quick_gelu_v(a: x_value_t) -> x_value_t { +fn quick_gelu_v(a: vec4) -> vec4 { let one = 1.0; let zero = 0.0; - let alpha_vec = x_value_t(uniforms.alpha); + let alpha_vec = vec4(uniforms.attr); let v = a * alpha_vec; - var x1 : x_value_t; + var x1 : vec4; for (var i = 0; i < 4; i = i + 1) { if (v[i] >= zero) { x1[i] = one / (one + exp(-v[i])); @@ -133,8 +158,5 @@ constexpr const char GeluExpr[] = "0.5 * a * (1.0 + erf_v(a * 0.7071067811865475 // fast GELU expression, depending on TanhImpl constexpr const char FastGeluExpr[] = "a * (0.5 + 0.5 * tanh_v(a * (0.035677408136300125 * a * a + 0.7978845608028654)))"; -// quick GELU expression, depending on QuickGeluImpl -constexpr const char QuickGeluExpr[] = "quick_gelu_v(a)"; - } // namespace webgpu -} // namespace onnxruntime +} // namespace onnxruntime \ No newline at end of file From 910e3f70fa879b159484bcaa8f926705a636904b Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Tue, 11 Mar 2025 16:42:20 -0700 Subject: [PATCH 5/5] comment --- onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc index bd2a05f58d615..7d669e140ef23 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/quick_gelu.cc @@ -3,7 +3,7 @@ #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" -#include "core/providers/webgpu/math/unary_elementwise_ops.h" +#include "core/providers/webgpu/math/unary_elementwise_ops.h" // contained Gelu definition #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" namespace onnxruntime {