diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 0b0172792c..5d39939d30 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -40,26 +40,29 @@ Status ReshapeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, c } /** -Apply Reshape Fusion. The fowllowing are subgraphs before and after fusion: +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.) Before fusion: - [Sub-graph Root Node ] - | / \ - | Shape Shape - | | | (one or two int64[] constant initializers) - | Gather(indice=0) Gather(indice=1) a[] b[] (optional) - | \ / / / - | Unsqueeze Unsqueeze / / - | \ / ___________________/ / - | \ / / ____________________________/ - | \ / / / + [Sub-graph Root Node ] + | / \ + | Shape Shape + | | | + | Gather(indices=0) a[] Gather(indices=2) b[] + | \ / / / + | Unsqueeze / Unsqueeze / + | \ / ___________/ / + | \ / / _______________________/ + | \ / / / \ Concat \ / Reshape After fusion: - [Sub-graph Root Node] (Constant Initializers, b is optional) - \ [0, 0, a, b] + [Sub-graph Root Node] (Constant Initializer) + \ [0, a, 0, b] \ / Reshape */ @@ -77,84 +80,57 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L } auto concat_input_count = concat.InputArgCount().front(); - if (concat_input_count < 3 || concat_input_count > 4 || concat.GetOutputEdgesCount() > 1) { + if (concat.GetOutputEdgesCount() > 1) { return false; } - // path 1: [Root] --> Shape --> Gather(indices=0) --> Unsqueeze (axes=0) --> Concat [input 0] - std::vector parent_path{ - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}}; + 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 an initializer. + // We do not check whether the initializer is constant. + // Some model uses constant initializer and some does not. + // Here we assume that no one will override the initializer using graph input. + if (optimizer_utils::AppendTensorFromInitializer(graph, *(concat.InputDefs()[i]), shape_value)) { + continue; + } - std::vector edges; - if (!graph_utils::FindPath(concat, true, parent_path, edges, logger)) { - return false; - } + // 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}}; - const Node& unsqueeze_1 = edges[0]->GetNode(); - const Node& gather_1 = edges[1]->GetNode(); - const Node& shape_1 = edges[2]->GetNode(); - if (unsqueeze_1.GetOutputEdgesCount() != 1 || gather_1.GetOutputEdgesCount() != 1 || shape_1.GetOutputEdgesCount() != 1) { - return false; - } - - if (graph_utils::GetInputNode(shape_1, 0) != p_root) { - return false; - } - - std::vector axes; - if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze_1, "axes", axes) && axes.size() == 1 && axes[0] == 0)) { - return false; - } - - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_1.InputDefs()[1]), int64_t(0), false)) { - return false; - } - - // path 2: [Root] --> Shape --> Gather(indices=1) --> Unsqueeze (axes=0) --> Concat [input 1] - std::vector parent_path2 { - {0, 1, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}}; - - if (!graph_utils::FindPath(concat, true, parent_path2, edges, logger)) { - return false; - } - - const Node& unsqueeze_2 = edges[0]->GetNode(); - const Node& gather_2 = edges[1]->GetNode(); - const Node& shape_2 = edges[2]->GetNode(); - if (unsqueeze_2.GetOutputEdgesCount() != 1 || gather_2.GetOutputEdgesCount() != 1 || shape_2.GetOutputEdgesCount() != 1) { - return false; - } - - if (graph_utils::GetInputNode(shape_2, 0) != p_root) { - return false; - } - - if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze_2, "axes", axes) && axes.size() == 1 && axes[0] == 0)) { - return false; - } - - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_2.InputDefs()[1]), int64_t(1), false)) { - return false; - } - - // Compose the shape value input for reshape op. - std::vector shape_value = {0, 0}; - - // We do not check whether the initializer is constant. - // Some model uses constant initializer and some does not. - // Here we assume that no one will override the initializer using graph input. - if (!optimizer_utils::AppendTensorFromInitializer(graph, *(concat.InputDefs()[2]), shape_value)) { - return false; - } - - if (concat_input_count > 3) { - if (!optimizer_utils::AppendTensorFromInitializer(graph, *(concat.InputDefs()[3]), shape_value)) { + std::vector edges; + if (!graph_utils::FindPath(concat, true, parent_path, edges, logger)) { return false; } + + const Node& unsqueeze = edges[0]->GetNode(); + const Node& gather = edges[1]->GetNode(); + const Node& shape = edges[2]->GetNode(); + + if (graph_utils::GetInputNode(shape, 0) != p_root) { + 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(i), 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()}); } // Create an initializer with the same name as the concat node output, and replace the concat node @@ -174,18 +150,12 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L return false; } - // Remove nodes not used anymore. - std::vector nodes_to_remove{ - graph.GetNode(unsqueeze_1.Index()), - graph.GetNode(gather_1.Index()), - graph.GetNode(shape_1.Index()), - graph.GetNode(unsqueeze_2.Index()), - graph.GetNode(gather_2.Index()), - graph.GetNode(shape_2.Index())}; - - for (Node* node : nodes_to_remove) { - graph_utils::RemoveNodeOutputEdges(graph, *node); - graph.RemoveNode(node->Index()); + // 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()); + } } return true; diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index dcb870b759..f5e0ec7ec5 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -959,6 +959,85 @@ TEST_F(GraphTransformationTests, ReshapeFusionOneConstTest) { } } +// Test Reshape Fusion with an internal node being the output of the graph. +TEST_F(GraphTransformationTests, ReshapeFusionInternalNodeIsOutput) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_internal_node_is_graph_output.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::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Shape"], 1); + ASSERT_EQ(op_to_count["Gather"], 1); + 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], 0); + EXPECT_EQ(val[1], 0); + EXPECT_EQ(val[2], -1); + } + } +} + +// Test Reshape Fusion where some of the internal nodes are reused: +// A Shape is used in two Gather's, and the third Gather is the graph output. +TEST_F(GraphTransformationTests, ReshapeFusionInternalReuseTest) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_internal_nodes_reused.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::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Shape"], 1); + ASSERT_EQ(op_to_count["Gather"], 1); + 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(), 5); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 0); + EXPECT_EQ(val[1], 128); + EXPECT_EQ(val[2], 0); + EXPECT_EQ(val[3], 0); + EXPECT_EQ(val[4], -1); + } else if (node.OpType() == "Shape") { + EXPECT_EQ(node.Name(), "shape2"); + } else if (node.OpType() == "Gather") { + EXPECT_EQ(node.Name(), "gather3"); + } + } +} + #ifndef DISABLE_CONTRIB_OPS static void ValidateAttention(Graph& graph) { diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py new file mode 100644 index 0000000000..3f01ffea8f --- /dev/null +++ b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py @@ -0,0 +1,72 @@ +import onnx +from onnx import helper +from onnx import TensorProto + +def save_model(graph, file_name): + model = helper.make_model(graph) + onnx.checker.check_model(model) + onnx.save(model, file_name) + +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["SubgraphRoot"], ["shape1_out"], "shape1"), + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), + helper.make_node("Gather", ["shape1_out", "indices0"], ["gather0_out"], "gather0", axis=0), + helper.make_node("Gather", ["shape1_out", "indices2"], ["gather2_out"], "gather2", axis=0), + helper.make_node("Gather", ["shape2_out", "indices3"], ["gather3_out"], "gather3", axis=0), + helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), + helper.make_node("Unsqueeze", ["gather2_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + helper.make_node("Unsqueeze", ["gather3_out"], ["unsqueeze3_out"], "unsqueeze3", axes=[0]), + + helper.make_node("Concat", ["unsqueeze0_out", "a1", "unsqueeze2_out", "unsqueeze3_out", "a4"], ["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, ['unk_0', 256, 'unk_2', 'unk_3']), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, ['unk_1', 128, 'unk_2', 'unk_3', 'unk_4']), + helper.make_tensor_value_info('gather3_out', TensorProto.INT64, []), + ], + [ # initializers + helper.make_tensor('a1', TensorProto.INT64, [1], [128]), + helper.make_tensor('a4', TensorProto.INT64, [1], [-1]), + helper.make_tensor('indices0', TensorProto.INT64, [], [0]), + helper.make_tensor('indices2', TensorProto.INT64, [], [2]), + helper.make_tensor('indices3', TensorProto.INT64, [], [3]), + ] +) + +save_model(graph, 'reshape_fusion_internal_nodes_reused.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("Concat", ["unsqueeze0_out", "unsqueeze1_out", "a"], ["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('gather0_out', TensorProto.INT64, []), + ], + [ # initializers + helper.make_tensor('a', TensorProto.INT64, [1], [-1]), + helper.make_tensor('indices0', TensorProto.INT64, [], [0]), + helper.make_tensor('indices1', TensorProto.INT64, [], [1]), + ] +) + +save_model(graph, 'reshape_fusion_internal_node_is_graph_output.onnx') + diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_node_is_graph_output.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_node_is_graph_output.onnx new file mode 100644 index 0000000000..a5fefccc7b Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_node_is_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_nodes_reused.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_nodes_reused.onnx new file mode 100644 index 0000000000..7b09ce1e14 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_internal_nodes_reused.onnx differ