From 8ed2928dd5ea9779549d66faf6286e16126a0773 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Wed, 13 Nov 2019 09:26:00 -0800 Subject: [PATCH] Fuse Add + Gelu (#2360) Implement the transformer to fuse add + gelu Implement the accurate kernel --- .../cuda/activation/activations_impl.cu | 2 +- .../cuda/math/binary_elementwise_ops.cc | 70 +++++++++++++ .../cuda/math/binary_elementwise_ops.h | 29 ++++++ .../cuda/math/binary_elementwise_ops_impl.cu | 94 ++++++++++++++++++ .../cuda/math/binary_elementwise_ops_impl.h | 38 +++++++ .../contrib_ops/cuda_contrib_kernels.cc | 9 +- .../core/graph/contrib_ops/contrib_defs.cc | 34 ++++++- onnxruntime/core/optimizer/add_gelu_fusion.cc | 70 +++++++++++++ onnxruntime/core/optimizer/add_gelu_fusion.h | 23 +++++ .../core/optimizer/graph_transformer_utils.cc | 19 ++-- .../core/providers/cuda/cu_inc/common.cuh | 5 + .../test/contrib_ops/element_wise_ops_test.cc | 85 +++++++++++++++- .../test/optimizer/graph_transform_test.cc | 22 ++++ .../transform/fusion/add_gelu_fusion.onnx | Bin 0 -> 419 bytes .../testdata/transform/fusion/gelu_add_gen.py | 33 ++++++ 15 files changed, 519 insertions(+), 14 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc create mode 100644 onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.h create mode 100644 onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.h create mode 100644 onnxruntime/core/optimizer/add_gelu_fusion.cc create mode 100644 onnxruntime/core/optimizer/add_gelu_fusion.h create mode 100644 onnxruntime/test/testdata/transform/fusion/add_gelu_fusion.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gelu_add_gen.py diff --git a/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu b/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu index 59fce9e081..62601a1c69 100644 --- a/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu +++ b/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu @@ -39,7 +39,7 @@ struct OP_ScaledTanh : public CtxScaledTanh { template struct OP_Gelu : public CtxGelu { __device__ __inline__ T operator()(const T& a) const { - return a * _Normcdf(a); + return _Gelu(a); } }; diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc new file mode 100644 index 0000000000..5449bbb491 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "binary_elementwise_ops.h" +#include "binary_elementwise_ops_impl.h" + +using namespace onnxruntime::common; +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define CONTRIB_BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(x, ver, T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + x, \ + kMSDomain, \ + ver, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + x); + +#define CONTRIB_BINARY_ELEMENTWISE_COMPUTE(x, T) \ + template <> \ + Status x::ComputeInternal(OpKernelContext* context) const { \ + BinaryElementwisePreparation prepare(this); \ + ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \ + ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); \ + Impl_##x::MappedType>( \ + prepare.output_rank_or_simple_broadcast, \ + prepare.lhs_padded_strides.GpuPtr(), \ + reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), \ + prepare.rhs_padded_strides.GpuPtr(), \ + reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), \ + prepare.fdm_output_strides.GpuPtr(), \ + prepare.fdm_H, \ + prepare.fdm_C, \ + reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), \ + prepare.output_tensor->Shape().Size()); \ + return Status::OK(); \ + } + +#define CONTRIB_BINARY_OP_TYPED(name, ver, T) \ + CONTRIB_BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, T) \ + CONTRIB_BINARY_ELEMENTWISE_COMPUTE(name, T) + +// since different ops has different types, we cannot use BINARY_OPS() directly +// the postfix of means the types supported by the op: +// B: uint8_t +// W: uint16_t +// U: uint32_t +// Z: uint64_t +// C: int8_t +// S: int16_t +// I: int32_t +// L: int64_t +// H: float16 +// F: float +// D: double +// O: bool + +#define CONTRIB_BINARY_OP_HFD(name, ver) \ + CONTRIB_BINARY_OP_TYPED(name, ver, MLFloat16) \ + CONTRIB_BINARY_OP_TYPED(name, ver, float) \ + CONTRIB_BINARY_OP_TYPED(name, ver, double) + +CONTRIB_BINARY_OP_HFD(AddGeluFusion, 1) + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.h b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.h new file mode 100644 index 0000000000..69dba848a5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/math/binary_elementwise_ops.h" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/shared_inc/fast_divmod.h" +#include "core/providers/cpu/tensor/utils.h" + +using namespace onnxruntime::cuda; + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// AddGelu fuse Add + Gelu +template +class AddGeluFusion final : public BinaryElementwise { + public: + AddGeluFusion(const OpKernelInfo& info) : BinaryElementwise(info) { + } + + Status ComputeInternal(OpKernelContext* context) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu new file mode 100644 index 0000000000..db4fcb18bb --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include "binary_elementwise_ops_impl.h" +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cu_inc/binary_elementwise_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define OP(name, expr) \ + template \ + struct OP_##name { \ + __device__ __inline__ T operator()(T a, T b) const { \ + return (expr); \ + } \ + }; + +#define CONTRIB_BINARY_ELEMENTWISE_IMPL(name) \ + CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name) { \ + BinaryElementWiseImpl(output_rank_or_simple_broadcast, \ + lhs_padded_strides, \ + lhs_data, \ + rhs_padded_strides, \ + rhs_data, \ + fdm_output_strides, \ + fdm_H, \ + fdm_C, \ + output_data, \ + OP_##name(), \ + count); \ + } + +#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \ + template void Impl_##x(size_t output_rank, \ + const int64_t* lhs_padded_strides, \ + const T* lhs_data, \ + const int64_t* rhs_padded_strides, \ + const T* rhs_data, \ + const onnxruntime::cuda::fast_divmod* fdm_output_strides, \ + const onnxruntime::cuda::fast_divmod& fdm_H, \ + const onnxruntime::cuda::fast_divmod& fdm_C, \ + T* output_data, size_t count); + +#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) + +#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_OIL(x) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, bool) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t) + +#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(x) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) + +// create declarations for op and impl +#define CONTRIB_BINARY_OP_NAME_EXPR(name, expr) \ + OP(name, expr) \ + CONTRIB_BINARY_ELEMENTWISE_IMPL(name) + +CONTRIB_BINARY_OPS() + +#undef CONTRIB_BINARY_OP_NAME_EXPR + +// create specialized impl +// the postfix of means the types supported by the op: +// B: uint8_t +// W: uint16_t +// U: uint32_t +// Z: uint64_t +// C: int8_t +// S: int16_t +// I: int32_t +// L: int64_t +// H: float16 +// F: float +// D: double +// O: bool + +CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(AddGeluFusion) + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.h b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.h new file mode 100644 index 0000000000..ee6cd680ae --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.h @@ -0,0 +1,38 @@ +#include +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +using namespace onnxruntime::cuda; + +namespace onnxruntime { +namespace contrib { +namespace cuda { +// These macros simplifies coding. To add a new op with following steps: +// 1. Add a new entry in CONTRIB_BINARY_OPS() list +// 2. (optional) Define templated single element operator in binary_elementwise_ops_impl.cu +// 3. (optional) Implement specialized single element operator +// 4. Add op kernel class definition in binary_elementwise_ops.h +// 5. Add op kernel registration and compute specialization in binary_elementwise_ops.cc +#define CONTRIB_BINARY_OPS() \ + CONTRIB_BINARY_OP_NAME_EXPR(AddGeluFusion, _Gelu(a + b)) + +// NOTE that cu files are compiled with nvcc and should not refer to any onnxruntime headers +// so struct BinaryElementwisePreparation cannot be used here +#define CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name) \ + template \ + void Impl_##name( \ + size_t output_rank_or_simple_broadcast, \ + const int64_t* lhs_padded_strides, \ + const T* lhs_data, \ + const int64_t* rhs_padded_strides, \ + const T* rhs_data, \ + const onnxruntime::cuda::fast_divmod* fdm_output_strides, \ + const onnxruntime::cuda::fast_divmod& fdm_H, \ + const onnxruntime::cuda::fast_divmod& fdm_C, \ + T* output_data, \ + size_t count) +#define CONTRIB_BINARY_OP_NAME_EXPR(name, expr) CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name); +CONTRIB_BINARY_OPS() +#undef CONTRIB_BINARY_OP_NAME_EXPR +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc index 799a86776f..31790b899b 100644 --- a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc @@ -14,6 +14,9 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Gelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Gelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Gelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, AddGeluFusion); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, AddGeluFusion); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, AddGeluFusion); // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility @@ -55,6 +58,9 @@ void RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility @@ -87,8 +93,7 @@ void RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo - }; + BuildKernelCreateInfo}; for (auto& function_table_entry : function_table) { kernel_registry.Register(function_table_entry()); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 82ef303efa..c05551b2ff 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -193,7 +193,6 @@ void RegisterNchwcSchemas() { } void RegisterBertSchemas() { - ONNX_CONTRIB_OPERATOR_SCHEMA(Attention) .SetDomain(kMSDomain) .SinceVersion(1) @@ -337,7 +336,6 @@ void RegisterBertSchemas() { .Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T") .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); - } void RegisterContribSchemas() { @@ -1926,11 +1924,18 @@ Example 4: RegisterNchwcSchemas(); } + static const char* Gelu_ver1_doc = + R"DOC(Gaussian Error Linear Unit. +A high-performing neural network activation function.The GELU nonlinearity is +the expected transformation of a stochastic regularizer which randomly applies +the identity or zero map to a neuron's input. The GELU nonlinearity weights +inputs by their magnitude, rather than gates inputs by their sign as in ReLUs.)DOC"; + ONNX_CONTRIB_OPERATOR_SCHEMA(Gelu) .SetDomain(kMSDomain) .SinceVersion(1) .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) - .SetDoc("Gelu") + .SetDoc(Gelu_ver1_doc) .Input(0, "X", "The input data as Tensor.", "T") .Output(0, "Y", "The output.", "T") .TypeConstraint( @@ -1939,6 +1944,29 @@ Example 4: "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); + ONNX_CONTRIB_OPERATOR_SCHEMA(AddGeluFusion) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) + .SetDoc("AddGeluFusion fuses Add+Gelu. The fused Add op is the parent node of the fused Gelu.") + .Input(0, "A", "The input data as Tensor that is the first input of fused Add.", "T") + .Input(1, "B", "The input data as Tensor that is the second input of fused Add.", "T") + .Output(0, "C", "The output.", "T") + .TypeConstraint( + "T", + {"tensor(float16)", "tensor(float)", "tensor(double)"}, + "Constrain input and output types to float tensors.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + + if (hasNInputShapes(ctx, 2)) { + bidirectionalBroadcastShapeInference( + ctx.getInputType(0)->tensor_type().shape(), + ctx.getInputType(1)->tensor_type().shape(), + *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()); + } + }); + RegisterBertSchemas(); #ifdef MICROSOFT_INTERNAL diff --git a/onnxruntime/core/optimizer/add_gelu_fusion.cc b/onnxruntime/core/optimizer/add_gelu_fusion.cc new file mode 100644 index 0000000000..133079fcf5 --- /dev/null +++ b/onnxruntime/core/optimizer/add_gelu_fusion.cc @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/initializer.h" +#include "core/optimizer/add_gelu_fusion.h" +#include "core/graph/graph_utils.h" +#include + +using namespace ONNX_NAMESPACE; +using namespace ::onnxruntime::common; +namespace onnxruntime { + +Status AddGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level) const { + GraphViewer graph_viewer(graph); + const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); + + for (auto node_index : node_topology_list) { + auto* node_ptr = graph.GetNode(node_index); + if (nullptr == node_ptr) + continue; // node was removed + + auto& node = *node_ptr; + + ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level)); + + if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Add", {7}) || + !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders()) || + node.GetOutputEdgesCount() != 1) { + continue; + } + + auto next_node_itr = node.OutputNodesBegin(); + if (next_node_itr == node.OutputNodesEnd()) { + continue; + } + + const Node& next_node = (*next_node_itr); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(next_node, "Gelu", {1}, kMSDomain) || + next_node.GetExecutionProviderType() != node.GetExecutionProviderType()) { + continue; + } + + if (!graph.GetNodeOutputsInGraphOutputs(node).empty()) { + continue; + } + + Node& add_node = node; + Node& gelu_node = const_cast(next_node); + + Node& gelu_add_fusion_node = graph.AddNode(graph.GenerateNodeName("AddGeluFusion"), + "AddGeluFusion", + "fused Add and Gelu", + {add_node.MutableInputDefs()}, + {}, + {}, + kMSDomain); + + // Assign provider to this new node. Provider should be same as the provider for old node. + gelu_add_fusion_node.SetExecutionProviderType(gelu_node.GetExecutionProviderType()); + + // move output definitions and edges from gelu_node to gelu_add_fusion_node + //delete add_node and gelu_node. + graph_utils::FinalizeNodeFusion(graph, {add_node, gelu_node}, gelu_add_fusion_node); + + modified = true; + } + + return Status::OK(); +} +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/add_gelu_fusion.h b/onnxruntime/core/optimizer/add_gelu_fusion.h new file mode 100644 index 0000000000..5c62046a4f --- /dev/null +++ b/onnxruntime/core/optimizer/add_gelu_fusion.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** +@Class GeluFusion +Fuse Add + Gelu to GeluFusion +*/ +class AddGeluFusion : public GraphTransformer { + public: + AddGeluFusion(const std::unordered_set& compatible_execution_providers = {}) noexcept + : GraphTransformer("AddGeluFusion", compatible_execution_providers) { + } + + Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 8e1fa16536..d26ddf86c4 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -18,6 +18,7 @@ #include "core/optimizer/shape_to_initializer.h" #include "core/optimizer/nchwc_transformer.h" #include "core/optimizer/free_dim_override_transformer.h" +#include "core/optimizer/add_gelu_fusion.h" #include "core/optimizer/gelu_fusion.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/mlas/inc/mlas.h" @@ -108,17 +109,23 @@ std::vector> GenerateTransformers(TransformerL } break; case TransformerLevel::Level2: { - std::unordered_set l2_execution_providers = {onnxruntime::kCpuExecutionProvider}; + std::unordered_set cpu_execution_providers = {onnxruntime::kCpuExecutionProvider}; // create rule based transformer consisting of all the level2 rewrite rules - rule_transformer = GenerateRuleBasedGraphTransformer(level, transformers_and_rules_to_enable, l2_execution_providers); + rule_transformer = GenerateRuleBasedGraphTransformer(level, transformers_and_rules_to_enable, cpu_execution_providers); // create standalone transformers #ifndef DISABLE_CONTRIB_OPS - transformers.emplace_back(onnxruntime::make_unique(l2_execution_providers)); - transformers.emplace_back(onnxruntime::make_unique(l2_execution_providers)); - transformers.emplace_back(onnxruntime::make_unique(l2_execution_providers)); - transformers.emplace_back(onnxruntime::make_unique(l2_execution_providers)); + transformers.emplace_back(onnxruntime::make_unique(cpu_execution_providers)); + transformers.emplace_back(onnxruntime::make_unique(cpu_execution_providers)); + + std::unordered_set cpu_cuda_execution_providers = {onnxruntime::kCpuExecutionProvider, onnxruntime::kCudaExecutionProvider}; + transformers.emplace_back(onnxruntime::make_unique(cpu_cuda_execution_providers)); + transformers.emplace_back(onnxruntime::make_unique(cpu_cuda_execution_providers)); + + std::unordered_set cuda_execution_providers = {onnxruntime::kCudaExecutionProvider}; + transformers.emplace_back(onnxruntime::make_unique(cuda_execution_providers)); + #endif } break; diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index b440519a8a..a685396a75 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -196,6 +196,11 @@ __device__ __inline__ double _Normcdf(double a) { return normcdf(a); } template <> __device__ __inline__ half _Normcdf(half a) { return half(normcdff((float)a)); } +template +__device__ __inline__ T _Gelu(T a) { + return a * _Normcdf(a); +} + // We would like to use 64-bit integer to support large matrices. However, CUDA seems to support only 32-bit integer // For now, use int32_t to ensure that both Linux and Windows see this as 32 bit integer type. diff --git a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc index 434b2f89e7..3cbc75ffe4 100644 --- a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc +++ b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc @@ -2,11 +2,16 @@ // Licensed under the MIT License. #include "gtest/gtest.h" -#include "test/providers/provider_test_utils.h" -#include "core/util/math.h" + #include #include +#include "core/util/math.h" +#include "core/mlas/inc/mlas.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" + using namespace onnxruntime::test; namespace onnxruntime { @@ -47,5 +52,81 @@ TEST(MathOpTest, Scale_Default) { test.Run(); } +std::vector Add_Simple(const std::vector& input_a_data, const std::vector& input_b_data) { + EXPECT_TRUE(input_a_data.size() % input_b_data.size() == 0 || input_b_data.size() % input_a_data.size() == 0); + const std::vector& input_large_size = input_a_data.size() >= input_b_data.size() ? input_a_data : input_b_data; + const std::vector& input_small_size = input_a_data.size() < input_b_data.size() ? input_a_data : input_b_data; + + std::vector output(input_large_size.size()); + for (int iter = 0; iter < input_large_size.size() / input_small_size.size(); iter++) { + std::transform(input_large_size.begin() + iter * input_small_size.size(), + input_large_size.begin() + (iter + 1) * input_small_size.size(), + input_small_size.begin(), + output.begin() + iter * input_small_size.size(), + [](float a, float b) { + return a + b; + }); + } + return output; +} + +const std::vector ComputeGeluWithErf(const std::vector& input_data) { + std::vector output(input_data.size()); + + std::transform(input_data.begin(), + input_data.end(), + output.begin(), + [](float x) { + float y = erf(x * static_cast(M_SQRT1_2)); + return x * 0.5f * (y + 1.0f); + }); + + return output; +} + +static void RunAddGeluFusionTest( + const std::vector& input_a_data, + const std::vector& input_b_data, + const std::vector& input_a_dims, + const std::vector& input_b_dims) { + if (HasCudaEnvironment(0)) { + std::vector output_data = ComputeGeluWithErf(Add_Simple(input_a_data, input_b_data)); + + OpTester tester("AddGeluFusion", 1, onnxruntime::kMSDomain); + + const std::vector& output_dims = input_a_dims.size() >= input_b_dims.size() ? input_a_dims : input_b_dims; + tester.AddInput("A", input_a_dims, input_a_data); + tester.AddInput("B", input_b_dims, input_b_data); + tester.AddOutput("C", output_dims, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +TEST(AddGeluFusionTest, Two_One_Dim) { + std::vector input_a_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector input_b_data = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + RunAddGeluFusionTest(input_a_data, input_b_data, {2, 4}, {4}); +} + +TEST(AddGeluFusionTest, Two_Two_Dim) { + std::vector input_a_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector input_b_data = { + -0.5f, 0.6f, 1.2f, 2.1f, + 0.4f, 0.6f, 0.2f, -0.4f}; + + RunAddGeluFusionTest(input_a_data, input_b_data, {2, 4}, {2, 4}); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 81e19a1751..0a568eb829 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -14,6 +14,7 @@ #include "core/optimizer/conv_activation_fusion.h" #include "core/optimizer/dropout_elimination.h" #include "core/optimizer/gemm_activation_fusion.h" +#include "core/optimizer/add_gelu_fusion.h" #include "core/optimizer/gelu_fusion.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/optimizer/graph_transformer.h" @@ -820,6 +821,27 @@ TEST(GraphTransformationTests, GeluFusionTest) { ASSERT_TRUE(op_to_count["Mul"] == 0); ASSERT_TRUE(op_to_count["Gelu"] == 1); } + +TEST(GraphTransformationTests, AddGeluFusionTest) { + string model_uri = MODEL_FOLDER + "fusion/add_gelu_fusion.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2); + ASSERT_TRUE(ret.IsOK()); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Div"] == 0); + ASSERT_TRUE(op_to_count["Add"] == 0); + ASSERT_TRUE(op_to_count["Erf"] == 0); + ASSERT_TRUE(op_to_count["Mul"] == 0); + ASSERT_TRUE(op_to_count["Gelu"] == 0); + ASSERT_TRUE(op_to_count["GeluFusion"] == 0); +} + TEST(GraphTransformationTests, LayerNormFusionTest) { string model_uri = MODEL_FOLDER + "fusion/layer_norm.onnx"; std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/add_gelu_fusion.onnx b/onnxruntime/test/testdata/transform/fusion/add_gelu_fusion.onnx new file mode 100644 index 0000000000000000000000000000000000000000..d495584a739e7ff0a51925fffd9ec25f1e9e3633 GIT binary patch literal 419 zcmaiwJ&wXK5QXD)P{K4IZ?TA<0#`sq*sW+*34awAps@gh6cG{$Fg54cdvG!?uww@Y zL__!V{QS+E@k_%$3JoA3kcjz9Df=Q*r@E^Zr=gV4fx?EejL%}7rN