diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index c546d42530..7ee51e8c10 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -4,6 +4,7 @@ #include "core/graph/graph_utils.h" #include "core/graph/graph.h" #include "core/common/logging/logging.h" +#include namespace onnxruntime { @@ -744,5 +745,41 @@ bool FindPath(const Node& node, bool is_input_edge, const std::vector q; + std::vector nodes_to_remove; + q.push(&start_node); + // From the current node, remove nodes bottom-up util it reaches a node with multiple outputs/graph output. + while (q.size() != 0) { + const Node& cur_node = *(q.front()); + q.pop(); + // Each eligible node in the subgraph must have less than one output edge and no output should be + // the graph output + if (cur_node.GetOutputEdgesCount() > 1 || !graph.GetNodeOutputsInGraphOutputs(cur_node).empty()) { + continue; + } + nodes_to_remove.push_back(cur_node.Index()); + // push the parents of current node to the queue. + for (unsigned int i = 0; i < cur_node.InputDefs().size(); ++i) { + const std::string& input_name = GetNodeInputName(cur_node, i); + if (IsInitializer(graph, input_name, true) || IsGraphInput(graph, cur_node.InputDefs()[i])) { + // skip initializers and graph inputs + continue; + } + q.push(GetInputNode(cur_node, i)); + } + } + if (nodes_to_remove.size() <= 0) { + // Nothing to remove + return false; + } + // Remove nodes that are not used anymore. + for (const auto& node_index : nodes_to_remove) { + Node* node = graph.GetNode(node_index); + RemoveNodeOutputEdges(graph, *node); + graph.RemoveNode(node->Index()); + } + return true; +} } // namespace graph_utils } // namespace onnxruntime diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index 77c3888bff..ec4a5388b2 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -244,5 +244,11 @@ struct EdgeEndToMatch { */ bool FindPath(const Node& node, bool is_input_edge, const std::vector& edges_to_match, std::vector& result, const logging::Logger& logger); +/** + * Remove nodes with only one output edge using bottom-up bfs traversal. + * @param node: The node to start with. + * @returns true if there is one or more node(s) removed by this function. Otherwise return false. +*/ +bool RemoveNodesWithOneOutputBottomUp(Graph& graph, const Node& node); } // namespace graph_utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 2183d4a00e..3b00e55e6e 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -39,18 +39,53 @@ Status ReshapeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, c return Status::OK(); } +/** + * Find the subgraph that matches [root] -> Shape -> Gather -> Unsqueeze + */ +bool ReshapeFusion::Fuse_Subgraph2(Graph& graph, const NodeArg& root_input, const Node& concat, + int index, std::vector shape_value, const logging::Logger& logger) { + std::vector parent_path{ + {0, index, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}}; + std::vector edges; + if (graph_utils::FindPath(concat, true, parent_path, edges, logger)) { + const Node& unsqueeze = edges[0]->GetNode(); + const Node& gather = edges[1]->GetNode(); + const Node& shape = edges[2]->GetNode(); + + const NodeArg& shape_input = *(shape.InputDefs()[0]); + if (shape_input.Name() != root_input.Name()) { + return false; + } + + std::vector axes; + if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze, "axes", axes) && axes.size() == 1 && axes[0] == 0)) { + return false; + } + + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) { + return false; + } + return true; + } + + return false; +} /** Apply Reshape Fusion. The following are subgraphs before and after fusion: (a[] and b[] are int64[] constant initializers; Concat may have any number of arguments, each of which is a constant initializer or a Shape->Gather->Unsqueeze chain with the -index corresponding to the index of the argument.) +index corresponding to the index of the argument, or a custom subgraph in which nodes +have only one output edge. Note the resulting shape value should contain no more than one +value of -1. Before fusion: [Sub-graph Root] | / \ | Shape Shape | | | - | Gather(indices=0) a[] Gather(indices=2) b[] + | Gather(indices=0) a[] Gather(indices=2) b[] or subgraph | \ / / / | Unsqueeze / Unsqueeze / | \ / ___________/ / @@ -85,51 +120,37 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L return false; } + // Loop through the inputs of concat node to calculate the shape_value for a potential reshape fusion. std::vector shape_value; shape_value.reserve(concat_input_count); - // Used to keep the following nodes in the order of their potential removal. - enum class NodeType { Unsqueeze, Gather, Shape }; - std::set> candidates_for_removal; for (int i = 0; i < concat_input_count; ++i) { // First check if the i-th argument is a constant initializer. if (optimizer_utils::AppendTensorFromInitializer(graph, *(concat.InputDefs()[i]), shape_value, true)) { continue; } - // Try to find path [Root] --> Shape --> Gather(indices=i) --> Unsqueeze (axes=0) --> Concat [input i] - std::vector parent_path{ - {0, i, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}}; - - std::vector edges; - if (!graph_utils::FindPath(concat, true, parent_path, edges, logger)) { + bool matched = ReshapeFusion::Fuse_Subgraph2(graph, root_input, concat, i, shape_value, logger); + const Node* p_cur_node = graph_utils::GetInputNode(concat, i); + if (matched) { + shape_value.push_back(0); + } else if (p_cur_node != nullptr) { + // This node could lead to a potential subgraph pattern fusion. Mark the shape value to -1. + shape_value.push_back(-1); + } else { return false; } + } - const Node& unsqueeze = edges[0]->GetNode(); - const Node& gather = edges[1]->GetNode(); - const Node& shape = edges[2]->GetNode(); - - const NodeArg& shape_input = *(shape.InputDefs()[0]); - if (shape_input.Name() != root_input.Name()) { - return false; + // Check how many -1 are there in shape_value. + int subgraph_cnt = 0; + for (auto it = shape_value.begin(); it < shape_value.end(); ++it) { + if ((*it) == -1) { + subgraph_cnt++; } - - std::vector axes; - if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze, "axes", axes) && axes.size() == 1 && axes[0] == 0)) { - return false; - } - - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) { - return false; - } - - shape_value.push_back(0); - - candidates_for_removal.insert({NodeType::Unsqueeze, unsqueeze.Index()}); - candidates_for_removal.insert({NodeType::Gather, gather.Index()}); - candidates_for_removal.insert({NodeType::Shape, shape.Index()}); + } + // If more than one "-1" value is present in shape_value, return false to exit current fusion. + if (subgraph_cnt > 1) { + return false; } // Create an initializer with the same name as the concat node output, and replace the concat node @@ -145,18 +166,18 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L shape_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); shape_initializer_proto.set_raw_data(shape_value.data(), shape_value.size() * sizeof(int64_t)); auto& new_node_arg = graph_utils::AddInitializer(graph, shape_initializer_proto); - if (!graph_utils::ReplaceNodeWithInitializer(graph, *graph.GetNode(concat.Index()), new_node_arg)) { - return false; - } - // Remove nodes that are not used anymore. - for (const auto& node_type_and_index : candidates_for_removal) { - Node* node = graph.GetNode(node_type_and_index.second); - if (node->GetOutputEdgesCount() == 0 && graph.GetNodeOutputsInGraphOutputs(*node).empty()) { - graph.RemoveNode(node->Index()); + // Safely remove concat parent nodes which have only one output + for (int i = 0; i < concat_input_count; ++i) { + const Node* p_cur_node = graph_utils::GetInputNode(concat, i); + if (p_cur_node != nullptr) { + graph_utils::RemoveNodesWithOneOutputBottomUp(graph, *p_cur_node); } } + if (!graph_utils::ReplaceNodeWithInitializer(graph, *graph.GetNode(concat.Index()), new_node_arg)) { + return false; + } return true; } diff --git a/onnxruntime/core/optimizer/reshape_fusion.h b/onnxruntime/core/optimizer/reshape_fusion.h index d906c556b1..aba7460800 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.h +++ b/onnxruntime/core/optimizer/reshape_fusion.h @@ -20,6 +20,7 @@ class ReshapeFusion : public GraphTransformer { private: static bool Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::Logger& logger); + static bool Fuse_Subgraph2(Graph& graph, const NodeArg& root_input, const Node& concat, int index, std::vector shape_value, const logging::Logger& looger); }; } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 4d67444fd8..af5c231c13 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1199,8 +1199,8 @@ TEST_F(GraphTransformationTests, ReshapeFusionGraphInputsTest) { ASSERT_EQ(op_to_count["Concat"], 1); ASSERT_EQ(op_to_count["Reshape"], 1); } - -TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerDoesntApplyTest) { + +TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerSubgraphTest) { auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx"; std::shared_ptr p_model; ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); @@ -1214,7 +1214,26 @@ TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerDoesntA // The optimization does not apply. std::map op_to_count = CountOpsInGraph(graph); - ASSERT_EQ(op_to_count_orig, op_to_count); + ASSERT_EQ(op_to_count["Shape"], 0); + ASSERT_EQ(op_to_count["Gather"], 0); + ASSERT_EQ(op_to_count["Unsqueeze"], 0); + ASSERT_EQ(op_to_count["Concat"], 0); + ASSERT_EQ(op_to_count["Reshape"], 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "Reshape") { + const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, node.InputDefs()[1]->Name()); + ASSERT_TRUE(tensor_proto != nullptr); + + auto initializer = onnxruntime::make_unique(*tensor_proto, graph.ModelPath()); + EXPECT_EQ(tensor_proto->data_type(), ONNX_NAMESPACE::TensorProto_DataType_INT64); + EXPECT_EQ(initializer->size(), 3); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 1); + EXPECT_EQ(val[1], 200); + EXPECT_EQ(val[2], -1); + } + } } TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerAppliesTest) { @@ -1256,7 +1275,6 @@ TEST_F(GraphTransformationTests, ReshapeFusionAnotherGraphInput) { std::shared_ptr p_model; ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); Graph& graph = p_model->MainGraph(); - std::map op_to_count_orig = CountOpsInGraph(graph); onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); @@ -1265,7 +1283,11 @@ TEST_F(GraphTransformationTests, ReshapeFusionAnotherGraphInput) { // The optimization does not apply. std::map op_to_count = CountOpsInGraph(graph); - ASSERT_EQ(op_to_count_orig, op_to_count); + ASSERT_EQ(op_to_count["Shape"], 0); + ASSERT_EQ(op_to_count["Gather"], 0); + ASSERT_EQ(op_to_count["Unsqueeze"], 0); + ASSERT_EQ(op_to_count["Concat"], 0); + ASSERT_EQ(op_to_count["Reshape"], 1); } TEST_F(GraphTransformationTests, ReshapeFusionOverridableInitializer) { @@ -1278,6 +1300,7 @@ TEST_F(GraphTransformationTests, ReshapeFusionOverridableInitializer) { onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + std::cout << "ret " << ret << std::endl; ASSERT_TRUE(ret.IsOK()); // The optimization does not apply. @@ -1285,6 +1308,79 @@ TEST_F(GraphTransformationTests, ReshapeFusionOverridableInitializer) { ASSERT_EQ(op_to_count_orig, op_to_count); } +TEST_F(GraphTransformationTests, ReshapeFusionConcatSubgraphMultipleOutputs) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_concat_subgraph_multiple_outputs.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + // The optimization applies but certain paths with multiple outputs/graph outputs are not removed. + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Shape"], 3); + ASSERT_EQ(op_to_count["Gather"], 1); + ASSERT_EQ(op_to_count["Unsqueeze"], 1); + ASSERT_EQ(op_to_count["Squeeze"], 1); + ASSERT_EQ(op_to_count["Div"], 1); + ASSERT_EQ(op_to_count["Concat"], 0); + ASSERT_EQ(op_to_count["Reshape"], 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "Reshape") { + const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, node.InputDefs()[1]->Name()); + ASSERT_TRUE(tensor_proto != nullptr); + + auto initializer = onnxruntime::make_unique(*tensor_proto, graph.ModelPath()); + EXPECT_EQ(tensor_proto->data_type(), ONNX_NAMESPACE::TensorProto_DataType_INT64); + EXPECT_EQ(initializer->size(), 3); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 0); + EXPECT_EQ(val[1], 0); + EXPECT_EQ(val[2], -1); + } + } +} + +TEST_F(GraphTransformationTests, ReshapeFusionConcatSubgraph) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_concat_subgraph.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Shape"], 0); + ASSERT_EQ(op_to_count["Gather"], 0); + ASSERT_EQ(op_to_count["Unsqueeze"], 0); + ASSERT_EQ(op_to_count["Squeeze"], 0); + ASSERT_EQ(op_to_count["Div"], 0); + ASSERT_EQ(op_to_count["Concat"], 0); + ASSERT_EQ(op_to_count["Reshape"], 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "Reshape") { + const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, node.InputDefs()[1]->Name()); + ASSERT_TRUE(tensor_proto != nullptr); + + auto initializer = onnxruntime::make_unique(*tensor_proto, graph.ModelPath()); + EXPECT_EQ(tensor_proto->data_type(), ONNX_NAMESPACE::TensorProto_DataType_INT64); + EXPECT_EQ(initializer->size(), 3); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 0); + EXPECT_EQ(val[1], 0); + EXPECT_EQ(val[2], -1); + } + } +} + TEST_F(GraphTransformationTests, ExpandElimination) { auto model_uri = MODEL_FOLDER "expand_elimination.onnx"; std::shared_ptr model; @@ -1303,6 +1399,8 @@ TEST_F(GraphTransformationTests, ExpandElimination) { ASSERT_TRUE(op_to_count["Expand"] == 3); } + + TEST_F(GraphTransformationTests, CastElimination) { auto model_uri = MODEL_FOLDER "cast_elimination.onnx"; std::shared_ptr model; diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph.onnx new file mode 100644 index 0000000000..1146a2d281 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph_multiple_outputs.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph_multiple_outputs.onnx new file mode 100644 index 0000000000..a59992a62d Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_concat_subgraph_multiple_outputs.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py index 2b3f3558d0..b32b1f2eac 100644 --- a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py @@ -187,3 +187,71 @@ graph = helper.make_graph( save_model(graph, 'reshape_fusion_with_graph_inputs.onnx') +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["SubgraphRoot"], ["shape0_out"], "shape0"), + helper.make_node("Shape", ["SubgraphRoot"], ["shape1_out"], "shape1"), + helper.make_node("Gather", ["shape0_out", "indices0"], ["gather0_out"], "gather0", axis=0), + helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), + helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), + helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), + + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), + helper.make_node("Squeeze", ["shape2_out"], ["squeeze_out"], "squeeze"), + helper.make_node("Div", ["squeeze_out", "div_init"], ["div_out"], "div"), + helper.make_node("Unsqueeze", ["div_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + + helper.make_node("Concat", ["unsqueeze0_out", "unsqueeze1_out", "unsqueeze2_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [10, 20, 30]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [10, 20, 'unk']), + ], + [ # initializers + helper.make_tensor('indices0', TensorProto.INT64, [], [0]), + helper.make_tensor('indices1', TensorProto.INT64, [], [1]), + helper.make_tensor('div_init', TensorProto.INT64, [], [1]), + ] +) + +save_model(graph, 'reshape_fusion_concat_subgraph.onnx') + +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["SubgraphRoot"], ["shape0_out"], "shape0"), + helper.make_node("Shape", ["SubgraphRoot"], ["shape1_out"], "shape1"), + helper.make_node("Gather", ["shape0_out", "indices0"], ["gather0_out"], "gather0", axis=0), + helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), + helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), + helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), + + helper.make_node("Shape", ["unsqueeze0_out"], ["dummy_out"], "dummy"), + + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), + helper.make_node("Squeeze", ["shape2_out"], ["squeeze_out"], "squeeze"), + helper.make_node("Div", ["squeeze_out", "div_init"], ["div_out"], "div"), + helper.make_node("Unsqueeze", ["div_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + + helper.make_node("Concat", ["unsqueeze0_out", "unsqueeze1_out", "unsqueeze2_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [10, 20, 30]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [10, 20, 'unk']), + helper.make_tensor_value_info('div_out', TensorProto.INT64, []), + ], + [ # initializers + helper.make_tensor('indices0', TensorProto.INT64, [], [0]), + helper.make_tensor('indices1', TensorProto.INT64, [], [1]), + helper.make_tensor('div_init', TensorProto.INT64, [], [1]), + ] +) + +save_model(graph, 'reshape_fusion_concat_subgraph_multiple_outputs.onnx') \ No newline at end of file