diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 6be9ac8a60..897dfba4a8 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -65,65 +65,146 @@ 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)) { - return false; - } - if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { - 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)) { + 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; } - for (size_t i = 0; i < edges.size(); i++) { - if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + 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; + } + // 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 (pg_edges[7]->GetNode().GetOutputEdgesCount() == 1) { + // In this case, the "Gather" node in "Path 1" must have 1 output edge. + if (!graph_utils::FindPath(expand_node, true, {{0, 1, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) { + return false; + } + 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)) { + 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; + 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; } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index ad108f64df..8bc0e8558f 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1354,6 +1354,35 @@ 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 0000000000..eddbf40b5b Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format4.onnx differ