diff --git a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc b/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc index c76781d28f..ccf9ea8aeb 100644 --- a/onnxruntime/core/optimizer/skip_layer_norm_fusion.cc +++ b/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> 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(&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(&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 diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index f5e0ec7ec5..caa5beb6bf 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/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 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(), TransformerLevel::Level2); + graph_transformation_mgr.Register(onnxruntime::make_unique(), 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& 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& 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 p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx new file mode 100644 index 0000000000..555c701239 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx differ