From c0e817ff161a558591180ab8d81412f39fa940d2 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Thu, 23 Apr 2020 14:18:59 -0700 Subject: [PATCH] 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. --- .../core/optimizer/skip_layer_norm_fusion.cc | 42 ++++++++++-------- .../test/optimizer/graph_transform_test.cc | 33 ++++++++++++++ .../skip_layer_norm_input_output_check.onnx | Bin 0 -> 1149 bytes 3 files changed, 56 insertions(+), 19 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/skip_layer_norm_input_output_check.onnx 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 0000000000000000000000000000000000000000..555c701239f9c9ac38a9ae9cd42e37d5b4b070b5 GIT binary patch literal 1149 zcma)5Z%k8H6z^?IdAmdCUUbUa$g(1pG{E}W(tq^cc5IDWWmJrAu$Z?3&8%X<^4iJ9 z7P26c!JIJ*YS9Sjvc;I}i+(}gF-&C(A=$F|Wkf$Xh=v#@A>o6%S#R5QZpoH#laq7a z`Q>-d@0{~4r{nWGy5#m)%U(evR#m(2a@c_LBhikyTqR<)#mpnSz}1E1x_FypHt4bu zUWbt#BO8w#Lf-4Gt&(WaFo$1yMnw)DIkTW83o?R}M=qhDA>0~o3D<=~QORn+&%R7+ zu{Lxl-04f@4jMuE_(jIU$i=Hgvyitl-YnS+8s@N>WkdnL(H`xTLs40(HgLI^-ISLI zwZ+3GPLBdwl}hCsL0Y^4vpblHjA;?_g6$tjc7ujF>||LB%Ub`X&an!0X%)7AsBr$T ziioLK4}? zna}j(z$o34Q_!WiB(ixkpF!B*(rvi}U1m$xOdVdIS$O5Xh+G*y)X^T3O(KW&s1}KV za+RvbyBf<;ENRN>j_Hhs!$rnsI?_VRS#vo5gC55dg7 zMRGNGlHNb9g|=899i1DXoYn}@l@a>LrKab1FDg^BU(o5vQhHWwqK#!fdU(rWW$XM7 zm@Ivl9OzY1=_fDz+V_N9awO;u^3N`e7v)qe1;Tc&Fk7Y<*lsE`EKDIE;O? z?9)+bE?*`OQe%pH;W$lB1(n8{W*C1Ifmlf+^<5X~o%um(>WGmt?+DzkNkZqs9y(R{ z95J3d0jeb<@E0ifyGxb$t)nzyZBk4{4Iui)uiWeJhwy|tYt_Z%eKfA)KN z=6c9^<1XcwZ(pU|vz;*Z+NkoCWe7G_4$$x|LN6441dT;QA?;ih}hh#e04L2u_PE#JghAWF+ v*!Iep^7cdtv>p3Rsl3z+O(lQ7c%H|44)Fz8&7r>xH>Y2m>%Wj%wf_01rN?X_ literal 0 HcmV?d00001