From e1c17fd126a474b1e21a7c160eaaab334e82a545 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 18 Nov 2019 10:07:10 -0800 Subject: [PATCH] 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 | 64 ++++- onnxruntime/core/graph/graph_utils.h | 56 +++++ onnxruntime/core/optimizer/gelu_fusion.cc | 50 +--- .../core/optimizer/graph_transformer_utils.cc | 2 + onnxruntime/core/optimizer/reshape_fusion.cc | 220 ++++++++++++++++++ onnxruntime/core/optimizer/reshape_fusion.h | 25 ++ onnxruntime/core/optimizer/utils.cc | 114 +++++++++ onnxruntime/core/optimizer/utils.h | 26 ++- .../test/optimizer/graph_transform_test.cc | 74 ++++++ .../testdata/transform/fusion/reshape.onnx | Bin 0 -> 3691 bytes .../transform/fusion/reshape_one_const.onnx | Bin 0 -> 525 bytes 11 files changed, 578 insertions(+), 53 deletions(-) create mode 100644 onnxruntime/core/optimizer/reshape_fusion.cc create mode 100644 onnxruntime/core/optimizer/reshape_fusion.h create mode 100644 onnxruntime/core/optimizer/utils.cc create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape_one_const.onnx diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index a37fd6812e..321d67a45b 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/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_listSinceVersion()) != versions.end(); } +bool MatchesOpSinceVersion(const Node& node, const std::vector& 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::vectorGetDstArgIndex()) { + 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& edges_to_match, std::vector& 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 diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index 4a9fe05dd1..36303479cf 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/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& versions); +bool MatchesOpSinceVersion(const Node& node, const std::vector& 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>& 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 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& edges_to_match, std::vector& result, const logging::Logger& logger); + } // namespace graph_utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/gelu_fusion.cc b/onnxruntime/core/optimizer/gelu_fusion.cc index 06a5830550..550fafd02b 100644 --- a/onnxruntime/core/optimizer/gelu_fusion.cc +++ b/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 @@ -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(*tensor_proto); - const auto data_type = tensor_proto->data_type(); - if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { - float* val = init_const->data(); - float diff = std::abs(val[0] - static_cast(expected_value)); - if (diff > FLT_EPSILON) { - return false; - } - } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE) { - double* val = init_const->data(); - double diff = std::abs(val[0] - static_cast(expected_value)); - if (diff > DBL_EPSILON) { - return false; - } - } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { - MLFloat16* val = init_const->data(); - float diff = std::abs(math::halfToFloat(val[0].val) - static_cast(expected_value)); - if (diff > FLT_EPSILON) { - return false; - } - } - - return true; -} - // Gelu supports limited data types. static std::vector 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(M_SQRT2))) { + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(div.InputDefs()[1]), static_cast(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; } diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 75b380eebd..c05d02687f 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/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> GenerateTransformers(TransformerL transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); + transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(free_dimension_overrides)); rule_transformer = GenerateRuleBasedGraphTransformer(level, transformers_and_rules_to_enable, l1_execution_providers); diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc new file mode 100644 index 0000000000..b8f14a6a48 --- /dev/null +++ b/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& data) { + const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr; + if (!graph.GetInitializedTensor(input_arg.Name(), tensor_proto)) { + return false; + } + + auto init_const = onnxruntime::make_unique(*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(); + 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(); + data.reserve(data.size() + init_const->size()); + for (int64_t i = 0; i < init_const->size(); i++) { + data.push_back(static_cast(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 parent_path{ + {0, 0, "Unsqueeze", {1}, kOnnxDomain}, + {0, 0, "Gather", {1}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}}; + + std::vector 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 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}, 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 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(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 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 diff --git a/onnxruntime/core/optimizer/reshape_fusion.h b/onnxruntime/core/optimizer/reshape_fusion.h new file mode 100644 index 0000000000..d906c556b1 --- /dev/null +++ b/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& 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 diff --git a/onnxruntime/core/optimizer/utils.cc b/onnxruntime/core/optimizer/utils.cc new file mode 100644 index 0000000000..1f6df8ca80 --- /dev/null +++ b/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 + +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(*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 diff = std::abs(val[0] - static_cast(expected_value)); + if (diff > FLT_EPSILON) { + return false; + } + } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE) { + const double* val = init_const->data(); + double diff = std::abs(val[0] - static_cast(expected_value)); + if (diff > DBL_EPSILON) { + return false; + } + } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + const MLFloat16* val = init_const->data(); + float diff = std::abs(math::halfToFloat(val[0].val) - static_cast(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(*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(); + if (val[0] != expected_value) { + return false; + } + } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT32) { + const int32_t* val = init_const->data(); + if (static_cast(val[0]) != expected_value) { + return false; + } + } else { + // Not expected data types. + return false; + } + + return true; +} + +} // namespace optimizer_utils +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/utils.h b/onnxruntime/core/optimizer/utils.h index d50a134c1f..4de2fb119e 100644 --- a/onnxruntime/core/optimizer/utils.h +++ b/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 diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 976840549c..93d40cbd9a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/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 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_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(*tensor_proto); + EXPECT_EQ(tensor_proto->data_type(), ONNX_NAMESPACE::TensorProto_DataType_INT64); + EXPECT_EQ(initializer->size(), 4); + + const int64_t* val = initializer->data(); + 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 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_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(*tensor_proto); + 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], 768); + } + } +} + #ifndef DISABLE_CONTRIB_OPS TEST(GraphTransformationTests, GeluFusionTest) { auto model_uri = MODEL_FOLDER "fusion/gelu.onnx"; diff --git a/onnxruntime/test/testdata/transform/fusion/reshape.onnx b/onnxruntime/test/testdata/transform/fusion/reshape.onnx new file mode 100644 index 0000000000000000000000000000000000000000..d19398414563ea16bb827782a36a31b7c8cb506a GIT binary patch literal 3691 zcmb7{dvI0N6~@<1NOE%$BngiMBIG1WfW!$w0_1gXHc+L84nkBkTHAAzTpk(|&Ao`A zFb9Do#3EuTtvu}E(E<_!Lq)*ko*YG`wiT^KN}19gs|7pEI8sn*iDQ2!*lK6$A7y4| zzt;M`Z>_!dJ{Nyt#_&Z=3!{289Es`=G&h7B_4(nZa7!rK+){q*TCGTnpHN)l{h6f2 zmy}HO{*+`eEzL`u-rQ6liWa1314551jAV)Bi_-#HTxqFRR2OcE7KNMY0cHa8VT!j1Ea9$FM`SyEKDFcis%D=o{5pW4u%KZ;eQ(qBAPNV zrsO5uwjlIiSR2b^a-N^aBhgS(v>;1McsSI!C|r=5mQ?MRFaKG-0a{wWtH<$m-?DG? z^?l7Z)ARN8^{1jNFR3OR@#r&#MHbb~ZwWoPK%dv#qW|aM3I_XqX-U((^QrcyehKu| zVo9}M{d@6nmfuWF^G*9dPpI0=60^k?bDS-P5fpc#EzZE0IJGg+J8f}Su$LYb=TboA zVoa>mfaq-jQSC8NL)aT^i+V30)=B)o9TZh!i#3<^%qr1u3yAqL_UkJ}a@l((C{})t zxHE%d466`1Tq%wp-@YEPR#b|47vpg;k)2?rV1u1A7M@?A;~iRkP{cvUvVge7wg@)P zfq)oM_OgRwE@Xc)+$*}l#_yJ(==;$SB%U1h4+g|RyI629hP#?^E%AN|_n%@{G2RC+ zhHWl!+>YNyd`_e52yq^Urx#9tujt3YuBOd_ZwKomS*r|+`v$o=hW78#l6%G3PafYV zp80T1kBK!L+YWMZ0IfE-h2ZJ*Z}o_>;Yg_xb3c6zeR_{Lui2ves>C6u#v|mfuu7~2 z#IEo^jmBAM91q4rv>ft00Q=u!w~IL1vAgAvxHn*%-P?~Ziy4zs<5PTddPQbq|5Qu_ zJ&Igeo#4IBK3sYs>nB+MEAuDFNmr#_ zCRRs}n9m2rXa(dhmMyqobogHH!{?~{P&zkxHJei@v(Xstj?De-+qY-p0LL2(DO z&e++4?@hEl*jyx@O`P!r=54ek@O%^ytB$=b_zfp-C-A#OTz>$68{Cy>g9g;($KR7%q-rTq2z3Mub5-7E9UJU$T*`)6z_&KVs`;7y;rPZob}DDjo}?~(RPL0 zzK5nw=wVIe)4I^I4PE3!evNMl<7>dA{@sc6&vWig@cfZ8dBKYZ%v9c|&xrFBmH5TL;Df<{uFIG;H(Wnoqtv(K(wlnTZcIsJiHiTK&cA0ZJ#J!ml6w~`&7;wG|7j^2sQz434RG*NSN$8k}Z6@#15$<3+ z=bwO{MsQoP;Vzl4v2UU8K6s~gi@OTEW!O}({|e{1jCixqhpjTP=i6uw!UrAJ8uC$y z-5lcDgU?>ZI|BW0f%yaO+zD0Ao%s&?r}4c4-aXXT=bZ0R@NeQ9{%TNU2DSLyRV-&D zJ@?G2h4Tx>tH9rdp6S$4pDk)GKKnRt?x%ML^zqCu5cA8#2EOUt1%7w9=g`oHp7CJu z_L+~7kMF{H9@{zCoC8l|oja;Nv}XqwZ;;wb-L;pCQ3d{F^5E^=g)KGjy}eEsV`^3= zVw(gG_fQ>yH;1@=*ih5P2b>MR0nFp9@1*Ar7@g$&R&=f8H%J!eHX1#}*mvOXalBYA z4&Po?V~cttMr`1P1N-L=YNuKh#4T$R?0FmwV`aLf3+FsPi3N}#ib?x3;Oq$mjD0& literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_one_const.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_one_const.onnx new file mode 100644 index 0000000000000000000000000000000000000000..7f9324f7d03e8cdb66b87f458c889f6cf5198284 GIT binary patch literal 525 zcmZ{h%TB{E5Jl^xO`*0k*Fn)#+q~Q%*fcGr$su`o)$`bi|kd*y@^PLW`AAD$kVgTcnx`xkxIBm7v64B9}1m>ws0r z>QAY9FNF9KD>tDM;i1*-o))dHJ=0fLy- z9;Q4O