Update Gelu Fusion to support new graph pattern from PyTorch 1.4 (#3148)

* update GeluFusion to support pattern from PyTorch 1.4; 
* Fix a bug that missing the check of an edge between mul2 and root.
* update script to fuse gelu from PyTorch 1.4
* Add test for python optimizer

onnxruntime/core/optimizer/gelu_fusion.cc

@@ -24,7 +24,25 @@ static bool IsSupportedDataType(const Node& node) {
   }
   return true;
 }
+/*
+     This function fuses subgraph like the following into one Gelu node.
+     Subgraph pattern 1:
+                   +-------Mul(0.5)---------------------+
+                   |                                    |
+                   |                                    v
+                [root] --> Div -----> Erf  --> Add --> Mul ==>
+                          (B=1.4142...)        (1)
 
+      Subgraph pattern 2:
+                   +------------------------------------+
+                   |                                    |
+                   |                                    v
+                [root] --> Div -----> Erf  --> Add --> Mul -->Mul ==>
+                          (B=1.4142...)        (1)            (0.5)
+
+       After Fusion:
+                [root]--> Gelu ==>
+*/
 Status GeluFusion::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();
@@ -68,13 +86,9 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
       continue;
     }
 
-    // Check the other input node(e.g. not of type Erf) is 1.0f.
-    const Node& add_first_input_node = *(add_node.InputNodesBegin());
-    int add_const_input_index = 0;
-    if (add_first_input_node.OpType().compare("Erf") == 0) {
-      add_const_input_index = 1;
-    }
-    const auto& add_const_input_arg = add_node.InputDefs()[add_const_input_index];
+    // Check the other input node (e.g. not the Erf) is 1.0f.
+    bool is_erf_first_input = (add_node.InputDefs()[0]->Name() == erf_node.MutableOutputDefs()[0]->Name());
+    const auto& add_const_input_arg = add_node.InputDefs()[is_erf_first_input ? 1 : 0];
     if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *add_const_input_arg, 1.0f, true)) {
       continue;
     }
@@ -87,35 +101,60 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
       continue;
     }
 
-    const Node* p_mul2_node = nullptr;
-    for (auto iter = mul_node.InputNodesBegin(); iter != mul_node.InputNodesEnd(); ++iter) {
-      if ((*iter).OpType().compare("Mul") == 0) {
-        // find the other input node of Mul
-        p_mul2_node = &(*iter);
-        break;
+    bool is_pattern_1 = true;
+    const Node* p_mul2_node = graph_utils::FirstParentByType(mul_node, "Mul");
+    if (p_mul2_node != nullptr) {
+      // Match subgraph pattern 1
+      Node& mul2_node = *graph.GetNode(p_mul2_node->Index());
+      if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul2_node, "Mul", {7}) ||
+          mul2_node.GetExecutionProviderType() != div.GetExecutionProviderType() ||
+          mul2_node.GetOutputEdgesCount() != 1 ||
+          !IsSupportedDataType(mul2_node)) {
+        continue;
       }
-    }
-    if (p_mul2_node == nullptr) {
-      continue;
-    }
 
-    Node& mul2_node = *graph.GetNode(p_mul2_node->Index());
-    if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul2_node, "Mul", {7}) ||
-        mul2_node.GetExecutionProviderType() != div.GetExecutionProviderType() ||
-        mul2_node.GetOutputEdgesCount() != 1 ||
-        !IsSupportedDataType(mul2_node)) {
-      continue;
-    }
+      // One input of mul2_node shall be the subgraph input
+      auto root_index = optimizer_utils::IndexOfNodeInput(*p_mul2_node, *div.InputDefs()[0]);
+      if (root_index < 0)
+        continue;
 
-    // Check the other input node(e.g. not of type Add) is 0.5f.
-    int mul_const_input_index = 0;
-    if (mul2_node.InputDefs()[0]->Name() == div.MutableInputDefs()[0]->Name()) {
-      mul_const_input_index = 1;
-    }
+      // Check the other input node is 0.5f.
+      int mul_const_input_index = (root_index == 0 ? 1 : 0);
 
-    const auto& mul_const_input_arg = mul2_node.InputDefs()[mul_const_input_index];
-    if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *mul_const_input_arg, 0.5f, true)) {
-      continue;
+      const auto& mul_const_input_arg = mul2_node.InputDefs()[mul_const_input_index];
+      if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *mul_const_input_arg, 0.5f, true)) {
+        continue;
+      }
+    } else {
+      is_pattern_1 = false;
+
+      // Match subgraph pattern 2
+      if (mul_node.GetOutputEdgesCount() != 1) {
+        continue;
+      }
+
+      // Another input of Mul node shall be the subgraph input.
+      auto root_index = optimizer_utils::IndexOfNodeInput(mul_node, *div.InputDefs()[0]);
+      if (root_index < 0)
+        continue;
+
+      Node& mul2_node = *graph.GetNode(mul_node.OutputNodesBegin()->Index());
+      if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul2_node, "Mul", {7}) ||
+          mul_node.GetExecutionProviderType() != div.GetExecutionProviderType() ||
+          !IsSupportedDataType(mul_node)) {
+        continue;
+      }
+
+      int mul_const_input_index = 0;
+      if (mul2_node.InputDefs()[0]->Name() == mul_node.MutableOutputDefs()[0]->Name()) {
+        mul_const_input_index = 1;
+      }
+      const auto& mul_const_input_arg = mul2_node.InputDefs()[mul_const_input_index];
+      if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *mul_const_input_arg, 0.5f, true)) {
+        continue;
+      }
+
+      p_mul2_node = &mul2_node;
     }
 
     const std::vector<NodeArg*> gelu_input_defs{div.MutableInputDefs()[0]};
@@ -131,7 +170,12 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
     // move input edges to div (first in list) across to the gelu_node.
     // move output definitions and output edges from mul_node (last in list) to gelu_node.
     // remove all the other nodes.
-    graph_utils::FinalizeNodeFusion(graph, {div, erf_node, add_node, mul2_node, mul_node}, gelu_node);
+    Node& mul2 = *graph.GetNode(p_mul2_node->Index());
+    if (is_pattern_1) {
+      graph_utils::FinalizeNodeFusion(graph, {div, erf_node, add_node, mul2, mul_node}, gelu_node);
+    } else {
+      graph_utils::FinalizeNodeFusion(graph, {div, erf_node, add_node, mul_node, mul2}, gelu_node);
+    }
 
     modified = true;
   }

onnxruntime/python/tools/bert/BertOnnxModel.py

@@ -234,15 +234,24 @@ class BertOnnxModel(OnnxModel):
 
     """
      Fuse Gelu with Erf into one node:
-                   +-------Mul(B=0.5)-------------------+
+     Pattern 1:
+                   +-------Mul(0.5)---------------------+
                    |                                    |
                    |                                    v
                 [root] --> Div -----> Erf  --> Add --> Mul -->
-                          (B=1.4142...)       (B=1)
+                          (B=1.4142...)       (1)
+
+      Pattern 2:
+                   +------------------------------------+
+                   |                                    |
+                   |                                    v
+                [root] --> Div -----> Erf  --> Add --> Mul -->Mul -->
+                          (B=1.4142...)       (1)            (0.5)
 
      Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
     """
     def fuse_gelu_with_elf(self, gelu_op_name):
+        logger.debug(f"start fuse_gelu_with_elf({gelu_op_name})")
         input_name_to_nodes = self.input_name_to_nodes()
         output_name_to_node = self.output_name_to_node()
 
@@ -276,25 +285,38 @@ class BertOnnxModel(OnnxModel):
             if self.find_constant_input(div, 1.4142, delta=0.001) != 1:
                 continue
 
-            root_node = self.get_parent(div, 0, output_name_to_node)
-            if root_node is None:
-                continue
+            subgraph_input = div.input[0]
 
-            mul_half = self.match_parent(mul_after_erf, 'Mul', None, output_name_to_node)
-            if mul_half is None:
-                continue
+            another = 1 if mul_after_erf.input[0] == add_after_erf.output[0] else 0
+            if subgraph_input == mul_after_erf.input[another]: # pattern 2
+                children = input_name_to_nodes[mul_after_erf.output[0]]
+                if len(children) != 1 or children[0].op_type != 'Mul':
+                    continue
+                mul_half = children[0]
+                if not self.has_constant_input(mul_half, 0.5):
+                    continue
+                subgraph_output = mul_half.output[0]
+            else: # pattern 1
+                mul_half = self.match_parent(mul_after_erf, 'Mul', another, output_name_to_node)
+                if mul_half is None:
+                    continue
             
-            if not self.has_constant_input(mul_half, 0.5):
-                continue
+                if not self.has_constant_input(mul_half, 0.5):
+                    continue
+
+                if subgraph_input not in mul_half.input:
+                    continue
+
+                subgraph_output = mul_after_erf.output[0]
 
             subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul_half]
-            if not self.is_safe_to_fuse_nodes(subgraph_nodes, [mul_after_erf.output[0]], input_name_to_nodes, output_name_to_node):
+            if not self.is_safe_to_fuse_nodes(subgraph_nodes, [subgraph_output], input_name_to_nodes, output_name_to_node):
                 continue
 
             nodes_to_remove.extend(subgraph_nodes)
             gelu_node = onnx.helper.make_node(gelu_op_name,
-                inputs=[root_node.output[0]],
-                outputs=[mul_after_erf.output[0]])
+                inputs=[subgraph_input],
+                outputs=[subgraph_output])
             gelu_node.domain = "com.microsoft"
             nodes_to_add.append(gelu_node)
 

onnxruntime/python/tools/bert/test_bert_optimization.py

@@ -22,6 +22,7 @@ from OnnxModel import OnnxModel
 BERT_TEST_MODELS = {
     "bert_pytorch_0": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_0.onnx',
     "bert_pytorch_1": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_1.onnx',
+    "bert_squad_pytorch1.4_opset10_fp32": 'test_data\\bert_squad_pytorch1.4_opset10_fp32\\BertForQuestionAnswering.onnx',
     "bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx'
 }
 
@@ -155,6 +156,13 @@ class TestBertOptimization(unittest.TestCase):
             }
         self.verify_node_count(bert_model, expected_node_count)
 
+    def test_pytorch_model_2_cpu(self):
+        input = BERT_TEST_MODELS['bert_squad_pytorch1.4_opset10_fp32']
+        bert_model = optimize_model(input, 'bert', gpu_only=False,
+                                    num_heads=2, hidden_size=8, sequence_length=10,
+                                    input_int32=False, float16=False)
+        self.assertTrue(bert_model.is_fully_optimized())
+
     def test_keras_model_1_cpu(self):
         input = BERT_TEST_MODELS['bert_keras_0']
 

onnxruntime/python/tools/bert/test_data/bert_squad_pytorch1.4_opset10_fp32/test_data_set_0/output_0.pb

@@ -0,0 +1,3 @@
+

+
BstartJ(÷ù¹A2I<Wö1<nR—<Gß;º^<Õq?;–‰‡<
+r¦;“ò–<

onnxruntime/python/tools/bert/test_data/bert_squad_pytorch1.4_opset10_fp32/test_data_set_0/output_1.pb

@@ -0,0 +1,2 @@
+

+
BendJ(îÿ”=FÚž=Lª‹=QR—=±wž=6\Â=¶<>=®‚¤=ùð•=Jg<4A>=

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -1131,6 +1131,63 @@ TEST(GraphTransformationTests, GeluFusionTest) {
   ASSERT_TRUE(op_to_count["Gelu"] == 1);
 }
 
+TEST(GraphTransformationTests, GeluFusionTestSwitchOrderFormat2) {
+  auto model_uri = MODEL_FOLDER "fusion/gelu_format2_0.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<GeluFusion>(), 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);
+  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"] == 1);
+}
+
+TEST(GraphTransformationTests, GeluFusionTestFormat2) {
+  auto model_uri = MODEL_FOLDER "fusion/gelu_format2_1.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<GeluFusion>(), 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);
+  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"] == 1);
+}
+
+TEST(GraphTransformationTests, GeluFusionTestFormat2GraphInput) {
+  auto model_uri = MODEL_FOLDER "fusion/gelu_format2_1_use_graph_input.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<GeluFusion>(), 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);
+  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"] == 1);
+}
+
 TEST(GraphTransformationTests, BiasGeluTest) {
   auto model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion.onnx";
   std::shared_ptr<Model> p_model;

onnxruntime/test/testdata/transform/fusion/gelu_gen.py

@@ -0,0 +1,149 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+"""
+Generate test model for Gelu subgraph pattern 2:
+                   +------------------------------------+
+                   |                                    |
+                   |                                    v
+                [root] --> Div -----> Erf  --> Add --> Mul -->Mul ==>
+                          (B=1.4142...)        (1)            (0.5)
+"""
+
+has_bias = True # change it to True to generate gelu_format2_*_with_bias.onnx
+gelu_use_graph_input = False # change it to False to let Gelu don't have graph inputs as inputs.
+switch_order = True # switch order of inputs for Mul and Add
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 64])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "seqlen", 64])
+
+value = (0.01 * np.arange(64)).astype(np.float32).reshape((64))
+bias_initializer = numpy_helper.from_array(value, "input_bias")
+
+value = np.asarray([1.4142099618911743]).astype(np.float32).reshape(())
+initializer_sqrt_2 = numpy_helper.from_array(value, "mul1_init")
+
+value = np.asarray([0.5]).astype(np.float32).reshape(())
+initializer_0_5 = numpy_helper.from_array(value, "mul3_init")
+
+value = np.asarray([1.0]).astype(np.float32).reshape(())
+initializer_1 = numpy_helper.from_array(value, "add1_init")
+
+
+nodes = []
+gelu_input = "input"
+if not gelu_use_graph_input:
+    leading_identity = helper.make_node(
+        'Identity', 
+        [gelu_input], 
+        ['identity_leading'], 
+        name="identity_leading"
+    )
+    gelu_input = "identity_leading"
+    nodes.append(leading_identity)
+
+gelu_root = gelu_input
+if has_bias:
+    add0 = helper.make_node(
+        'Add', 
+        [gelu_input, bias_initializer.name] if switch_order else [bias_initializer.name, gelu_input],
+        ['add0'], 
+        name="add0_node"
+    )
+    gelu_root = "add0"
+    nodes.append(add0)
+
+
+div = helper.make_node(
+    'Div', 
+    [gelu_root, initializer_sqrt_2.name], 
+    ['div'], 
+    name="div_node"
+)
+nodes.append(div)
+
+erf = helper.make_node(
+    'Erf', 
+    ['div'], 
+    ['erf'], 
+    name="erf_node"
+)
+nodes.append(erf)
+
+add1 = helper.make_node(
+    'Add', 
+    ['erf', initializer_1.name] if switch_order else [initializer_1.name, 'erf'],
+    ['add1'], 
+    name="add1"
+)
+nodes.append(add1)
+
+mul = helper.make_node(
+    'Mul', 
+    [gelu_root, 'add1'] if switch_order else ['add1', gelu_root],
+    ['mul'], 
+    name="mul_node"
+)
+nodes.append(mul)
+
+mul2 = helper.make_node(
+    'Mul', 
+    ['mul', initializer_0_5.name] if switch_order else [initializer_0_5.name, 'mul'],
+    ['mul2'], 
+    name="mul2_node"
+)
+nodes.append(mul2)
+
+ending_identity = helper.make_node(
+    'Identity', 
+    ['mul2'], 
+    ['output'], 
+    name="identity_ending"
+)
+nodes.append(ending_identity)
+
+initializers = []
+if has_bias:
+    initializers = [bias_initializer]
+
+initializers.extend([initializer_sqrt_2, initializer_1, initializer_0_5])
+
+# Create the graph (GraphProto)
+graph_def = helper.make_graph(
+    nodes,
+    'gelu_pattern_2',
+    [X],
+    [Y],
+    initializers
+)
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 10
+onnxdomain.domain = "" # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs={}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+file_name = "gelu_format2_0" if switch_order else "gelu_format2_1"
+if has_bias:
+    file_name += "_with_bias"
+
+if gelu_use_graph_input:
+    file_name += "_use_graph_input"
+
+file_name += ".onnx"
+onnx.save(model_def, file_name)
+print(file_name)
+