Fix a bug in skiplayernorm fusion pattern 2 (#3660)

For skiplayernorm fusion pattern 2, its input[0] should be equal to the input[0] of Add_1, but is overridden by the input[0] of Add_2.

onnxruntime/core/optimizer/skip_layer_norm_fusion.cc

@@ -57,13 +57,18 @@ static bool CheckFirstAdd(Node& add, ProviderType providertype) {
 // Add2 is the 2nd add of the to be fused sub-graph
 // The 1st input should be a 3D tensor
 // The 2nd input should be a 1D constant value
-static bool CheckSecondAdd(Node& add, ProviderType providertype) {
+static bool CheckSecondAdd(Graph& graph, Node& add, ProviderType providertype) {
   if (providertype != add.GetExecutionProviderType() ||
       !IsSupportedDataType(add) ||
       add.GetOutputEdgesCount() != 1) {
     return false;
   }
 
+  // The 2nd input should be a constant value
+  if (!graph_utils::NodeArgIsConstant(graph, *(add.MutableInputDefs()[1]))) {
+    return false;
+  }
+
   // Check the input dimensions of the "Add" node.
   const TensorShapeProto* add_input1_shape = add.MutableInputDefs()[0]->Shape();
   const TensorShapeProto* add_input2_shape = add.MutableInputDefs()[1]->Shape();
@@ -84,8 +89,8 @@ Skip Layer Normalization will fuse Add + LayerNormalization into one node, and a
 
 Before fusion:
 Format 1:
-      [Sub1]   [Sub2]
-         \       /
+    [Sub1]  C    [Sub2]
+        \  /     /
         Add2    /
            \   /
             Add1
@@ -93,10 +98,10 @@ Format 1:
      LayerNormalization
 
 Format 2:
-      [Sub1]   [Sub2]
-         \       /
-          \    Add2
-           \   /
+      [Sub1] [Sub2]  C
+         \      \   /
+          \     Add2
+           \    /
             Add1
              |
      LayerNormalization
@@ -131,10 +136,9 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le
   const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
   std::vector<std::reference_wrapper<Node>> nodes_to_remove;
   for (auto node_index : node_topology_list) {
-    nodes_to_remove.clear();
     Node* p_layernorm = graph.GetNode(node_index);
     if (p_layernorm == nullptr)
-      continue;  // we removed the node as part of an earlier fusion.
+      continue;  // node was removed in an earlier fusion.
 
     Node& ln_node = *p_layernorm;
     ORT_RETURN_IF_ERROR(Recurse(ln_node, modified, graph_level, logger));
@@ -167,7 +171,7 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le
       p_add2 = const_cast<Node*>(&edges[1]->GetNode());
 
       if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) &&
-          CheckSecondAdd(*p_add2, ln_node.GetExecutionProviderType())) {
+          CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) {
         matched_format = Format::Format1;
       }
     }
@@ -183,7 +187,7 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le
         p_add2 = const_cast<Node*>(&edges[1]->GetNode());
 
         if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) &&
-            CheckSecondAdd(*p_add2, ln_node.GetExecutionProviderType())) {
+            CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) {
           matched_format = Format::Format2;
         }
       }
@@ -230,18 +234,18 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le
                                                "SkipLayerNormalization",
                                                "fused SkipLayerNorm subgraphs ",
                                                skip_layer_norm_input_defs,
-                                               {}, {}, kMSDomain);
+                                               ln_node.MutableOutputDefs(), {}, kMSDomain);
 
     // Assign provider to this new node. Provider should be same as the provider for old node.
     skip_layer_norm_node.SetExecutionProviderType(ln_node.GetExecutionProviderType());
-
-    // move input edges to add (first in list) across to the layer_norm_node.
-    // move output definitions and output edges from mul_node (last in list) to layer_norm_node.
-    // remove all the other nodes.
-    graph_utils::FinalizeNodeFusion(graph, nodes_to_remove, skip_layer_norm_node);
-
-    modified = true;
   }
+  for (const auto& node : nodes_to_remove) {
+    graph_utils::RemoveNodeOutputEdges(graph, node);
+    graph.RemoveNode(node.get().Index());
+  }
+
+  modified = true;
+
   return Status::OK();
 }
 }  // namespace onnxruntime

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -1609,6 +1609,39 @@ TEST_F(GraphTransformationTests, SkipLayerNormFusionTest) {
   TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format3_no_fusion.onnx", 1, 1, 0, logger_.get());
 }
 
+TEST_F(GraphTransformationTests, SkipLayerNormFusion_Input_Output_Check) {
+  auto model_uri = MODEL_FOLDER "fusion/skip_layer_norm_input_output_check.onnx";
+  std::shared_ptr<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(onnxruntime::make_unique<LayerNormFusion>(), TransformerLevel::Level2);
+  graph_transformation_mgr.Register(onnxruntime::make_unique<SkipLayerNormFusion>(), TransformerLevel::Level2);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_);
+  ASSERT_TRUE(ret.IsOK());
+
+  for (Node& node : graph.Nodes()) {
+    if (node.OpType() == "SkipLayerNormalization") {
+      // check inputs
+      std::vector<NodeArg*>& input_defs = node.MutableInputDefs();
+      EXPECT_EQ(input_defs.size(), 5u) << "SkipLayerNormalization number of inputs does not equal to 5. Got:" << node.InputDefs().size();
+      EXPECT_EQ(input_defs[0]->Name(), "input.1");
+      EXPECT_EQ(input_defs[1]->Name(), "6");
+      EXPECT_EQ(input_defs[2]->Name(), "1");
+      EXPECT_EQ(input_defs[3]->Name(), "2");
+      EXPECT_EQ(input_defs[4]->Name(), "4");
+
+      // check outputs
+      std::vector<NodeArg*>& output_defs = node.MutableOutputDefs();
+      EXPECT_EQ(node.OutputDefs().size(), 1u) << "SkipLayerNormalization number of outputs does not equal to 1. Got:" << node.OutputDefs().size();
+      EXPECT_EQ(output_defs[0]->Name(), "19");
+    } else {
+      EXPECT_EQ(node.OpType(), "MatMul") << "Unexpected node: " << node.OpType() << "," << node.Name();
+    }
+  }
+}
+
 TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat1) {
   auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format1.onnx";
   std::shared_ptr<Model> p_model;