Generalize reshape fusion (#3554)

* Generalize reshape fusion

* Allow arbitrary number of Concat arguments
* Apply fusion even when an output of an internal node is used elsewhere
* Fix a bug when an internal node's output is the subgraph output
* Simplify code

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<graph_utils::EdgeEndToMatch> parent_path{
-      {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
-      {0, 0, "Gather", {1, 11}, kOnnxDomain},
-      {0, 0, "Shape", {1}, kOnnxDomain}};
+  std::vector<int64_t> 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<std::pair<NodeType, NodeIndex>> 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<const Node::EdgeEnd*> 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<graph_utils::EdgeEndToMatch> 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<int64_t> 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<graph_utils::EdgeEndToMatch> 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<int64_t> 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<const Node::EdgeEnd*> 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<int64_t> 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<Node*> 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;

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<Model> 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<ReshapeFusion>(), TransformerLevel::Level1);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Initializer>(*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<int64_t>();
+      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<Model> 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<ReshapeFusion>(), TransformerLevel::Level1);
+  auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger());
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> 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<Initializer>(*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<int64_t>();
+      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) {

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')
+