match reshape fusion for distilbert (#4844)

* reshape fusion for distilbert

* Update reshape_fusion.cc

* Update reshape_fusion.cc

* fix reshape

* resolve comments

* Update reshape_fusion.cc

* review comments

* review comments

* rename

* Update reshape_fusion.cc

onnxruntime/core/optimizer/reshape_fusion.cc

@@ -39,9 +39,99 @@ Status ReshapeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, c
   return Status::OK();
 }
 
+
+/**
+Provide check for Reshape Fusion for DistilBert. The following are subgraphs that
+match the pattern for DistilBert
+
+DistilBert reshape pattern:
+             [Root]
+             /   \ _ _ _ _ _ _ _
+         Shape                  \
+           |                  MatMul(w * w)
+    Gather(indices=0)             |
+            \                     |
+        Unsqueeze               Add(w)
+              \                  /
+             Concat     _ _ _ _ /
+                \      /
+                Reshape
+
+                                  - -> Shape
+                                  |
+Check the subgraph that matches [root] -> MatMul(w * w) -> Add(w) -> Reshape.
+*/
+static bool Match_Linear_Subgraph_1(Graph& graph, const Node& concat, const Node& root, const logging::Logger& logger) {
+  if (!optimizer_utils::CheckOutputEdges(graph, concat, 1)) {
+    return false;
+  }
+  auto reshape_itr = concat.OutputNodesBegin();
+  if ((*reshape_itr).OpType().compare("Reshape") != 0) {
+    return false;
+  }
+  const Node& reshape = *reshape_itr;
+
+  std::vector<graph_utils::EdgeEndToMatch> linear_path{
+    {0, 0, "Add", {7}, kOnnxDomain},
+    {0, 0, "MatMul", {1, 9}, kOnnxDomain}};
+  std::vector<const Node::EdgeEnd*> edges;
+  if (!graph_utils::FindPath(reshape, true, linear_path, edges, logger)) {
+    return false;
+  }
+
+  const Node& linear_path_add = edges[0]->GetNode();
+  const Node& linear_path_matmul = edges[1]->GetNode();
+
+  const Node* p_node_before_matmul = graph_utils::GetInputNode(linear_path_matmul, 0);
+  if (p_node_before_matmul != nullptr && p_node_before_matmul->Index() != root.Index()) {
+    return false;
+  }
+
+  if (linear_path_add.InputDefs().size() < 2) {
+    return false;
+  }
+  const NodeArg& linear_path_add_b = *(linear_path_add.InputDefs()[1]);
+  if (!graph_utils::IsInitializer(graph, linear_path_add_b.Name(), true)) {
+    return false;
+  }
+
+  if (!optimizer_utils::IsShapeKnownOnAllDims(linear_path_add_b, 1)) {
+    return false;
+  }
+  int64_t hidden_size = linear_path_add_b.Shape()->dim(0).dim_value();
+
+  if (!optimizer_utils::ValidateShape(*(linear_path_matmul.InputDefs()[1]), {hidden_size, hidden_size})) {
+    return false;
+  }
+
+  return true;
+}
+
+static bool Match_Shape(Graph& graph, const Node& concat, const Node& shape, const NodeArg& root_input, const logging::Logger& logger) {
+  const NodeArg& shape_input = *(shape.InputDefs()[0]);
+  if (shape_input.Name() == root_input.Name()) {
+    return true;
+  }
+
+  const ONNX_NAMESPACE::TensorShapeProto* shape_input_shape = shape_input.Shape();
+  const ONNX_NAMESPACE::TensorShapeProto* root_input_shape = root_input.Shape();
+
+  if (shape_input_shape != nullptr && root_input_shape != nullptr)
+    return optimizer_utils::CompareShape(*shape_input_shape, *root_input_shape);
+
+  const Node* p_node_before_shape = graph_utils::GetInputNode(shape, 0);
+  if (p_node_before_shape == nullptr) {
+    return false;
+  }
+  if (Match_Linear_Subgraph_1(graph, concat, *p_node_before_shape, logger)) {
+    return true;
+  }
+  return false;
+}
+
 /**
  * Find the subgraph that matches [root] -> Shape -> Gather -> Unsqueeze.
- * If checkOneElementOnly is set to true, this function only checks if the matched subgraph produces a 
+ * If checkOneElementOnly is set to true, this function only checks if the matched subgraph produces a
  * one element output(skip the Gather input indices check).
  */
 bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const NodeArg& root_input, const Node& concat,
@@ -57,11 +147,6 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node
     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<int64_t> axes;
     if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze, "axes", axes) && axes.size() == 1 && axes[0] == 0)) {
       return false;
@@ -74,6 +159,10 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node
     if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) {
       return false;
     }
+
+    if (!Match_Shape(graph, concat, shape, root_input, logger)) {
+      return false;
+    }
     return true;
   }
 
@@ -81,7 +170,7 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node
 }
 
 /**
- * Find the subgraph that matches [root] -> Shape -> Slice -> Squeeze. Check the inputs of slice 
+ * Find the subgraph that matches [root] -> Shape -> Slice -> Squeeze. Check the inputs of slice
  * to make sure the graph produces output with exactly one element.
  */
 bool ReshapeFusion::Match_One_Element_Output_Subgraph_2(Graph& graph, const NodeArg& root_input, const Node& cur_node,
@@ -117,7 +206,7 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_2(Graph& graph, const Node
 }
 
 /**
- * Check if the i-th input of the current node contains exactly one element by checking 
+ * Check if the i-th input of the current node contains exactly one element by checking
  * its inferred shape.
  */
 bool ReshapeFusion::Is_One_Element_Input(const Node& cur_node, int index) {
@@ -140,8 +229,8 @@ bool ReshapeFusion::Is_One_Element_Input(const Node& cur_node, int index) {
 }
 
 /**
- * Search all known patterns of one element subgraphs, which include - 
- * 1. A concat input with inferred shape that can only contain one element. 
+ * Search all known patterns of one element subgraphs, which include -
+ * 1. A concat input with inferred shape that can only contain one element.
  * 2. [root] -> Shape -> Gather(any 1d indice) -> Unsqueeze -> [Concat]
  * 3. [root] -> Shape -> Slice (slice to one element) -> Squeeze -> (Div/Mul) -> Unsqueeze -> [Concat]
  *                                                                      |
@@ -191,7 +280,7 @@ bool ReshapeFusion::Is_One_Element_Output_Subgraph(Graph& graph, const NodeArg&
     auto input_count = binary_node.InputArgCount().front();
 
     for (int i = 0; i < input_count; ++i) {
-      // For each input, look for "one-element subgraph -> concat" or "shape -> slice -> squeeze" path for 
+      // For each input, look for "one-element subgraph -> concat" or "shape -> slice -> squeeze" path for
       // a potential match.
       if (!ReshapeFusion::Is_One_Element_Input(binary_node, i) &&
           !ReshapeFusion::Match_One_Element_Output_Subgraph_2(graph, root_input, binary_node, i, logger)) {
@@ -207,7 +296,7 @@ bool ReshapeFusion::Is_One_Element_Output_Subgraph(Graph& graph, const NodeArg&
 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, or a custom subgraph in which nodes 
+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.
 

onnxruntime/core/optimizer/utils.cc

@@ -195,6 +195,24 @@ bool ValidateShape(const NodeArg& node_arg, const std::initializer_list<int64_t>
   return true;
 }
 
+bool CompareShape(const ONNX_NAMESPACE::TensorShapeProto& node_arg_shape, const ONNX_NAMESPACE::TensorShapeProto& node_arg_other_shape) {
+  if (node_arg_shape.dim_size() != node_arg_other_shape.dim_size())
+    return false;
+
+  if (node_arg_shape.dim_size() < 1)
+    return false;
+
+  for (int i = 0; i < node_arg_shape.dim_size(); ++i) {
+    const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim = node_arg_shape.dim(i);
+    const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim_other = node_arg_other_shape.dim(i);
+    if (!utils::HasDimValue(dim) || !utils::HasDimValue(dim_other))
+      return false;
+    if (dim.dim_value() != dim_other.dim_value())
+      return false;
+  }
+  return true;
+}
+
 bool IsShapeKnownOnAllDims(const NodeArg& node_arg, int expected_dim_size) {
   auto shape = node_arg.Shape();
   if (shape == nullptr || shape->dim_size() != expected_dim_size) {

onnxruntime/core/optimizer/utils.h

@@ -50,6 +50,11 @@ bool AppendTensorFromInitializer(const Graph& graph, const NodeArg& input_arg, s
 */
 bool ValidateShape(const NodeArg& node_arg, const std::initializer_list<int64_t>& expected_dim_values);
 
+/** Compare Shape of node input or output.
+@remarks exactly compare two TensorShapeProtos. Return true if they are same
+*/
+bool CompareShape(const ONNX_NAMESPACE::TensorShapeProto& node_arg_shape, const ONNX_NAMESPACE::TensorShapeProto& node_arg_other_shape);
+
 /** Check check whether each dimension is known for shape of node_arg
 @returns false when shape is nullptr, or total dimension is not same as expected_dim_size length,
          or any dim is unknown (without dim value).

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -1558,6 +1558,42 @@ TEST_F(GraphTransformationTests, ReshapeFusionConcatSubgraphWithMul) {
   }
 }
 
+TEST_F(GraphTransformationTests, ReshapeFusionDistilBertTest) {
+  auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_distillbert.onnx";
+  std::shared_ptr<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  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, *logger_);
+  ASSERT_TRUE(ret.IsOK());
+
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Shape"] == 0);
+  ASSERT_TRUE(op_to_count["Gather"] == 0);
+  ASSERT_TRUE(op_to_count["Unsqueeze"] == 0);
+  ASSERT_TRUE(op_to_count["Concat"] == 0);
+  ASSERT_TRUE(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(), 4);
+
+      const int64_t* val = initializer->data<int64_t>();
+      EXPECT_EQ(val[0], 0);
+      EXPECT_EQ(val[1], -1);
+      EXPECT_EQ(val[2], 2);
+      EXPECT_EQ(val[3], 4);
+    }
+  }
+}
+
 TEST_F(GraphTransformationTests, ExpandElimination) {
   auto model_uri = MODEL_FOLDER "expand_elimination.onnx";
   std::shared_ptr<Model> model;

onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py

@@ -195,7 +195,7 @@ graph = helper.make_graph(
         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("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"),
 
@@ -227,9 +227,9 @@ graph = helper.make_graph(
         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("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"),
 
@@ -248,7 +248,7 @@ graph = helper.make_graph(
         helper.make_tensor('indices0', TensorProto.INT64, [], [0]),
         helper.make_tensor('indices1', TensorProto.INT64, [], [1]),
         helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]),
-        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])        
+        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])
     ]
 )
 
@@ -262,7 +262,7 @@ graph = helper.make_graph(
         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("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"),
         helper.make_node("Pad", ["slice_out", "pads"], ["pad0_out"], "pad0", mode = "constant"),
@@ -280,9 +280,9 @@ graph = helper.make_graph(
     [  # initializers
         helper.make_tensor('indices0', TensorProto.INT64, [], [0]),
         helper.make_tensor('indices1', TensorProto.INT64, [], [1]),
-        helper.make_tensor('pads', TensorProto.INT64, [2], [1, 0]),        
+        helper.make_tensor('pads', TensorProto.INT64, [2], [1, 0]),
         helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]),
-        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])        
+        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])
     ]
 )
 
@@ -296,7 +296,7 @@ graph = helper.make_graph(
         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("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"),
         helper.make_node("Squeeze", ["slice_out"], ["squeeze0_out"], "squeeze0", axes=[0]),
@@ -316,9 +316,9 @@ graph = helper.make_graph(
     [  # initializers
         helper.make_tensor('indices0', TensorProto.INT64, [], [0]),
         helper.make_tensor('indices1', TensorProto.INT64, [], [1]),
-        helper.make_tensor('div_init', TensorProto.INT64, [], [1]),        
+        helper.make_tensor('div_init', TensorProto.INT64, [], [1]),
         helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]),
-        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])        
+        helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3])
     ]
 )
 
@@ -332,7 +332,7 @@ graph = helper.make_graph(
         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("Slice", ["shape2_out", "slice_starts_0", "slice_ends_0"], ["slice0_out"], "slice0"),
         helper.make_node("Squeeze", ["slice0_out"], ["squeeze0_out"], "squeeze0", axes=[0]),
@@ -358,8 +358,62 @@ graph = helper.make_graph(
         helper.make_tensor('slice_starts_0', TensorProto.INT64, [1], [2]),
         helper.make_tensor('slice_ends_0', TensorProto.INT64, [1], [3]),
         helper.make_tensor('slice_starts_1', TensorProto.INT64, [1], [1]),
-        helper.make_tensor('slice_ends_1', TensorProto.INT64, [1], [2])     
+        helper.make_tensor('slice_ends_1', TensorProto.INT64, [1], [2])
     ]
 )
 
-save_model(graph, 'reshape_fusion_concat_subgraph_mul.onnx')
+save_model(graph, 'reshape_fusion_concat_subgraph_mul.onnx')
+
+matmul_weights = [
+    -0.04888916015625, 0.0143280029296875, 0.066650390625,-0.0343017578125,
+    -0.0010356903076171875, -0.00048232078552246094, 0.07470703125, -0.04736328125,
+    0.01454925537109375, -0.0086669921875, -0.051971435546875, -0.0201568603515625,
+    0.040435791015625, -0.019256591796875, 0.0205078125, 0.0111541748046875,
+    0.0071868896484375, -0.0298309326171875, -0.0306549072265625, -0.0225372314453125,
+    -0.04193115234375, 0.07073974609375, -0.048065185546875, 0.0198822021484375,
+    -0.035552978515625, -0.022796630859375, 0.03839111328125, 0.007099151611328125,
+    -0.0080108642578125, -0.0017957687377929688, 0.0266265869140625,-0.028289794921875,
+    0.0032901763916015625, 0.0208740234375, -0.01529693603515625, -0.046600341796875,
+    -0.034637451171875, 0.011322021484375, -0.026458740234375, 0.04656982421875,
+    -0.0091705322265625, 0.017913818359375, -0.019256591796875, -0.001216888427734375,
+    -0.08245849609375, -0.023162841796875, -0.04132080078125, -0.03363037109375,
+    0.0029315948486328125, 0.03173828125, -0.004024505615234375, 0.04534912109375,
+    -0.0036163330078125, -0.03912353515625, -0.00800323486328125, 0.058197021484375,
+    0.05572509765625, 0.01165771484375, 0.06756591796875, 0.05816650390625,
+    -0.0654296875, -0.0241851806640625, 0.0205535888671875, -0.031707763671875
+]
+
+add_weight = [-0.23681640625, -0.16552734375, 0.2191162109375, -0.1756591796875,
+                -0.03460693359375, -0.05316162109375, -0.336181640625, -0.253662109375]
+
+graph = helper.make_graph(
+    [ # nodes
+        helper.make_node("Add", ["Input", "Bias"], ["add0_out"], "add0"),
+        helper.make_node("Shape", ["add0_out"], ["shape0_out"], "shape0"),
+        helper.make_node("Gather", ["shape0_out", "indices0"], ["gather0_out"], "gather0", axis=0),
+        helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]),
+        helper.make_node("Concat", ["unsqueeze0_out", "dim_-1", "dim_2", "dim_4"], ["concat_out"], "concat", axis=0),
+        helper.make_node("MatMul", ["add0_out", "matmul_weight"], ["matmul_out"], "matmul"),
+        helper.make_node("Add", ["matmul_out", "add_weight"], ["add1_out"], "add1"),
+        helper.make_node("Reshape", ["add1_out", "concat_out"], ["Result"], "reshape"),
+    ],
+    "Reshape_Fusion",  #name
+    [  # inputs
+        helper.make_tensor_value_info('Input', TensorProto.FLOAT, [1, 8]),
+    ],
+    [  # outputs
+        helper.make_tensor_value_info('Result', TensorProto.FLOAT, [1, -1, 2, 4]),
+    ],
+    [  # initializers
+        helper.make_tensor('Bias', TensorProto.FLOAT, [8], add_weight),
+        helper.make_tensor('dim_-1', TensorProto.INT64, [1], [-1]),
+        helper.make_tensor('dim_2', TensorProto.INT64, [1], [2]),
+        helper.make_tensor('dim_4', TensorProto.INT64, [1], [4]),
+        helper.make_tensor('indices0', TensorProto.INT64, [], [0]),
+        helper.make_tensor('matmul_weight', TensorProto.FLOAT, [8, 8], matmul_weights),
+        helper.make_tensor('add_weight', TensorProto.FLOAT, [8], add_weight),
+    ]
+)
+
+save_model(graph, 'reshape_fusion_distillbert.onnx')
+