Add Reshape Fusion (#2395)

* Add reshape fusion

* Add some comments

* update comments

* update comment format

* update according to feedback

* update for recent logger change

* fix build error

* (1) Support both input and output edges in find path in graphutils
(2) Add a test case of only one constant initializer of Concat input.
(3) Refactor ReshapeFusion class to allow add more subgraph fusion in the future.

* fix error

* (1) loose constraint on initializer: non constant is allowed for reshape fusion.
(2) Change versions type to vector.
(3) Add logging.
(4) Return false when multiple output edges matched in FindPath. Add comments.

* only allow one direction (input or output) in FindPath

onnxruntime/core/graph/graph_utils.cc

@@ -191,7 +191,7 @@ static bool CanUpdateImplicitInputNameInSubgraphs(const Graph& graph,
       const Node& output_edge_node = *graph.GetNode(output_edge.dst_node);
       if (!CanUpdateImplicitInputNameInSubgraph(output_edge_node, output_edge.arg_name, new_arg_name)) {
         LOGS(logger, WARNING) << " Implicit input name " << output_edge.arg_name
-                                 << " cannot be safely updated to " << new_arg_name << " in one of the subgraphs.";
+                              << " cannot be safely updated to " << new_arg_name << " in one of the subgraphs.";
         return false;
       }
     }
@@ -287,6 +287,10 @@ bool MatchesOpSinceVersion(const Node& node, const std::initializer_list<ONNX_NA
   return std::find(versions.begin(), versions.end(), node.Op()->SinceVersion()) != versions.end();
 }
 
+bool MatchesOpSinceVersion(const Node& node, const std::vector<ONNX_NAMESPACE::OperatorSetVersion>& versions) {
+  return std::find(versions.begin(), versions.end(), node.Op()->SinceVersion()) != versions.end();
+}
+
 bool MatchesOpSetDomain(const Node& node, const std::string& domain) {
   const auto& node_domain = node.Domain();
   // We do a special check for the ONNX domain, as it has two aliases.
@@ -563,7 +567,6 @@ const Node* FirstParentByType(Node& node, const std::string& parent_type) {
   return nullptr;
 }
 
-
 NodeArg& AddInitializer(Graph& graph, const ONNX_NAMESPACE::TensorProto& new_initializer) {
   // sanity check as AddInitializedTensor silently ignores attempts to add a duplicate initializer
   const ONNX_NAMESPACE::TensorProto* existing = nullptr;
@@ -662,5 +665,62 @@ void FinalizeNodeFusion(Graph& graph, const std::vector<std::reference_wrapper<N
   }
 }
 
+const Node::EdgeEnd*
+GetInputEdge(const Node& node, int arg_index) {
+  for (auto it = node.InputEdgesBegin(), end = node.InputEdgesEnd(); it != end; ++it) {
+    if (arg_index == it->GetDstArgIndex()) {
+      return &(*it);
+    }
+  }
+  return nullptr;
+}
+
+const Node* GetInputNode(const Node& node, int arg_index) {
+  const Node::EdgeEnd* edge = GetInputEdge(node, arg_index);
+  if (nullptr == edge) {
+    return nullptr;
+  }
+  return &(edge->GetNode());
+}
+
+bool FindPath(const Node& node, bool is_input_edge, const std::vector<EdgeEndToMatch>& edges_to_match, std::vector<const Node::EdgeEnd*>& result, const logging::Logger& logger) {
+  result.clear();
+  result.reserve(edges_to_match.size());
+
+  const Node* current_node = &node;
+  for (const auto& edge : edges_to_match) {
+    const Node::EdgeEnd* edge_found = nullptr;
+
+    auto edges_begin = is_input_edge ? current_node->InputEdgesBegin() : current_node->OutputEdgesBegin();
+    auto edges_end = is_input_edge ? current_node->InputEdgesEnd() : current_node->OutputEdgesEnd();
+    for (auto it = edges_begin; it != edges_end; ++it) {
+
+      if (edge.dst_arg_index == it->GetDstArgIndex() && edge.src_arg_index == it->GetSrcArgIndex() && edge.op_type == it->GetNode().OpType() && MatchesOpSinceVersion(it->GetNode(), edge.versions) && MatchesOpSetDomain(it->GetNode(), edge.domain)) {
+        // For output edge, there could be multiple edges matched.
+        // This function will return failure in such case by design.
+        if (nullptr != edge_found) {
+          LOGS(logger, WARNING) << "Failed since multiple edges matched:" << current_node->OpType() << "->" << edge.op_type;
+          return false;
+        }
+        edge_found = &(*it);
+
+        // For input edge, each dst_arg_index only accepts one input edge so only there is at most one match.
+        if (is_input_edge) {
+          break;
+        }
+      }
+    }
+
+    if (nullptr == edge_found) {
+      return false;
+    }
+
+    result.push_back(edge_found);
+    current_node = &(edge_found->GetNode());
+  }
+
+  return true;
+}
+
 }  // namespace graph_utils
 }  // namespace onnxruntime

onnxruntime/core/graph/graph_utils.h

@@ -18,6 +18,7 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node,
 
 /** Checks if the node has the same operator since version as the given one. */
 bool MatchesOpSinceVersion(const Node& node, const std::initializer_list<ONNX_NAMESPACE::OperatorSetVersion>& versions);
+bool MatchesOpSinceVersion(const Node& node, const std::vector<ONNX_NAMESPACE::OperatorSetVersion>& versions);
 
 /** Checks if the node has the same op set domain as the given one. */
 bool MatchesOpSetDomain(const Node& node, const std::string& domain);
@@ -187,5 +188,60 @@ void FinalizeNodeFusion(Graph& graph, Node& first_node, Node& second_node);
 */
 void FinalizeNodeFusion(Graph& graph, const std::vector<std::reference_wrapper<Node>>& nodes, Node& replacement_node);
 
+/** Find the input edge of a node for a specified input index.
+@returns nullptr when not found.
+*/
+const Node::EdgeEnd* GetInputEdge(const Node& node, int arg_index);
+
+/** Find the source node of an input edge for a specified input index.
+@returns nullptr when not found.
+*/
+const Node* GetInputNode(const Node& node, int arg_index);
+
+
+/** Expected edge end information for matching input or output edge.
+    For input edge, the node in the edge end refers to the source node, otherwise the destination node.
+*/
+struct EdgeEndToMatch {
+  // Source arg index of edge.
+  int src_arg_index;
+
+  // Destination arg index of edge.
+  int dst_arg_index;
+
+  // Expected operator type of the node in the edge end.
+  std::string op_type;
+
+  // Expected version of the operator of node in the edge end.
+  std::vector<ONNX_NAMESPACE::OperatorSetVersion> versions;
+
+  // Expected domain of the operator of node in the edge end.
+  std::string domain;
+};
+
+/** Find a path that matches the specified edges.
+    A path is a sequence of adjacent edges, and the result is returned as a list of EdgeEnd items.
+@param node is the current node to start matching.
+@param is_input_edge is a flag to indicate  whether the edges are input or output edges.
+@param edges_to_match has information of a sequence of adjacent edges in the path to be matched one by one.
+@param result stores edges that are found.
+@returns false when one edge has multiple candidates, or not all edges are found.
+@remarks matching an EdgeEndToMatch might get multiple candidates in output edges. 
+    When such case is encountered, this function will return false. This is by design to reduce complexity.
+    Here is an example graph:
+                  Add
+                 /    \
+               Mul     Mul
+                 \    /
+                  Sub
+    For example, you want to match path from top to bottom: Add-->Mul-->Sub. 
+    When matching the first edge Add-->Mul, the algorithm found two matches.
+    Then it returns false, and output a warning log entry.
+
+    It is recommended to match path from bottom to top direction to avoid such issue.
+    It is because each node input (dst_arg_index) only accepts one input edge.
+*/
+bool FindPath(const Node& node, bool is_input_edge, const std::vector<EdgeEndToMatch>& edges_to_match, std::vector<const Node::EdgeEnd*>& result, const logging::Logger& logger);
+
 }  // namespace graph_utils
 }  // namespace onnxruntime

onnxruntime/core/optimizer/gelu_fusion.cc

@@ -3,6 +3,7 @@
 
 #include "core/optimizer/initializer.h"
 #include "core/optimizer/gelu_fusion.h"
+#include "core/optimizer/utils.h"
 #include "core/graph/graph_utils.h"
 #include "float.h"
 #include <deque>
@@ -11,49 +12,6 @@ using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;
 namespace onnxruntime {
 
-static bool CheckConstantInput(const Graph& graph, const NodeArg& input_arg, float expected_value) {
-  auto shape = input_arg.Shape();
-  if (shape == nullptr) {
-    // shape inferencing wasn't able to populate shape information for this NodeArg
-    return false;
-  }
-
-  auto dim_size = shape->dim_size();
-  if (dim_size != 0) {
-    // only check scalar.
-    return false;
-  }
-
-  const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, input_arg.Name());
-  if (tensor_proto == nullptr) {
-    return false;
-  }
-
-  auto init_const = onnxruntime::make_unique<Initializer>(*tensor_proto);
-  const auto data_type = tensor_proto->data_type();
-  if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
-    float* val = init_const->data<float>();
-    float diff = std::abs(val[0] - static_cast<float>(expected_value));
-    if (diff > FLT_EPSILON) {
-      return false;
-    }
-  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE) {
-    double* val = init_const->data<double>();
-    double diff = std::abs(val[0] - static_cast<double>(expected_value));
-    if (diff > DBL_EPSILON) {
-      return false;
-    }
-  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
-    MLFloat16* val = init_const->data<MLFloat16>();
-    float diff = std::abs(math::halfToFloat(val[0].val) - static_cast<float>(expected_value));
-    if (diff > FLT_EPSILON) {
-      return false;
-    }
-  }
-
-  return true;
-}
-
 // Gelu supports limited data types.
 static std::vector<std::string> supported_data_types{"tensor(float16)", "tensor(float)", "tensor(double)"};
 
@@ -87,7 +45,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
     }
 
     // Check second input is sqrt(2)
-    if (!CheckConstantInput(graph, *(div.MutableInputDefs()[1]), static_cast<float>(M_SQRT2))) {
+    if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(div.InputDefs()[1]), static_cast<float>(M_SQRT2), true)) {
       continue;
     }
 
@@ -114,7 +72,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
       add_const_input_index = 1;
     }
     const auto& add_const_input_arg = add_node.InputDefs()[add_const_input_index];
-    if (!CheckConstantInput(graph, *add_const_input_arg, 1.0f)) {
+    if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *add_const_input_arg, 1.0f, true)) {
       continue;
     }
 
@@ -153,7 +111,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons
     }
 
     const auto& mul_const_input_arg = mul2_node.InputDefs()[mul_const_input_index];
-    if (!CheckConstantInput(graph, *mul_const_input_arg, 0.5f)) {
+    if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *mul_const_input_arg, 0.5f, true)) {
       continue;
     }
 

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -22,6 +22,7 @@
 #include "core/optimizer/gelu_fusion.h"
 #include "core/optimizer/layer_norm_fusion.h"
 #include "core/optimizer/skip_layer_norm_fusion.h"
+#include "core/optimizer/reshape_fusion.h"
 #include "core/mlas/inc/mlas.h"
 #include "core/session/inference_session.h"
 
@@ -104,6 +105,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
 
       transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(l1_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<MatMulAddFusion>(l1_execution_providers));
+      transformers.emplace_back(onnxruntime::make_unique<ReshapeFusion>(l1_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<FreeDimensionOverrideTransformer>(free_dimension_overrides));
 
       rule_transformer = GenerateRuleBasedGraphTransformer(level, transformers_and_rules_to_enable, l1_execution_providers);

onnxruntime/core/optimizer/reshape_fusion.cc

@@ -0,0 +1,220 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/graph/graph_utils.h"
+#include "core/optimizer/initializer.h"
+#include "core/optimizer/reshape_fusion.h"
+#include "core/optimizer/utils.h"
+
+using namespace ONNX_NAMESPACE;
+using namespace onnxruntime::common;
+namespace onnxruntime {
+
+// Get values of integer tensor from initializer, and append them to a vector.
+static bool LoadIntegerTensor(const Graph& graph, const NodeArg& input_arg, std::vector<int64_t>& data) {
+  const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
+  if (!graph.GetInitializedTensor(input_arg.Name(), tensor_proto)) {
+    return false;
+  }
+
+  auto init_const = onnxruntime::make_unique<Initializer>(*tensor_proto);
+  const auto data_type = tensor_proto->data_type();
+  if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64) {
+    const int64_t* val = init_const->data<int64_t>();
+    data.reserve(data.size() + init_const->size());
+    data.insert(data.end(), val, val + init_const->size());
+  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT32) {
+    const int32_t* val = init_const->data<int32_t>();
+    data.reserve(data.size() + init_const->size());
+    for (int64_t i = 0; i < init_const->size(); i++) {
+      data.push_back(static_cast<int64_t>(val[i]));
+    }
+  } else {
+    return false;
+  }
+
+  return true;
+}
+
+Status ReshapeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const {
+  GraphViewer graph_viewer(graph);
+  const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
+
+  int fused_count = 0;
+  for (auto node_index : node_topology_list) {
+    auto* p_reshape = graph.GetNode(node_index);
+    if (p_reshape == nullptr)
+      continue;  // we removed the node as part of an earlier fusion
+
+    Node& reshape = *p_reshape;
+    ORT_RETURN_IF_ERROR(Recurse(reshape, modified, graph_level, logger));
+
+    if (!graph_utils::IsSupportedOptypeVersionAndDomain(reshape, "Reshape", {5}) ||
+        !graph_utils::IsSupportedProvider(reshape, GetCompatibleExecutionProviders())) {
+      continue;
+    }
+
+    if (ReshapeFusion::Fuse_Subgraph1(reshape, graph, logger)) {
+      fused_count++;
+      LOGS(logger, INFO) << "Fused reshape node: " << reshape.OutputDefs()[0]->Name();
+      modified = true;
+    }
+  }
+  LOGS(logger, INFO) << "Total fused reshape node count: " << fused_count;
+
+  return Status::OK();
+}
+
+/**
+Apply Reshape Fusion. The fowllowing are subgraphs before and after fusion:
+
+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              /          /
+    |         \        /  ___________________/          /
+    |          \      /  / ____________________________/
+    |           \    /  / /
+     \            Concat
+      \          /
+         Reshape
+
+After fusion:
+    [Sub-graph Root Node]   (Constant Initializers, b is optional)
+                  \         [0, 0, a, b]
+                   \        /
+                    Reshape
+*/
+bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::Logger& logger) {
+  const Node* p_root = graph_utils::GetInputNode(reshape, 0);
+
+  const Node* p_concat = graph_utils::GetInputNode(reshape, 1);
+  if (nullptr == p_concat) {
+    return false;
+  }
+  const Node& concat = *p_concat;
+
+  if (!graph_utils::IsSupportedOptypeVersionAndDomain(concat, "Concat", {1, 4})) {
+    return false;
+  }
+
+  auto concat_input_count = concat.InputArgCount().front();
+  if (concat_input_count < 3 || concat_input_count > 4 || 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}, kOnnxDomain},
+      {0, 0, "Gather", {1}, kOnnxDomain},
+      {0, 0, "Shape", {1}, kOnnxDomain}};
+
+  std::vector<const Node::EdgeEnd*> edges;
+  if (!graph_utils::FindPath(concat, true, parent_path, edges, logger)) {
+    return false;
+  }
+
+  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}, kOnnxDomain},
+      {0, 0, "Gather", {1}, 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 (!LoadIntegerTensor(graph, *(concat.InputDefs()[2]), shape_value)) {
+    return false;
+  }
+
+  if (concat_input_count > 3) {
+    if (!LoadIntegerTensor(graph, *(concat.InputDefs()[3]), shape_value)) {
+      return false;
+    }
+  }
+
+  // Create an initializer with the same name as the concat node output, and replace the concat node
+  const auto& new_initializer_name = concat.OutputDefs()[0]->Name();
+  if (!graph_utils::CanReplaceNodeWithInitializer(graph, concat, new_initializer_name, logger)) {
+    LOGS(logger, WARNING) << "Cannot replace concat node with initializer:" << new_initializer_name;
+    return false;
+  }
+  const auto* shape_def = concat.OutputDefs()[0];
+  ONNX_NAMESPACE::TensorProto shape_initializer_proto;
+  shape_initializer_proto.set_name(shape_def->Name());
+  shape_initializer_proto.add_dims(static_cast<int64_t>(shape_value.size()));
+  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 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());
+  }
+
+  return true;
+}
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/reshape_fusion.h

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/optimizer/graph_transformer.h"
+
+namespace onnxruntime {
+
+/**
+@Class ReshapeFusion
+Rewrite graph fusing reshape subgraph to a single Reshape node.
+*/
+class ReshapeFusion : public GraphTransformer {
+ public:
+  ReshapeFusion(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept
+      : GraphTransformer("ReshapeFusion", compatible_execution_providers) {}
+
+  Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
+
+private:
+  static bool Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::Logger& logger);
+};
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/utils.cc

@@ -0,0 +1,114 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#include "core/common/make_unique.h"
+#include "core/graph/onnx_protobuf.h"
+#include "core/graph/graph_utils.h"
+#include "core/framework/tensorprotoutils.h"
+#include "core/optimizer/initializer.h"
+#include "core/framework/utils.h"
+#include "core/optimizer/utils.h"
+#include "float.h"
+//#include <deque>
+
+using namespace onnxruntime;
+
+namespace onnxruntime {
+namespace optimizer_utils {
+
+bool IsFloatingPointDataType(const ONNX_NAMESPACE::TensorProto& tensor_proto) {
+  return tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT
+      || tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16
+      || tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE;
+}
+
+inline bool IsScalar(const NodeArg& input_arg) {
+  auto shape = input_arg.Shape();
+  if (shape == nullptr) {
+    // shape inferencing wasn't able to populate shape information for this NodeArg
+    return false;
+  }
+
+  auto dim_size = shape->dim_size();
+  return dim_size == 0 || (dim_size == 1 && shape->dim(0).has_dim_value() && shape->dim(0).dim_value() == 1);
+}
+
+// Check whether input is a constant scalar with expected float value.
+bool IsInitializerWithExpectedValue(const Graph& graph, const NodeArg& input_arg, float expected_value, bool is_constant) {
+  if (!IsScalar(input_arg)) {
+    return false;
+  }
+
+  const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
+  if (is_constant) {
+    tensor_proto = graph_utils::GetConstantInitializer(graph, input_arg.Name());
+  } else if (!graph.GetInitializedTensor(input_arg.Name(), tensor_proto)) {
+    return false;
+  }
+  
+  if (tensor_proto == nullptr) {
+    return false;
+  }
+
+  auto init_const = onnxruntime::make_unique<Initializer>(*tensor_proto);
+  const auto data_type = tensor_proto->data_type();
+  if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
+    const float* val = init_const->data<float>();
+    float diff = std::abs(val[0] - static_cast<float>(expected_value));
+    if (diff > FLT_EPSILON) {
+      return false;
+    }
+  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE) {
+    const double* val = init_const->data<double>();
+    double diff = std::abs(val[0] - static_cast<double>(expected_value));
+    if (diff > DBL_EPSILON) {
+      return false;
+    }
+  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
+    const MLFloat16* val = init_const->data<MLFloat16>();
+    float diff = std::abs(math::halfToFloat(val[0].val) - static_cast<float>(expected_value));
+    if (diff > FLT_EPSILON) {
+      return false;
+    }
+  } else {
+    // Not expected data types.
+    return false;
+  }
+
+  return true;
+}
+
+// Check whether input is a constant scalar with expected intger value.
+bool IsInitializerWithExpectedValue(const Graph& graph, const NodeArg& input_arg, int64_t expected_value, bool is_constant) {
+  if (!IsScalar(input_arg)) {
+    return false;
+  }
+
+  const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
+  if (is_constant) {
+    tensor_proto = graph_utils::GetConstantInitializer(graph, input_arg.Name());
+  } else if (!graph.GetInitializedTensor(input_arg.Name(), tensor_proto)) {
+    return false;
+  }
+
+  auto init_const = onnxruntime::make_unique<Initializer>(*tensor_proto);
+  const auto data_type = tensor_proto->data_type();
+  if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64) {
+    const int64_t* val = init_const->data<int64_t>();
+    if (val[0] != expected_value) {
+      return false;
+    }
+  } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT32) {
+    const int32_t* val = init_const->data<int32_t>();
+    if (static_cast<int64_t>(val[0]) != expected_value) {
+      return false;
+    }
+  } else {
+    // Not expected data types.
+    return false;
+  }
+
+  return true;
+}
+
+}  // namespace optimizer_utils
+}  // namespace onnxruntime

onnxruntime/core/optimizer/utils.h

@@ -4,15 +4,31 @@
 #pragma once
 
 #include "core/graph/onnx_protobuf.h"
+#include "core/graph/graph.h"
 
 namespace onnxruntime {
+class Graph;
+class NodeArg;
+
 namespace optimizer_utils {
 
 // Check if TensorProto contains a floating point type.
-static bool IsFloatingPointDataType(const ONNX_NAMESPACE::TensorProto& tensor_proto) {
-  return !(tensor_proto.data_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
-           tensor_proto.data_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 &&
-           tensor_proto.data_type() != ONNX_NAMESPACE::TensorProto_DataType_DOUBLE);
-}
+bool IsFloatingPointDataType(const ONNX_NAMESPACE::TensorProto& tensor_proto);
+
+/** Check whether a input is initializer with specified float value.
+@param expected_value is the expected value of the initializer.
+@param is_constant means whether the initializer is required to be constant.
+@remarks only support float16, float and double scalar.
+*/
+bool IsInitializerWithExpectedValue(const onnxruntime::Graph& graph, const onnxruntime::NodeArg& input_arg, float expected_value, bool is_constant);
+
+
+/** Check whether a input is initializer with specified integer value.
+@param expected_value is the expected value of the initializer.
+@param is_constant means whether the initializer is required to be constant.
+@remarks only support int32 and int64 scalar.
+*/
+bool IsInitializerWithExpectedValue(const onnxruntime::Graph& graph, const onnxruntime::NodeArg& input_arg, int64_t expected_value, bool is_constant);
+
 }  // namespace optimizer_utils
 }  // namespace onnxruntime

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -28,6 +28,7 @@
 #include "core/optimizer/shape_to_initializer.h"
 #include "core/optimizer/slice_elimination.h"
 #include "core/optimizer/unsqueeze_elimination.h"
+#include "core/optimizer/reshape_fusion.h"
 #include "core/platform/env.h"
 #include "core/util/math.h"
 #include "test/capturing_sink.h"
@@ -803,6 +804,79 @@ TEST(GraphTransformationTests, ReluClip11Fusion) {
   }
 }
 
+// Test Reshape Fusion with 2 constant initializers for Concat inputs.
+TEST(GraphTransformationTests, ReshapeFusionTest) {
+  auto model_uri = MODEL_FOLDER "fusion/reshape.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_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);
+      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], 0);
+      EXPECT_EQ(val[2], 12);
+      EXPECT_EQ(val[3], 64);
+    }
+  }
+}
+
+// Test Reshape Fusion with one constant initializer for Concat inputs.
+TEST(GraphTransformationTests, ReshapeFusionOneConstTest) {
+    auto model_uri = MODEL_FOLDER "fusion/reshape_one_const.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_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);
+            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], 768);
+        }
+    }
+}
+
 #ifndef DISABLE_CONTRIB_OPS
 TEST(GraphTransformationTests, GeluFusionTest) {
   auto model_uri = MODEL_FOLDER "fusion/gelu.onnx";