diff --git a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc index d4c93b8553..946be98776 100644 --- a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc @@ -172,7 +172,9 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le p_add2 = const_cast(&edges[1]->GetNode()); if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) && - CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) { + CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType()) && + graph.GetNodeOutputsInGraphOutputs(*p_add1).empty() && + graph.GetNodeOutputsInGraphOutputs(*p_add2).empty()) { matched_format = Format::Format1; } } @@ -188,7 +190,9 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le p_add2 = const_cast(&edges[1]->GetNode()); if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) && - CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType())) { + CheckSecondAdd(graph, *p_add2, ln_node.GetExecutionProviderType()) && + graph.GetNodeOutputsInGraphOutputs(*p_add1).empty() && + graph.GetNodeOutputsInGraphOutputs(*p_add2).empty()) { matched_format = Format::Format2; } } @@ -202,7 +206,8 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le if (graph_utils::FindPath(ln_node, true, format3_parent_path, edges, logger)) { p_add1 = const_cast(&edges[0]->GetNode()); - if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType())) { + if (CheckFirstAdd(*p_add1, ln_node.GetExecutionProviderType()) && + graph.GetNodeOutputsInGraphOutputs(*p_add1).empty()) { matched_format = Format::Format3; } } @@ -237,7 +242,7 @@ Status SkipLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le skip_layer_norm_input_defs, ln_node.MutableOutputDefs(), {}, kMSDomain); - // Get attribute "epsilon" from "LayerNormalization" node if available. Else, default value + // Get attribute "epsilon" from "LayerNormalization" node if available. Else, default value // will be used. NodeAttributes ln_attrs = ln_node.GetAttributes(); NodeAttributes::const_iterator epsilon = ln_attrs.find("epsilon"); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index d7f0fbeac2..4d67444fd8 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1199,7 +1199,7 @@ TEST_F(GraphTransformationTests, ReshapeFusionGraphInputsTest) { ASSERT_EQ(op_to_count["Concat"], 1); ASSERT_EQ(op_to_count["Reshape"], 1); } - + TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerDoesntApplyTest) { auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx"; std::shared_ptr p_model; @@ -1914,9 +1914,14 @@ TEST_F(GraphTransformationTests, SkipLayerNormFusionTest) { TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format1.onnx", 0, 0, 1, logger_.get()); TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format2.onnx", 0, 0, 1, logger_.get()); TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format3.onnx", 0, 0, 1, logger_.get()); + TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format1_partial.onnx", 1, 0, 1, logger_.get()); TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format2_partial.onnx", 1, 0, 1, logger_.get()); TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format3_no_fusion.onnx", 1, 1, 0, logger_.get()); + + TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format1_graph_output.onnx", 1, 0, 1, logger_.get()); + TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format2_graph_output.onnx", 1, 0, 1, logger_.get()); + TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format3_graph_output.onnx", 1, 1, 0, logger_.get()); } TEST_F(GraphTransformationTests, SkipLayerNormFusion_Input_Output_Check) { diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format1_graph_output.onnx b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format1_graph_output.onnx new file mode 100644 index 0000000000..973baaf3dc Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format1_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format2_graph_output.onnx b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format2_graph_output.onnx new file mode 100644 index 0000000000..eea1f9a362 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format2_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format3_graph_output.onnx b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format3_graph_output.onnx new file mode 100644 index 0000000000..a6164fa63f Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_format3_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_gen.py b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_gen.py index ea66b56374..c88e3903e4 100644 --- a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_gen.py @@ -10,7 +10,7 @@ class Format(Enum): Format3 = 3 -def GenerateModel(format, model_name, multi_output_add=False): +def GenerateModel(format, model_name, multi_output_add=False, add_output_in_graph_output=False): nodes = [ # LayerNorm subgraph helper.make_node("ReduceMean", ["ln_in"], ["rd1_out"], "reduce1", axes=[-1], keepdims=1), helper.make_node("Sub", ["ln_in", "rd1_out"], ["sb1_out"], "sub1"), @@ -67,6 +67,10 @@ def GenerateModel(format, model_name, multi_output_add=False): helper.make_tensor_value_info('C', TensorProto.FLOAT, [16, 32, 4]), ], initializers) + + if add_output_in_graph_output: + extra_output = "ln_in" if format is Format.Format3 else "add3_out" + graph.output.extend([helper.make_tensor_value_info(extra_output, TensorProto.FLOAT, [16, 32, 4])]) model = helper.make_model(graph) onnx.save(model, model_name) @@ -75,6 +79,10 @@ def GenerateModel(format, model_name, multi_output_add=False): GenerateModel(Format.Format1, 'skip_layer_norm_format1.onnx') GenerateModel(Format.Format2, 'skip_layer_norm_format2.onnx') GenerateModel(Format.Format3, 'skip_layer_norm_format3.onnx') -GenerateModel(Format.Format1, 'skip_layer_norm_format1_partial.onnx', True) -GenerateModel(Format.Format2, 'skip_layer_norm_format2_partial.onnx', True) -GenerateModel(Format.Format3, 'skip_layer_norm_format3_no_fusion.onnx', True) +GenerateModel(Format.Format1, 'skip_layer_norm_format1_partial.onnx', multi_output_add = True) +GenerateModel(Format.Format2, 'skip_layer_norm_format2_partial.onnx', multi_output_add = True) +GenerateModel(Format.Format3, 'skip_layer_norm_format3_no_fusion.onnx', multi_output_add = True) + +GenerateModel(Format.Format1, 'skip_layer_norm_format1_graph_output.onnx', add_output_in_graph_output = True) +GenerateModel(Format.Format2, 'skip_layer_norm_format2_graph_output.onnx', add_output_in_graph_output = True) +GenerateModel(Format.Format3, 'skip_layer_norm_format3_graph_output.onnx', add_output_in_graph_output = True) \ No newline at end of file