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.

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

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<std::string> 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<std::string>& 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

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<std::string>& 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

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<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
 
       std::unordered_set<std::string> cuda_execution_providers = {onnxruntime::kCudaExecutionProvider};
       transformers.emplace_back(onnxruntime::make_unique<AddGeluFusion>(cuda_execution_providers));
+      transformers.emplace_back(onnxruntime::make_unique<GeluApproximation>(cuda_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<SkipLayerNormFusion>(cuda_execution_providers));
 
 #endif
@@ -142,7 +144,6 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
         transformers.emplace_back(onnxruntime::make_unique<NchwcTransformer>());
       }
 #endif
-
     } break;
 
     default:

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");

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<int>(TransformerLevel::Level1); i < static_cast<int>(TransformerLevel::MaxTransformerLevel); i++) {
+  for (int i = static_cast<int>(TransformerLevel::Level1); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
     ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformer_mgr.ApplyTransformers(graph, static_cast<TransformerLevel>(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<uint32_t>(graph_optimization_level)));
 
-  if ((graph_optimization_level >= TransformerLevel::Level1) || !custom_list.empty()) {
-    add_transformers(TransformerLevel::Level1);
+  for (int i = static_cast<int>(TransformerLevel::Level1); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
+    TransformerLevel level = static_cast<TransformerLevel>(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) {

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<int>(TransformerLevel::Default); i < static_cast<int>(TransformerLevel::MaxTransformerLevel); i++) {
+  for (int i = static_cast<int>(TransformerLevel::Default); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
     SessionOptions so;
     so.session_logid = "InferenceSessionTests.TestL1AndL2Transformers";
     so.graph_optimization_level = static_cast<TransformerLevel>(i);

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<Model> 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<GeluApproximation>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Model> 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<GeluApproximation>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Model> 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<GeluApproximation>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Model> 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<GeluApproximation>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Model> p_model;

onnxruntime/test/testdata/transform/approximation/gelu.onnx

@@ -0,0 +1,13 @@
+:~
+#
+
A
CGelu_1"Gelu:
com.microsoftGelu_NoBiasZ$
+
A
+

+batch
+	seq_len
+€b$
+
C
+

+batch
+	seq_len
+€B

onnxruntime/test/testdata/transform/approximation/gelu_add_bias.onnx

@@ -0,0 +1,17 @@
+:¦

+8
+
A
+
B
CAddGeluFusion_1"
AddGeluFusion:
com.microsoftGelu_AddBiasZ$
+
A
+

+batch
+	seq_len
+€Z
+
B
+	

+€b$
+
C
+

+batch
+	seq_len
+€B

onnxruntime/test/testdata/transform/approximation/gelu_add_matmul.onnx

@@ -0,0 +1,25 @@
+:â

+
+
A
+
B
CMatMul_1"MatMul
+8
+
C
+
D
EAddGeluFusion_1"
AddGeluFusion:
com.microsoftMatMul_AddGeluFusionZ$
+
A
+

+batch
+	seq_len
+
xZ
+
B
+

+
+€

+€Z
+
D
+	

+€b$
+
E
+

+batch
+	seq_len
+€B

onnxruntime/test/testdata/transform/approximation/gelu_add_shape_not_match.onnx

@@ -0,0 +1,19 @@
+:Ä

+8
+
A
+
B
CAddGeluFusion_1"
AddGeluFusion:
com.microsoftGelu_Add_ShapeNotMatchZ$
+
A
+

+batch
+	seq_len
+€Z$
+
B
+

+batch
+	seq_len
+€b$
+
C
+

+batch
+	seq_len
+€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')
+