From e57b735bb9a4912eb38f9de9d272d42c03779569 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 27 Nov 2019 10:15:50 -0800 Subject: [PATCH] Add a transformer to use Gelu approximation for cuda provider (#2480) * Add Gelu Approximation Transformer to convert Gelu or AddGeluFusion to FastGelu to get better inference performance. --- .../core/optimizer/graph_transformer_level.h | 5 +- .../core/optimizer/gelu_approximation.cc | 119 ++++++++++++++++++ .../core/optimizer/gelu_approximation.h | 24 ++++ .../core/optimizer/graph_transformer_utils.cc | 3 +- .../core/session/abi_session_options.cc | 2 +- onnxruntime/core/session/inference_session.cc | 20 ++- .../test/framework/inference_session_test.cc | 2 +- .../test/optimizer/graph_transform_test.cc | 70 +++++++++++ .../transform/approximation/gelu.onnx | 13 ++ .../approximation/gelu_add_bias.onnx | 17 +++ .../approximation/gelu_add_matmul.onnx | 25 ++++ .../gelu_add_shape_not_match.onnx | 19 +++ .../approximation/gelu_approximation_gen.py | 86 +++++++++++++ 13 files changed, 387 insertions(+), 18 deletions(-) create mode 100644 onnxruntime/core/optimizer/gelu_approximation.cc create mode 100644 onnxruntime/core/optimizer/gelu_approximation.h create mode 100644 onnxruntime/test/testdata/transform/approximation/gelu.onnx create mode 100644 onnxruntime/test/testdata/transform/approximation/gelu_add_bias.onnx create mode 100644 onnxruntime/test/testdata/transform/approximation/gelu_add_matmul.onnx create mode 100644 onnxruntime/test/testdata/transform/approximation/gelu_add_shape_not_match.onnx create mode 100644 onnxruntime/test/testdata/transform/approximation/gelu_approximation_gen.py diff --git a/include/onnxruntime/core/optimizer/graph_transformer_level.h b/include/onnxruntime/core/optimizer/graph_transformer_level.h index 4f2d5b305c..7aeb00ba66 100644 --- a/include/onnxruntime/core/optimizer/graph_transformer_level.h +++ b/include/onnxruntime/core/optimizer/graph_transformer_level.h @@ -12,9 +12,8 @@ enum class TransformerLevel : int { Level1, Level2, Level3, - // Convenience enum to always get the max available value. - // This way when we add more levels code which iterates over this enum does not need to change. - MaxTransformerLevel + // The max level should always be same as the last level. + MaxLevel = Level3 }; } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/gelu_approximation.cc b/onnxruntime/core/optimizer/gelu_approximation.cc new file mode 100644 index 0000000000..9982a3836d --- /dev/null +++ b/onnxruntime/core/optimizer/gelu_approximation.cc @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/initializer.h" +#include "core/optimizer/gelu_approximation.h" +#include "core/framework/tensorprotoutils.h" +#include "core/optimizer/utils.h" +#include "core/graph/graph_utils.h" +#include "float.h" + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; +namespace onnxruntime { + +// FastGelu supports limited data types. +static std::vector supported_data_types{"tensor(float16)", "tensor(float)"}; + +static bool IsSupportedDataType(const Node& node) { + for (const auto& input_arg : node.InputDefs()) { + if (std::find(supported_data_types.begin(), supported_data_types.end(), + *(input_arg->Type())) == supported_data_types.end()) { + return false; + } + } + return true; +} + +static bool CheckInputShape(const Node& node, const NodeArg& input, const NodeArg& bias) { + const TensorShapeProto* bias_shape = bias.Shape(); + if (nullptr == bias_shape || + bias_shape->dim_size() != 1 || + !utils::HasDimValue(bias_shape->dim(0))) { + return false; + } + auto bias_length = bias_shape->dim(0).dim_value(); + + const TensorShapeProto* input_shape = input.Shape(); + if (nullptr != input_shape) { + if (input_shape->dim_size() >= 1) { + int last_dim = input_shape->dim_size() - 1; + if (utils::HasDimValue(input_shape->dim(last_dim)) && + input_shape->dim(last_dim).dim_value() == bias_length) { + return true; + } + } + return false; + } + + // Input does not have shape. We will check its parent node. + // When the parent is MatMul and its 2nd input has shape like {*, bias_length}, + // it means that the shape of MatMul output is good for FastGelu. + const Node* parent_node = graph_utils::GetInputNode(node, 0); + if (nullptr != parent_node && + graph_utils::IsSupportedOptypeVersionAndDomain(*parent_node, "MatMul", {1, 9}, kOnnxDomain)) { + const NodeArg& input_b = *(parent_node->InputDefs()[1]); + if (optimizer_utils::ValidateShape(input_b, {-1, bias_length})) { + return true; + } + } + + return false; +} + +static bool CheckGeluInputShape(const NodeArg& input) { + const TensorShapeProto* input_shape = input.Shape(); + return nullptr != input_shape && input_shape->dim_size() >= 1; +} + +static bool IsCandidateNode(Node& node, const std::unordered_set& compatible_providers) { + if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "AddGeluFusion", {1}, kMSDomain)) { + return graph_utils::IsSupportedProvider(node, compatible_providers) && + IsSupportedDataType(node) && + CheckInputShape(node, *(node.InputDefs()[0]), *(node.InputDefs()[1])); + } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Gelu", {1}, kMSDomain)) { + return graph_utils::IsSupportedProvider(node, compatible_providers) && + IsSupportedDataType(node) && + CheckGeluInputShape(*(node.InputDefs()[0])); + } + return false; +} + +Status GeluApproximation::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { + GraphViewer graph_viewer(graph); + const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); + + int count = 0; + for (auto node_index : node_topology_list) { + auto* p_node = graph.GetNode(node_index); + if (p_node == nullptr) + continue; // we removed the node as part of an earlier fusion + + Node& node = *p_node; + ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger)); + + if (IsCandidateNode(node, GetCompatibleExecutionProviders())) { + Node& fastgelu = graph.AddNode( + graph.GenerateNodeName("FastGelu"), + "FastGelu", + "Gelu approximation", + node.MutableInputDefs(), + node.MutableOutputDefs(), nullptr, kMSDomain); + + fastgelu.SetExecutionProviderType(node.GetExecutionProviderType()); + + graph_utils::RemoveNodeOutputEdges(graph, node); + graph.RemoveNode(node.Index()); + + count++; + } + } + + if (count > 0) { + modified = true; + LOGS(logger, INFO) << "Total Gelu Approximation (FastGelu) node count: " << count; + } + + return Status::OK(); +} +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/gelu_approximation.h b/onnxruntime/core/optimizer/gelu_approximation.h new file mode 100644 index 0000000000..def16eec92 --- /dev/null +++ b/onnxruntime/core/optimizer/gelu_approximation.h @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** +@Class GeluApproximation + +Rewrite graph to replace Gelu or AddGeluFusion by FastGelu node. FastGelu uses approximation for Gelu, +and it is faster. +*/ +class GeluApproximation : public GraphTransformer { + public: + GeluApproximation(const std::unordered_set& compatible_execution_providers={}) noexcept + : GraphTransformer("GeluApproximation", compatible_execution_providers) {} + + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 8845752d2c..3f0b6c1b43 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -20,6 +20,7 @@ #include "core/optimizer/free_dim_override_transformer.h" #include "core/optimizer/add_gelu_fusion.h" #include "core/optimizer/gelu_fusion.h" +#include "core/optimizer/gelu_approximation.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/reshape_fusion.h" @@ -130,6 +131,7 @@ std::vector> GenerateTransformers(TransformerL std::unordered_set cuda_execution_providers = {onnxruntime::kCudaExecutionProvider}; transformers.emplace_back(onnxruntime::make_unique(cuda_execution_providers)); + transformers.emplace_back(onnxruntime::make_unique(cuda_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(cuda_execution_providers)); #endif @@ -142,7 +144,6 @@ std::vector> GenerateTransformers(TransformerL transformers.emplace_back(onnxruntime::make_unique()); } #endif - } break; default: diff --git a/onnxruntime/core/session/abi_session_options.cc b/onnxruntime/core/session/abi_session_options.cc index 026e7425fc..d032e2dbd5 100644 --- a/onnxruntime/core/session/abi_session_options.cc +++ b/onnxruntime/core/session/abi_session_options.cc @@ -130,7 +130,7 @@ ORT_API_STATUS_IMPL(OrtApis::SetSessionGraphOptimizationLevel, _In_ OrtSessionOp options->value.graph_optimization_level = onnxruntime::TransformerLevel::Level2; break; case ORT_ENABLE_ALL: - options->value.graph_optimization_level = onnxruntime::TransformerLevel::Level3; + options->value.graph_optimization_level = onnxruntime::TransformerLevel::MaxLevel; break; default: return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "graph_optimization_level is not valid"); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 5b13df4387..c9aad03a7d 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -404,7 +404,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, // apply transformers except default transformers // Default transformers are required for correctness and they are owned and run by inference session - for (int i = static_cast(TransformerLevel::Level1); i < static_cast(TransformerLevel::MaxTransformerLevel); i++) { + for (int i = static_cast(TransformerLevel::Level1); i <= static_cast(TransformerLevel::MaxLevel); i++) { ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformer_mgr.ApplyTransformers(graph, static_cast(i), *session_logger_)); } @@ -1141,21 +1141,17 @@ void InferenceSession::AddPredefinedTransformers(GraphTransformerManager& transf } }; - ORT_ENFORCE(graph_optimization_level < TransformerLevel::MaxTransformerLevel, - "Allowed values are 1 and 2. Current level is set to " + + ORT_ENFORCE(graph_optimization_level <= TransformerLevel::MaxLevel, + "Exceeded max transformer level. Current level is set to " + std::to_string(static_cast(graph_optimization_level))); - if ((graph_optimization_level >= TransformerLevel::Level1) || !custom_list.empty()) { - add_transformers(TransformerLevel::Level1); + for (int i = static_cast(TransformerLevel::Level1); i <= static_cast(TransformerLevel::MaxLevel); i++) { + TransformerLevel level = static_cast(i); + if ((graph_optimization_level >= level) || !custom_list.empty()) { + add_transformers(level); + } } - if ((graph_optimization_level >= TransformerLevel::Level2) || !custom_list.empty()) { - add_transformers(TransformerLevel::Level2); - } - - if ((graph_optimization_level >= TransformerLevel::Level3) || !custom_list.empty()) { - add_transformers(TransformerLevel::Level3); - } } common::Status InferenceSession::WaitForNotification(Notification* p_executor_done, int64_t timeout_in_ms) { diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index 44f221aa3e..e5abaa7176 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -1428,7 +1428,7 @@ TEST(InferenceSessionTests, TestCopyToFromDevices) { TEST(InferenceSessionTests, TestRegisterTransformers) { string model_uri = "testdata/transform/fusion/fuse-conv-bn-mul-add-unsqueeze.onnx"; - for (int i = static_cast(TransformerLevel::Default); i < static_cast(TransformerLevel::MaxTransformerLevel); i++) { + for (int i = static_cast(TransformerLevel::Default); i <= static_cast(TransformerLevel::MaxLevel); i++) { SessionOptions so; so.session_logid = "InferenceSessionTests.TestL1AndL2Transformers"; so.graph_optimization_level = static_cast(i); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 9ef55cab71..9915505477 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -16,6 +16,7 @@ #include "core/optimizer/gemm_activation_fusion.h" #include "core/optimizer/add_gelu_fusion.h" #include "core/optimizer/gelu_fusion.h" +#include "core/optimizer/gelu_approximation.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/graph_transformer.h" @@ -1101,6 +1102,75 @@ TEST(GraphTransformationTests, AddGeluFusionTest) { ASSERT_TRUE(op_to_count["GeluFusion"] == 0); } +// Test Gelu -> FastGelu +TEST(GraphTransformationTests, GeluApproximation_Gelu) { + auto model_uri = MODEL_FOLDER "approximation/gelu.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["Gelu"], 0); + EXPECT_EQ(op_to_count["FastGelu"], 1); +} + +// Test AddGeluFusion -> FastGelu +TEST(GraphTransformationTests, GeluApproximation_Gelu_Add_Bias) { + auto model_uri = MODEL_FOLDER "approximation/gelu_add_bias.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["AddGeluFusion"], 0); + EXPECT_EQ(op_to_count["FastGelu"], 1); +} + +// Test MatMul & AddGeluFusion -> MatMul & FastGelu +TEST(GraphTransformationTests, GeluApproximation_Gelu_Add_MatMul) { + auto model_uri = MODEL_FOLDER "approximation/gelu_add_matmul.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["AddGeluFusion"], 0); + EXPECT_EQ(op_to_count["MatMul"], 1); + EXPECT_EQ(op_to_count["FastGelu"], 1); +} + +// Test AddGeluFusion with mis-match bias shape cannot convert to FastGelu. +TEST(GraphTransformationTests, GeluApproximation_Gelu_Add_Shape_Not_Match) { + auto model_uri = MODEL_FOLDER "approximation/gelu_add_shape_not_match.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["AddGeluFusion"], 1); + EXPECT_EQ(op_to_count["FastGelu"], 0); +} + TEST(GraphTransformationTests, LayerNormFusionTest) { auto model_uri = MODEL_FOLDER "fusion/layer_norm.onnx"; std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/approximation/gelu.onnx b/onnxruntime/test/testdata/transform/approximation/gelu.onnx new file mode 100644 index 0000000000..b183a1fc84 --- /dev/null +++ b/onnxruntime/test/testdata/transform/approximation/gelu.onnx @@ -0,0 +1,13 @@ +:~ +# +ACGelu_1"Gelu: com.microsoft Gelu_NoBiasZ$ +A + +batch + seq_len +€b$ +C + +batch + seq_len +€B \ No newline at end of file diff --git a/onnxruntime/test/testdata/transform/approximation/gelu_add_bias.onnx b/onnxruntime/test/testdata/transform/approximation/gelu_add_bias.onnx new file mode 100644 index 0000000000..b96cb4752a --- /dev/null +++ b/onnxruntime/test/testdata/transform/approximation/gelu_add_bias.onnx @@ -0,0 +1,17 @@ +:¦ +8 +A +BCAddGeluFusion_1" AddGeluFusion: com.microsoft Gelu_AddBiasZ$ +A + +batch + seq_len +€Z +B +  +€b$ +C + +batch + seq_len +€B \ No newline at end of file diff --git a/onnxruntime/test/testdata/transform/approximation/gelu_add_matmul.onnx b/onnxruntime/test/testdata/transform/approximation/gelu_add_matmul.onnx new file mode 100644 index 0000000000..688fa1cd68 --- /dev/null +++ b/onnxruntime/test/testdata/transform/approximation/gelu_add_matmul.onnx @@ -0,0 +1,25 @@ +:â + +A +BCMatMul_1"MatMul +8 +C +DEAddGeluFusion_1" AddGeluFusion: com.microsoftMatMul_AddGeluFusionZ$ +A + +batch + seq_len +xZ +B + + +€ +€Z +D +  +€b$ +E + +batch + seq_len +€B \ No newline at end of file diff --git a/onnxruntime/test/testdata/transform/approximation/gelu_add_shape_not_match.onnx b/onnxruntime/test/testdata/transform/approximation/gelu_add_shape_not_match.onnx new file mode 100644 index 0000000000..5061df5da8 --- /dev/null +++ b/onnxruntime/test/testdata/transform/approximation/gelu_add_shape_not_match.onnx @@ -0,0 +1,19 @@ +:Ä +8 +A +BCAddGeluFusion_1" AddGeluFusion: com.microsoftGelu_Add_ShapeNotMatchZ$ +A + +batch + seq_len +€Z$ +B + +batch + seq_len +€b$ +C + +batch + seq_len +€B \ No newline at end of file diff --git a/onnxruntime/test/testdata/transform/approximation/gelu_approximation_gen.py b/onnxruntime/test/testdata/transform/approximation/gelu_approximation_gen.py new file mode 100644 index 0000000000..65dae0cbe8 --- /dev/null +++ b/onnxruntime/test/testdata/transform/approximation/gelu_approximation_gen.py @@ -0,0 +1,86 @@ +import onnx +from onnx import helper +from onnx import TensorProto + +graph = helper.make_graph( + [ # nodes + # Add node before Gelu + helper.make_node("Gelu", ["A"], ["C"], "Gelu_1", domain="com.microsoft"), + ], + "Gelu_NoBias", #name + [ # inputs + helper.make_tensor_value_info('A', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + ], + [ # outputs + helper.make_tensor_value_info('C', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + ], + [ # initializers + ] +) + +model = helper.make_model(graph) +onnx.save(model, r'gelu.onnx') + +graph = helper.make_graph( + [ # nodes + # Add node before Gelu + helper.make_node("AddGeluFusion", ["A", "B"], ["C"], "AddGeluFusion_1", domain="com.microsoft"), + ], + "Gelu_AddBias", #name + [ # inputs + helper.make_tensor_value_info('A', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + helper.make_tensor_value_info('B', TensorProto.FLOAT, [3072]), + ], + [ # outputs + helper.make_tensor_value_info('C', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + ], + [ # initializers + ] +) + +model = helper.make_model(graph) +onnx.save(model, r'gelu_add_bias.onnx') + +graph = helper.make_graph( + [ # nodes + # Add node before Gelu + helper.make_node("AddGeluFusion", ["A", "B"], ["C"], "AddGeluFusion_1", domain="com.microsoft"), + ], + "Gelu_Add_ShapeNotMatch", #name + [ # inputs + helper.make_tensor_value_info('A', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + helper.make_tensor_value_info('B', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), # Bias shape not matched for FastGelu + ], + [ # outputs + helper.make_tensor_value_info('C', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + ], + [ # initializers + ] +) + +model = helper.make_model(graph) +onnx.save(model, r'gelu_add_shape_not_match.onnx') + + +graph = helper.make_graph( + [ # nodes + # Add node before Gelu + helper.make_node("MatMul", ["A", "B"], ["C"], "MatMul_1"), + helper.make_node("AddGeluFusion", ["C", "D"], ["E"], "AddGeluFusion_1", domain="com.microsoft"), + ], + "MatMul_AddGeluFusion", #name + [ # inputs + helper.make_tensor_value_info('A', TensorProto.FLOAT, ['batch', 'seq_len', 'x']), + helper.make_tensor_value_info('B', TensorProto.FLOAT, [128, 3072]), + helper.make_tensor_value_info('D', TensorProto.FLOAT, [3072]), + ], + [ # outputs + helper.make_tensor_value_info('E', TensorProto.FLOAT, ['batch', 'seq_len', 3072]), + ], + [ # initializers + ] +) + +model = helper.make_model(graph) +onnx.save(model, r'gelu_add_matmul.onnx') +