From 21c2365f1db40556f8d36c7033aa456dabc42a6d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 9 Dec 2019 14:56:40 -0800 Subject: [PATCH 1/6] Support opset11 in reshape fusion --- onnxruntime/core/optimizer/reshape_fusion.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 0b54e46745073..a53dd3ca14592 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -84,7 +84,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L // path 1: [Root] --> Shape --> Gather(indices=0) --> Unsqueeze (axes=0) --> Concat [input 0] std::vector parent_path{ {0, 0, "Unsqueeze", {1}, kOnnxDomain}, - {0, 0, "Gather", {1}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, {0, 0, "Shape", {1}, kOnnxDomain}}; std::vector edges; @@ -115,7 +115,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L // path 2: [Root] --> Shape --> Gather(indices=1) --> Unsqueeze (axes=0) --> Concat [input 1] std::vector parent_path2 { {0, 1, "Unsqueeze", {1}, kOnnxDomain}, - {0, 0, "Gather", {1}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, {0, 0, "Shape", {1}, kOnnxDomain}}; if (!graph_utils::FindPath(concat, true, parent_path2, edges, logger)) { From fa100d0cae76d12a2c3f98d9ab925d49fb8f40e0 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 9 Dec 2019 17:08:05 -0800 Subject: [PATCH 2/6] add opset 11 for all ops --- onnxruntime/core/optimizer/reshape_fusion.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index a53dd3ca14592..0b0172792cd17 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -72,7 +72,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L } const Node& concat = *p_concat; - if (!graph_utils::IsSupportedOptypeVersionAndDomain(concat, "Concat", {1, 4})) { + if (!graph_utils::IsSupportedOptypeVersionAndDomain(concat, "Concat", {1, 4, 11})) { return false; } @@ -83,7 +83,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L // path 1: [Root] --> Shape --> Gather(indices=0) --> Unsqueeze (axes=0) --> Concat [input 0] std::vector parent_path{ - {0, 0, "Unsqueeze", {1}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, {0, 0, "Gather", {1, 11}, kOnnxDomain}, {0, 0, "Shape", {1}, kOnnxDomain}}; @@ -114,7 +114,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L // path 2: [Root] --> Shape --> Gather(indices=1) --> Unsqueeze (axes=0) --> Concat [input 1] std::vector parent_path2 { - {0, 1, "Unsqueeze", {1}, kOnnxDomain}, + {0, 1, "Unsqueeze", {1, 11}, kOnnxDomain}, {0, 0, "Gather", {1, 11}, kOnnxDomain}, {0, 0, "Shape", {1}, kOnnxDomain}}; From 9db442402e043d528ac416c0a561e17978479985 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 10 Dec 2019 10:37:18 -0800 Subject: [PATCH 3/6] Support opset 11 Squad model --- .../core/optimizer/embed_layer_norm_fusion.cc | 282 +++++++++++++----- .../test/optimizer/graph_transform_test.cc | 26 ++ .../fusion/embed_layer_norm_format3.onnx | Bin 0 -> 2239 bytes .../transform/fusion/embed_layer_norm_gen.py | 75 +++++ 4 files changed, 307 insertions(+), 76 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/embed_layer_norm_gen.py diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 970de353f5392..b018e196aff39 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -65,6 +65,194 @@ static bool CheckInput(NodeArg* input, const logging::Logger& logger) { return true; } +static bool MatchPositionEmbeddingSubgraph1( + Graph& graph, + Node& position_gather_node, + NodeArg* input_ids, + const logging::Logger& logger, + std::vector& matched_edges) { + // Match two paths. + // Match Shape --> Expand path if needed. + std::vector position_parent_nodes; + std::vector position_embedding_path_symbolic{ + {0, 1, "Expand", {8}, kOnnxDomain}, + {0, 1, "Shape", {1}, kOnnxDomain}}; + std::vector edges; + if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + return false; + } + if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + auto p_expand_node = graph.GetNode(edges[0]->GetNode().Index()); + auto p_shape_node = graph.GetNode(edges[1]->GetNode().Index()); + // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand + Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); + std::vector pg_parent_path{ + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Cast", {9}, kOnnxDomain}, + {0, 0, "Squeeze", {1}, kOnnxDomain}, + {0, 0, "Transpose", {1}, kOnnxDomain}, + {0, 0, "NonZero", {9}, kOnnxDomain}, + {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + matched_edges = edges; + + if (!graph_utils::FindPath(expand_node, true, pg_parent_path, edges, logger)) { + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + } + // Check if the second input of the Gather node in the path has a constant input of 1 + Node& gather_node = *graph.GetNode(edges[edges.size() - 2]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { + DEBUG_LOG("Second input of Gather should be a constant with value 1. "); + return false; + } + + // Check if the parent of "shape" is the input_ids + Node& shape_node_2 = *graph.GetNode(edges[edges.size() - 1]->GetNode().Index()); + if (shape_node_1.MutableInputDefs()[0] != input_ids || + shape_node_1.MutableInputDefs()[0] != input_ids) { + return false; + } + + matched_edges.insert(matched_edges.end(), edges.begin(), edges.end()); + return true; +} + +/** Match subgraph like the following: + (input_ids) + / \ + Shape Shape + | | + Gather (indice=0) Gather (indice=1)--+ + | | | + Unsqueeze Unsqueeze Cast + \ / | + \ / Range(start=0, delta=1) + \ / | + Concat Unsqueeze + | | + +--|----------------------------+ + | | + Expand + | + Gather + + Note that position gather node is the node in the bottom of above sub-graph. +*/ + +static bool MatchPositionEmbeddingSubgraph2( + Graph& graph, + Node& position_gather_node, + NodeArg* input_ids, + const logging::Logger& logger, + std::vector& matched_edges) { + + // Match Gather <-- Expand <-- Unsqueeze <-- Range <-- Cast <-- Gather <-- Shape + std::vector position_parent_nodes; + std::vector position_embedding_path_symbolic{ + {0, 1, "Expand", {8}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Range", {11}, kOnnxDomain}, + {0, 1, "Cast", {9}, kOnnxDomain}, + {0, 0, "Gather", {11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + std::vector edges; + if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + DEBUG_LOG("Failed to find path 1."); + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != (i == 4 ? 2 : 1)) { + DEBUG_LOG("Output edge count not expected for nodes in path 1."); + return false; + } + } + matched_edges = edges; + + Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& range_node = *graph.GetNode(edges[2]->GetNode().Index()); + Node& gather_node_1 = *graph.GetNode(edges[4]->GetNode().Index()); + Node& shape_node_1 = *graph.GetNode(edges[5]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[0]), int64_t(0), true)) { + DEBUG_LOG("The first input of Range should be a constant with value 0."); + return false; + } + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[2]), int64_t(1), true)) { + DEBUG_LOG("The third input of Range should be a constant with value 1."); + return false; + } + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_1.InputDefs()[1]), int64_t(1), true)) { + DEBUG_LOG("The second input of Gather in path1 should be a constant with value 1."); + return false; + } + + std::vector expand_parent_path1{ + {0, 1, "Concat", {11}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + if (!graph_utils::FindPath(expand_node, true, expand_parent_path1, edges, logger)) { + DEBUG_LOG("Failed to find path 2."); + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + DEBUG_LOG("Output edge count not expected for nodes in path 2."); + return false; + } + } + + Node& concat_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& gather_node_0 = *graph.GetNode(edges[2]->GetNode().Index()); + Node& shape_node_0 = *graph.GetNode(edges[3]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_0.InputDefs()[1]), int64_t(0), true)) { + DEBUG_LOG("Second input of Gather in path2 should be a constant with value 0."); + return false; + } + matched_edges.insert(matched_edges.end(), edges.begin(), edges.end()); + + std::vector concat_parent_path{ + {0, 1, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain} + }; + if (!graph_utils::FindPath(concat_node, true, concat_parent_path, edges, logger)) { + DEBUG_LOG("Failed to find path 3."); + return false; + } + // Two paths share the gather node (with second input indices==1) + if (edges[1]->GetNode().Index() != gather_node_1.Index()) { + DEBUG_LOG(" Gather nodes in path 1 and 3 expected to be same node."); + return false; + } + if (edges[0]->GetNode().GetOutputEdgesCount() != 1) { + DEBUG_LOG("Output edge count not expected for nodes in path 3."); + return false; + } + + // Check if the two paths of position gather lead to the same input. + if (shape_node_0.MutableInputDefs()[0] != input_ids || + shape_node_1.MutableInputDefs()[0] != input_ids) { + DEBUG_LOG("Two shape nodes are expected to be input_ids."); + return false; + } + + // Do not add the gather node since it has been added in another path. + matched_edges.push_back(edges[0]); + return true; +} + /** Embed Layer Normalization will fuse embeddings and mask processing into one node : The embeddings before conversion: @@ -176,22 +364,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l continue; } + NodeArg* input_ids = word_gather_node.MutableInputDefs()[1]; + // Check the second input of position gather. If it's not initializer, check for two paths. - Node* p_expand_node = nullptr; - Node* p_shape_node = nullptr; std::vector pg_edges; - bool isValidEmbedSubNode = true; if (graph_utils::IsConstantInitializer(graph, position_gather_node.MutableInputDefs()[1]->Name())) { - // Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0. + // Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0. std::vector data; auto expected_shape = word_gather_node.MutableInputDefs()[1]->Shape(); - if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data) - || !utils::HasDimValue(expected_shape->dim()[0]) - || !utils::HasDimValue(expected_shape->dim()[1]) - || static_cast(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) { + if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data) || !utils::HasDimValue(expected_shape->dim()[0]) || !utils::HasDimValue(expected_shape->dim()[1]) || static_cast(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) { continue; } int64_t expected_value = 0; + bool isValidEmbedSubNode = true; for (size_t i = 0; i < data.size(); i++) { if (data[i] != expected_value) { isValidEmbedSubNode = false; @@ -202,68 +387,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l expected_value = 0; } } - } else { - // Match two paths. - // Match Shape --> Expand path if needed. - std::vector position_parent_nodes; - std::vector position_embedding_path_symbolic{ - {0, 1, "Expand", {8}, kOnnxDomain}, - {0, 1, "Shape", {1}, kOnnxDomain}}; - if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + if (!isValidEmbedSubNode) { continue; } - if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { - continue; - } - p_expand_node = graph.GetNode(edges[0]->GetNode().Index()); - p_shape_node = graph.GetNode(edges[1]->GetNode().Index()); - // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand - Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); - Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); - std::vector pg_parent_path{ - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Cast", {9}, kOnnxDomain}, - {0, 0, "Squeeze", {1}, kOnnxDomain}, - {0, 0, "Transpose", {1}, kOnnxDomain}, - {0, 0, "NonZero", {9}, kOnnxDomain}, - {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}, - }; - if (!graph_utils::FindPath(expand_node, true, pg_parent_path, pg_edges, logger)) { - continue; - } - for (size_t i = 0; i < pg_edges.size(); i++) { - if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) { - isValidEmbedSubNode = false; - break; + } else { + if (!MatchPositionEmbeddingSubgraph1(graph, position_gather_node, input_ids, logger, pg_edges)) { + pg_edges.clear(); + if (!MatchPositionEmbeddingSubgraph2(graph, position_gather_node, input_ids, logger, pg_edges)) { + continue; } } - // Check if the second input of the Gather node in the path has a constant input of 1 - Node& gather_node = *graph.GetNode(pg_edges[pg_edges.size() - 2]->GetNode().Index()); - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { - DEBUG_LOG("Second input of Gather should be a constant with value 1. "); - - continue; - } - // Check if the two paths of position gather lead to the same input. - Node& shape_node_2 = *graph.GetNode(pg_edges[pg_edges.size() - 1]->GetNode().Index()); - if (shape_node_1.MutableInputDefs()[0] != shape_node_2.MutableInputDefs()[0]) { - continue; - } - // Check if the parent of "shape" is the parent of "word gather" - if (shape_node_1.MutableInputDefs()[0] != word_gather_node.MutableInputDefs()[1]) { - continue; - } - - } - if (!isValidEmbedSubNode) { - continue; } // Get input "input_ids" from node. - NodeArg* input_ids = word_gather_node.MutableInputDefs()[1]; if (!CheckInput(input_ids, logger)) { DEBUG_LOG("Input id is not valid. "); continue; @@ -283,32 +419,30 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l continue; } - if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != - utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) { + if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != + utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) { DEBUG_LOG("Input_ids and segment id should have the same shape. "); continue; } - if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != - utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) { + if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != + utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) { DEBUG_LOG("Input_ids and mask should have the same shape. "); continue; } NodeArg* gamma = layer_norm_node.MutableInputDefs()[1]; NodeArg* beta = layer_norm_node.MutableInputDefs()[2]; - if (gamma->Shape() == nullptr - || gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { + if (gamma->Shape() == nullptr || gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { DEBUG_LOG("Gamma should be of shape (hidden_size). "); continue; } - if (beta->Shape() == nullptr - || beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { + if (beta->Shape() == nullptr || beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { DEBUG_LOG("Beta should be of shape (hidden_size). "); continue; } - // Cast input_ids, segment_ids, and mask to int32 if needed. + // Cast input_ids, segment_ids, and mask to int32 if needed. input_ids = CastToInt32(graph, input_ids, layer_norm_node.GetExecutionProviderType()); segment_ids = CastToInt32(graph, segment_ids, layer_norm_node.GetExecutionProviderType()); mask = CastToInt32(graph, mask, layer_norm_node.GetExecutionProviderType()); @@ -339,10 +473,6 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l for (size_t i = 0; i < pg_edges.size(); i++) { nodes_to_remove.push_back(pg_edges[i]->GetNode().Index()); } - if (p_shape_node != nullptr && p_expand_node != nullptr) { - nodes_to_remove.push_back(p_shape_node->Index()); - nodes_to_remove.push_back(p_expand_node->Index()); - } nodes_to_remove.push_back(word_gather_node.Index()); nodes_to_remove.push_back(position_gather_node.Index()); nodes_to_remove.push_back(segment_gather_node.Index()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index bd7665fdede74..ad108f64df52b 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1328,6 +1328,32 @@ TEST(GraphTransformationTests, EmbedLayerNormFusionFormat2) { ASSERT_TRUE(op_to_count["SkipLayerNormalization"] == 0); ASSERT_TRUE(op_to_count["EmbedLayerNormalization"] == 1); } + +TEST(GraphTransformationTests, EmbedLayerNormFusionFormat3) { + auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format3.onnx"; + std::shared_ptr 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(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["Shape"], 0); + EXPECT_EQ(op_to_count["Expand"], 0); + EXPECT_EQ(op_to_count["Gather"], 0); + EXPECT_EQ(op_to_count["Unsqueeze"], 0); + EXPECT_EQ(op_to_count["LayerNormalization"], 0); + EXPECT_EQ(op_to_count["SkipLayerNormalization"], 0); + EXPECT_EQ(op_to_count["ReduceSum"], 0); + EXPECT_EQ(op_to_count["MatMul"], 1); + EXPECT_EQ(op_to_count["Add"], 2); + EXPECT_EQ(op_to_count["Cast"], 3); + EXPECT_EQ(op_to_count["Attention"], 1); + EXPECT_EQ(op_to_count["EmbedLayerNormalization"], 1); +} #endif } // namespace test diff --git a/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx new file mode 100644 index 0000000000000000000000000000000000000000..e58ffa62a142f9652f9e6f396bb174cac9d8f9de GIT binary patch literal 2239 zcmb_e%}(1u5RQ`&+mjN)wzQ}x>PkgeDk1?yMay3httvqeP*jPFwZuze5!)eN7s`o4 zUxJ=^fS&sXy}>iJ`;*uk18NUMu{HC}Oy-;2nNf_=@~?4t1*R8tW9oQrgdl3;E~zn)33>HB#@)*}^lE|4uTF+f+`dtj`3d4Gqp)jTlYXgR|tUklqZHB_GFqHl&!`i(JKSHK1MVG${ z%`j+UDrrmEDs2i(7VAZmtk^uv4$e!)8BH9?s3MoF;wF3BCgf>2sdN8D&?Aa_)R}|C z7j&0&D0XVd?BSqAMAi`@S|ht0fi*}?O3CtE(~|4=mi$I3Q69KbA_8*O#epl!m9kaZ z5*WM%vl~VZ@ehd$iX*CUSoV(ygj7#1YX#sm?797DrE|vwtVEX;APIjDSwF~RMO6`p zUDx61`5-(}5v#D_y6_Inh_w1F%1rc0ASbB8afJq?dctA}ZY6y#xWsX*Yo=$(}m$Op?q@uwSOcL7+m3p0hwe-$5f71 zfeDtmK`FaLehyVIgV=Z41iR5>o)KRjZHE4=?=^d26dqAnfI%!Ti28A-pAJD^6$@0# zR%r**owx%tpjgF43SNu{Yc?KGV1m5p*DQeqGqs7nkW*sif0s&j=pXY@WvL@v*0+>&^A`%FXrm>W|CIRVJ@CvSVfM zKUj5YtJ+id Date: Tue, 10 Dec 2019 11:21:11 -0800 Subject: [PATCH 4/6] Fix build warning. --- onnxruntime/core/optimizer/embed_layer_norm_fusion.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index b018e196aff39..6be9ac8a60dfd 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -84,8 +84,7 @@ static bool MatchPositionEmbeddingSubgraph1( if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { return false; } - auto p_expand_node = graph.GetNode(edges[0]->GetNode().Index()); - auto p_shape_node = graph.GetNode(edges[1]->GetNode().Index()); + // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); @@ -120,7 +119,7 @@ static bool MatchPositionEmbeddingSubgraph1( // Check if the parent of "shape" is the input_ids Node& shape_node_2 = *graph.GetNode(edges[edges.size() - 1]->GetNode().Index()); if (shape_node_1.MutableInputDefs()[0] != input_ids || - shape_node_1.MutableInputDefs()[0] != input_ids) { + shape_node_2.MutableInputDefs()[0] != input_ids) { return false; } @@ -244,7 +243,7 @@ static bool MatchPositionEmbeddingSubgraph2( // Check if the two paths of position gather lead to the same input. if (shape_node_0.MutableInputDefs()[0] != input_ids || shape_node_1.MutableInputDefs()[0] != input_ids) { - DEBUG_LOG("Two shape nodes are expected to be input_ids."); + DEBUG_LOG("The parent of two shape nodes are expected to be input_ids."); return false; } From edf2082b4c12c1db7168d5e273c9075364879e6f Mon Sep 17 00:00:00 2001 From: CeciliaLiu Date: Tue, 10 Dec 2019 13:51:53 -0800 Subject: [PATCH 5/6] embed formatt 4 --- .../core/optimizer/embed_layer_norm_fusion.cc | 153 +++++++++++++----- .../test/optimizer/graph_transform_test.cc | 28 ++++ .../fusion/embed_layer_norm_format4.onnx | Bin 0 -> 2162 bytes 3 files changed, 145 insertions(+), 36 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format4.onnx diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 6be9ac8a60dfd..71173104a1e73 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -65,65 +65,147 @@ static bool CheckInput(NodeArg* input, const logging::Logger& logger) { return true; } +/** Match subgraph like the following: + (input_ids) + / \ + Shape Shape + | | + ^Gather (indice=0)^ Gather (indice=1)--+ + ^|^ ^|^ | + ^Unsqueeze^ ^Unsqueeze^ Unsqueeze + ^\^ ^/^ | + ^\^ ^/^ ConstantOfShape + ^\^ ^/^ | + ^Concat^ NonZero + | | + | Transpose + | | + | Squeeze + | | + | Cast + | | + | Unsqueeze + +--|----------------------------+ + | | + Expand + | + Gather + + Note that position gather node is the node in the bottom of above sub-graph. + Paths in ^^ are alternative path to be matched if path input_ids -> Shape -> Expand -> Gather is not found. +*/ static bool MatchPositionEmbeddingSubgraph1( Graph& graph, Node& position_gather_node, NodeArg* input_ids, const logging::Logger& logger, std::vector& matched_edges) { - // Match two paths. - // Match Shape --> Expand path if needed. - std::vector position_parent_nodes; - std::vector position_embedding_path_symbolic{ - {0, 1, "Expand", {8}, kOnnxDomain}, - {0, 1, "Shape", {1}, kOnnxDomain}}; - std::vector edges; - if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + + if (position_gather_node.OpType() != "Gather") { return false; } - if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { + std::vector pg_edges; + // Find the "Expand" node + if (!graph_utils::FindPath(position_gather_node, true, {{0, 1, "Expand", {8}, kOnnxDomain}}, pg_edges, logger)) { return false; } - - // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand - Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); - Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); - std::vector pg_parent_path{ - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Cast", {9}, kOnnxDomain}, - {0, 0, "Squeeze", {1}, kOnnxDomain}, - {0, 0, "Transpose", {1}, kOnnxDomain}, - {0, 0, "NonZero", {9}, kOnnxDomain}, - {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}, - }; - matched_edges = edges; - - if (!graph_utils::FindPath(expand_node, true, pg_parent_path, edges, logger)) { + Node& expand_node = *graph.GetNode(pg_edges[0]->GetNode().Index()); + const Node::EdgeEnd* expand_edge = pg_edges[0]; + // Look for Path 1: + // Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand + if (!graph_utils::FindPath(expand_node, true, + {{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Cast", {9}, kOnnxDomain}, + {0, 0, "Squeeze", {1}, kOnnxDomain}, + {0, 0, "Transpose", {1}, kOnnxDomain}, + {0, 0, "NonZero", {9}, kOnnxDomain}, + {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}}, + pg_edges, logger)) { return false; } - for (size_t i = 0; i < edges.size(); i++) { - if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + // All nodes in Path 1 except the "Gather" node must have only 1 output edge. + for (size_t i = 0; i < pg_edges.size() - 2; i++) { + if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) { return false; } } + if (pg_edges[8]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } // Check if the second input of the Gather node in the path has a constant input of 1 - Node& gather_node = *graph.GetNode(edges[edges.size() - 2]->GetNode().Index()); + Node& gather_node = *graph.GetNode(pg_edges[pg_edges.size() - 2]->GetNode().Index()); if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { DEBUG_LOG("Second input of Gather should be a constant with value 1. "); return false; } - // Check if the parent of "shape" is the input_ids - Node& shape_node_2 = *graph.GetNode(edges[edges.size() - 1]->GetNode().Index()); - if (shape_node_1.MutableInputDefs()[0] != input_ids || - shape_node_2.MutableInputDefs()[0] != input_ids) { + // Match Shape --> Expand path if needed. + std::vector pg_edges_2; + Node* p_shape_node_2 = nullptr; + const Node::EdgeEnd* unsqueeze_edge = nullptr; + if (graph_utils::FindPath(expand_node, true, {{0, 1, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) { + p_shape_node_2 = graph.GetNode(pg_edges_2[0]->GetNode().Index()); + // In this case, the "Gather" node in "Path 1" must have 1 output edge. + if (pg_edges[7]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + } else if (graph_utils::FindPath(expand_node, true, + {{0, 1, "Concat", {4}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}}, + pg_edges_2, logger)) { + p_shape_node_2 = graph.GetNode(pg_edges_2[3]->GetNode().Index()); + // In this case, the "Gather" node in "Path 1" must have 2 output edges. + if (pg_edges[7]->GetNode().GetOutputEdgesCount() != 2) { + return false; + } + // Check for Unsqueeze --> Concat + Node& concat_node = *graph.GetNode(pg_edges_2[0]->GetNode().Index()); + std::vector pg_edges_3; + if (!graph_utils::FindPath(concat_node, true, + {{0, 1, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}}, + pg_edges_3, logger)) { + return false; + } + unsqueeze_edge = pg_edges_3[0]; + if (pg_edges_3[0]->GetNode().GetOutputEdgesCount() != 1 || pg_edges_3[1]->GetNode().GetOutputEdgesCount() != 2) { + return false; + } + // The gather node must be the same gather node in path 1. + if (graph.GetNode(pg_edges_3[1]->GetNode().Index()) != graph.GetNode(pg_edges[7]->GetNode().Index())) { + return false; + } + } else { + return false; + } + for (size_t i = 0; i < pg_edges_2.size(); i++) { + if (pg_edges_2[i]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + } + // Check if the two paths of position gather lead to the same input. + Node& shape_node_1 = *graph.GetNode(pg_edges[pg_edges.size() - 1]->GetNode().Index()); + Node& shape_node_2 = *graph.GetNode(p_shape_node_2->Index()); + if (shape_node_1.MutableInputDefs()[0] != shape_node_2.MutableInputDefs()[0]) { + return false; + } + // Check if the parent of "shape" is the parent of "word gather" + if (shape_node_1.MutableInputDefs()[0] != input_ids) { return false; } - matched_edges.insert(matched_edges.end(), edges.begin(), edges.end()); + // Add all the valid nodes to result. + matched_edges = pg_edges; + matched_edges.push_back(expand_edge); + matched_edges.insert(matched_edges.end(), pg_edges_2.begin(), pg_edges_2.end()); + if (unsqueeze_edge != nullptr) { + matched_edges.push_back(unsqueeze_edge); + } return true; } @@ -251,7 +333,6 @@ static bool MatchPositionEmbeddingSubgraph2( matched_edges.push_back(edges[0]); return true; } - /** Embed Layer Normalization will fuse embeddings and mask processing into one node : The embeddings before conversion: diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index ad108f64df52b..0e7f1b201b3f4 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1354,6 +1354,34 @@ TEST(GraphTransformationTests, EmbedLayerNormFusionFormat3) { EXPECT_EQ(op_to_count["Attention"], 1); EXPECT_EQ(op_to_count["EmbedLayerNormalization"], 1); } + +TEST(GraphTransformationTests, EmbedLayerNormFusionFormat4) { + auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format4.onnx"; + std::shared_ptr 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(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Shape"] == 0); + ASSERT_TRUE(op_to_count["Expand"] == 0); + ASSERT_TRUE(op_to_count["Gather"] == 0); + ASSERT_TRUE(op_to_count["Concat"] == 0); + ASSERT_TRUE(op_to_count["Unsqueeze"] == 0); + ASSERT_TRUE(op_to_count["ConstantOfShape"] == 0); + ASSERT_TRUE(op_to_count["NonZero"] == 0); + ASSERT_TRUE(op_to_count["Transpose"] == 0); + ASSERT_TRUE(op_to_count["Squeeze"] == 0); + ASSERT_TRUE(op_to_count["Add"] == 0); + ASSERT_TRUE(op_to_count["ReduceSum"] == 0); + ASSERT_TRUE(op_to_count["Attention"] == 1); + ASSERT_TRUE(op_to_count["SkipLayerNormalization"] == 0); + ASSERT_TRUE(op_to_count["EmbedLayerNormalization"] == 1); +} #endif } // namespace test diff --git a/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format4.onnx b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format4.onnx new file mode 100644 index 0000000000000000000000000000000000000000..eddbf40b5bd5e03901e2319edfd08fc994ab8ddb GIT binary patch literal 2162 zcmZ`)O>7fK6!tofy?%BQXF?+r`Xd>7SwSL#D@vm51g8jr+JFcMuqwG5?>OGX-ZkDe zKWVGN1thri2B(%@dO}52xpC+LaiGVF8%Vh!S``FBAXVz2vp;XWyS5M3dcODn-n?(b zVKQN~tvyy-(aoCCZYq{$=;qwx9PEVT1*WalEfCI-(KW@Yv-aMK(Xm<`tB^eGoFQ)u zvu0gsF%%82*eBR-OGX^YY(!LqFb)~ldqvYV3quvr2t&<*mjp^kuA`ghI*e`lvD2pz zkus1_E;4g5c3VinJoM<2$to|@Qgq9xh4LR;^6+Cdx01auK$VQ@hQ7(#2Bw-?HuN$- z!ybr;NY3eS%P_mSp0s;zBk(Eoh{dgwU8F8<^QF4($^Rj6<8pp`j!Yavv%W3h5jB6Yu7WP+^eq zQRpgTqyU8WWQ1RXs^WJn$*0ecZkIJP&#zA$4g?IbQ4lsHBkW>J(N(M|dB{1!jK4); z1?(8Q^EYZ1JBGW4XVfGaMO9H)7&@q`!Bstrj9>BgafT~ znxWe*je)qqSWDHK-XN|C(2We+4q6t|ZNjRY+c;nG*ouTvJ6-6FNs=y;O{IPHLiewf zgS5``t20G4^I3ie=J>dq(@m|~HcX>tkvwazGPTmMYp>|M_=?M(`=o86APPC5kch=% zizSakC6Y@@f=~)Foit%%u~KN0q;WJ}Tr9aQmvRBfX}@$`5QoGx57PMU*E;?2%LlSi z`+&yo?5C%TWqSF{*YxSBf92`y3;Fe(irOyAyov+(wx@8zGj4$aCc}-l&hXbDqTQ00~kKrf7NyG6NeIGmtw+KH39zGE!ayR_;K0MX+;}W({#Jwgy c7xaoE9e6Lwey+tw{T9t`chbosUW(`7e`Rc}&;S4c literal 0 HcmV?d00001 From 2767ddd341c7b1f0a9e3721eaad80e0b61e5c26b Mon Sep 17 00:00:00 2001 From: CeciliaLiu Date: Tue, 10 Dec 2019 16:50:07 -0800 Subject: [PATCH 6/6] pr --- .../core/optimizer/embed_layer_norm_fusion.cc | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 27cd7ef0e8425..897dfba4a8eda 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -101,9 +101,6 @@ static bool MatchPositionEmbeddingSubgraph1( const logging::Logger& logger, std::vector& matched_edges) { - if (position_gather_node.OpType() != "Gather") { - return false; - } std::vector pg_edges; // Find the "Expand" node if (!graph_utils::FindPath(position_gather_node, true, {{0, 1, "Expand", {8}, kOnnxDomain}}, pg_edges, logger)) { @@ -147,23 +144,23 @@ static bool MatchPositionEmbeddingSubgraph1( std::vector pg_edges_2; Node* p_shape_node_2 = nullptr; const Node::EdgeEnd* unsqueeze_edge = nullptr; - if (graph_utils::FindPath(expand_node, true, {{0, 1, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) { - p_shape_node_2 = graph.GetNode(pg_edges_2[0]->GetNode().Index()); + if (pg_edges[7]->GetNode().GetOutputEdgesCount() == 1) { // In this case, the "Gather" node in "Path 1" must have 1 output edge. - if (pg_edges[7]->GetNode().GetOutputEdgesCount() != 1) { + if (!graph_utils::FindPath(expand_node, true, {{0, 1, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) { return false; - } - } else if (graph_utils::FindPath(expand_node, true, - {{0, 1, "Concat", {4}, kOnnxDomain}, + } + p_shape_node_2 = graph.GetNode(pg_edges_2[0]->GetNode().Index()); + } else if (pg_edges[7]->GetNode().GetOutputEdgesCount() == 2) { + // In this case, the "Gather" node in "Path 1" must have 2 output edges. + if (!graph_utils::FindPath(expand_node, true, + {{0, 1, "Concat", {4, 11}, kOnnxDomain}, {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, {0, 0, "Gather", {1, 11}, kOnnxDomain}, {0, 0, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) { - p_shape_node_2 = graph.GetNode(pg_edges_2[3]->GetNode().Index()); - // In this case, the "Gather" node in "Path 1" must have 2 output edges. - if (pg_edges[7]->GetNode().GetOutputEdgesCount() != 2) { return false; } + p_shape_node_2 = graph.GetNode(pg_edges_2[3]->GetNode().Index()); // Check for Unsqueeze --> Concat Node& concat_node = *graph.GetNode(pg_edges_2[0]->GetNode().Index()); std::vector pg_edges_3;