From 61b733ce6d587a89b4d0dacdfdf6e6ac83ed1609 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Mon, 15 Jul 2019 12:41:01 +1000 Subject: [PATCH] Update optimizers to be able to utilize a constant initializer from an ancestor graph (#1346) * Now that we check for a constant initializer in an ancestor graph we also need to be able to retrieve and replace that initializer. Add helpers to do so. Update optimizers to use the new helpers. Fix bug in UnsqueezeElimination where it wasn't checking if the initializer it was replacing was constant. --- include/onnxruntime/core/graph/graph.h | 7 + onnxruntime/core/graph/graph.cc | 9 ++ onnxruntime/core/graph/graph_utils.cc | 135 +++++++++--------- onnxruntime/core/graph/graph_utils.h | 26 +++- .../core/optimizer/constant_folding.cc | 13 +- onnxruntime/core/optimizer/conv_add_fusion.cc | 27 ++-- onnxruntime/core/optimizer/conv_bn_fusion.cc | 46 +++--- onnxruntime/core/optimizer/conv_mul_fusion.cc | 25 ++-- .../core/optimizer/unsqueeze_elimination.cc | 21 +-- .../test/framework/session_state_test.cc | 2 +- onnxruntime/test/framework/test_utils.cc | 27 +++- .../test/optimizer/graph_transform_test.cc | 77 +++++++++- 12 files changed, 253 insertions(+), 162 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 4abd12499e..d671e21991 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -279,6 +279,10 @@ class Node { return !attr_to_subgraph_map_.empty(); } + /** Get the const subgraphs from a node. + @remarks Creates a new vector so calling ContainsSubgraphs first is preferred. */ + std::vector> GetSubgraphs() const; + /** Gets a map of attribute name to the mutable Graph instances for all subgraphs of the Node. @returns Map of the attribute name that defines the subgraph to the subgraph's Graph instance. nullptr if the Node has no subgraphs. @@ -756,6 +760,9 @@ class Graph { /** Returns the parent graph if this is a subgraph */ const Graph* ParentGraph() const { return parent_graph_; } + /** Returns the mutable parent graph if this is a subgraph */ + Graph* MutableParentGraph() { return parent_graph_; } + /** Construct a Graph instance for a subgraph that is created from a GraphProto attribute in a Node. Inherits some properties from the parent graph. @param parent_graph The Graph containing the Node which has a GraphProto attribute. diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 605aaba84f..af46870030 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -549,6 +549,15 @@ const Graph* Node::GetGraphAttribute(const std::string& attr_name) const { return const_cast(this)->GetMutableGraphAttribute(attr_name); } +std::vector> Node::GetSubgraphs() const { + std::vector> subgraphs; + subgraphs.reserve(attr_to_subgraph_map_.size()); + std::transform(attr_to_subgraph_map_.cbegin(), attr_to_subgraph_map_.cend(), std::back_inserter(subgraphs), + [](const auto& entry) { return entry.second; }); + + return subgraphs; +} + void Node::ForEachDef(std::function func, bool include_missing_optional_defs) const { for (const auto* arg : InputDefs()) { diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index 2ca9358ac2..15d2613b99 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -54,18 +54,20 @@ static bool OutputEdgeProvidesImplicitInput(const Graph& graph, const GraphEdge& /** Checks if new_output_name can be used to replace removed_output_name in the subgraph input. If there is an existing NodeArg in a subgraph that implicitly consumes removed_output_name, it is not safe. */ -static bool CanUpdateImplicitInputNameInSubgraph(Node& node, +static bool CanUpdateImplicitInputNameInSubgraph(const Node& node, const std::string& removed_output_name, const std::string& new_output_name) { - for (auto& attr_subgraph_pair : node.GetAttributeNameToMutableSubgraphMap()) { - Graph& subgraph = *attr_subgraph_pair.second; + if (!node.ContainsSubgraph()) + return true; + + for (const gsl::not_null& subgraph : node.GetSubgraphs()) { // if we have an existing NodeArg in the subgraph with the new_output_name that would override an implicit input // with the same name - if (subgraph.GetNodeArg(new_output_name) != nullptr) { + if (subgraph->GetNodeArg(new_output_name) != nullptr) { return false; } - for (auto& subgraph_node : attr_subgraph_pair.second->Nodes()) { + for (auto& subgraph_node : subgraph->Nodes()) { // recurse if this node also consumes removed_output_name as an implicit input (i.e. there are multiple levels of nested // subgraphs, and at least one level lower uses removed_output_name as an implicit input const auto subgraph_node_implicit_inputs = subgraph_node.ImplicitInputDefs(); @@ -248,6 +250,26 @@ static bool RemoveNodeWithSingleInitializerIn(Graph& graph, Node& node) { return true; } +static bool ReplaceInitializerImpl(Graph& graph, const std::string& original_name, + const ONNX_NAMESPACE::TensorProto& initializer, bool check_outer_scope) { + bool replaced = false; + const ONNX_NAMESPACE::TensorProto* old_initializer = nullptr; + if (graph.GetInitializedTensor(original_name, old_initializer)) { + // Be conservative and only remove if the name matches. Graph::CleanupUnusedInitializers can take care + // of removing anything unused after optimization + if (original_name == initializer.name()) { + graph.RemoveInitializedTensor(original_name); + } + graph.AddInitializedTensor(initializer); + replaced = true; + + } else if (check_outer_scope && graph.IsSubgraph()) { + replaced = ReplaceInitializerImpl(*graph.MutableParentGraph(), original_name, initializer, check_outer_scope); + } + + return replaced; +} + //---------------------------- //--- end of local helpers --- //---------------------------- @@ -292,52 +314,6 @@ bool IsSupportedProvider(const Node& node, compatible_providers.find(node.GetExecutionProviderType()) == compatible_providers.end()); } -Status ForAllMutableSubgraphs(Graph& graph, std::function func) { - Status status = Status::OK(); - - for (auto& node : graph.Nodes()) { - for (auto& attr_name_to_subgraph_pair : node.GetAttributeNameToMutableSubgraphMap()) { - Graph* subgraph = attr_name_to_subgraph_pair.second; - ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved."); - - status = func(*subgraph); - ORT_RETURN_IF_ERROR(status); - - // recurse - status = ForAllMutableSubgraphs(*subgraph, func); - ORT_RETURN_IF_ERROR(status); - } - } - - return status; -} - -Status ForAllSubgraphs(const Graph& graph, std::function func) { - Status status = Status::OK(); - - for (auto& node : graph.Nodes()) { - for (auto& attribute : node.GetAttributes()) { - auto& name = attribute.first; - auto& proto = attribute.second; - - // check if it has a subgraph - if (proto.has_g()) { - const Graph* subgraph = node.GetGraphAttribute(name); - ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved."); - - status = func(*subgraph); - ORT_RETURN_IF_ERROR(status); - - // recurse - status = ForAllSubgraphs(*subgraph, func); - ORT_RETURN_IF_ERROR(status); - } - } - } - - return status; -} - bool IsSingleInSingleOutNode(const Node& node) { return node.InputDefs().size() == 1 && node.ImplicitInputDefs().empty() && node.OutputDefs().size() == 1; } @@ -399,48 +375,69 @@ bool IsGraphInput(const Graph& graph, const NodeArg* input) { return std::find(graph_inputs.begin(), graph_inputs.end(), input) != graph_inputs.end(); } -bool IsConstantInitializer(const Graph& graph, const std::string& initializer_name, bool check_outer_scope) { - const onnx::TensorProto* initializer = nullptr; - bool is_local_initializer = graph.GetInitializedTensor(initializer_name, initializer); - - // if we know it's an initializer we assume it's constant initially. - // otherwise it's not an initializer so can't be a constant initializer by definition. - bool constant_initializer = is_local_initializer; - - if (is_local_initializer) { +const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const Graph& graph, const std::string& initializer_name, + bool check_outer_scope) { + const ONNX_NAMESPACE::TensorProto* initializer = nullptr; + if (graph.GetInitializedTensor(initializer_name, initializer)) { if (graph.CanOverrideInitializer()) { const auto& graph_inputs = graph.GetInputsIncludingInitializers(); - constant_initializer = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(), - [&initializer_name](const NodeArg* input) { - return input->Name() == initializer_name; - }); + bool is_constant = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(), + [&initializer_name](const NodeArg* input) { + return input->Name() == initializer_name; + }); + + if (!is_constant) { + initializer = nullptr; + } } } else if (check_outer_scope && graph.IsSubgraph()) { - constant_initializer = IsConstantInitializer(*graph.ParentGraph(), initializer_name, check_outer_scope); + initializer = GetConstantInitializer(*graph.ParentGraph(), initializer_name); } - return constant_initializer; + return initializer; +} + +bool IsConstantInitializer(const Graph& graph, const std::string& initializer_name, bool check_outer_scope) { + const onnx::TensorProto* initializer = GetConstantInitializer(graph, initializer_name, check_outer_scope); + return initializer != nullptr; } bool NodeArgIsConstant(const Graph& graph, const NodeArg& node_arg) { return IsConstantInitializer(graph, node_arg.Name(), true); } -bool AllNodeInputsAreConstant(const Graph& graph, const Node& node) { +bool AllNodeInputsAreConstant(const Graph& graph, const Node& node, InitializedTensorSet& constant_inputs) { + // clear so we have a known state. if we fail part way through we go back to this state. + constant_inputs.clear(); + + // only initializers can be constant, and there's no edge from a node to an initializer + // so the input edges count must be 0 if (node.GetInputEdgesCount() > 0) { return false; } + for (const auto* input_def : node.InputDefs()) { // Important note: when an initializer appears in the graph's input, this input will not be considered constant, - // because it can be overriden by the user at runtime. For constant folding to be applied, the initializer should + // because it can be overridden by the user at runtime. For constant folding to be applied, the initializer should // not appear in the graph's inputs (that is the only way to guarantee it will always be constant). - if (!NodeArgIsConstant(graph, *input_def)) { + const ONNX_NAMESPACE::TensorProto* initializer = GetConstantInitializer(graph, input_def->Name(), true); + if (initializer) { + constant_inputs.insert({input_def->Name(), initializer}); + } else { + constant_inputs.clear(); return false; } } + return true; } +void ReplaceInitializer(Graph& graph, const std::string& original_name, const ONNX_NAMESPACE::TensorProto& initializer, + bool check_outer_scope) { + ORT_ENFORCE(ReplaceInitializerImpl(graph, original_name, initializer, check_outer_scope), + "Failed to replace initializer. Original initializer was not found. Name:", original_name); +} + size_t RemoveNodeOutputEdges(Graph& graph, Node& node) { std::vector output_edges = GetNodeOutputEdges(node); RemoveGraphEdges(graph, output_edges); diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index a3be17f266..9d60e70fbd 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -39,16 +39,31 @@ bool IsOutputUsed(const Node& node, int index); bool IsGraphInput(const Graph& graph, const NodeArg* input); /** returns true if 'name' is an initializer, and is constant and cannot be overridden at runtime. -@param check_outer_scope If true and the graph is a subgraph, check parent graph/s for 'name' if not found in 'graph'. +@param check_outer_scope If true and the graph is a subgraph, check ancestor graph/s for 'name' if not found in 'graph'. */ -bool IsConstantInitializer(const Graph& graph, const std::string& name, bool check_outer_scope); +bool IsConstantInitializer(const Graph& graph, const std::string& name, bool check_outer_scope = true); -/** Checks if the given node has only constant inputs (initializers). */ -bool AllNodeInputsAreConstant(const Graph& graph, const Node& node); +/** returns the initializer's TensorProto if 'name' is an initializer, and is constant and +cannot be overridden at runtime. If the initializer is not found or is not constant a nullptr is returned. +@param check_outer_scope If true and the graph is a subgraph, check ancestor graph/s for 'name' if not found in 'graph'. +*/ +const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const Graph& graph, const std::string& name, + bool check_outer_scope = true); + +/** Find the initializer called 'original_name' in 'graph', or its ancestors if check_outer_scope is true, + and replace with 'initializer' in the current graph. + Does NOT look in any subgraphs. Requires original_name to match an initializer. + */ +void ReplaceInitializer(Graph& graph, const std::string& original_name, const ONNX_NAMESPACE::TensorProto& initializer, + bool check_outer_scope = true); /** Checks if the given NodeArg is constant, i.e., it appears in the graph's initializers but not in its inputs. */ bool NodeArgIsConstant(const Graph& graph, const NodeArg& node_arg); +/** Checks if the given node has only constant inputs (initializers) and if so returns them in constant_inputs as they +may come from outer scope. */ +bool AllNodeInputsAreConstant(const Graph& graph, const Node& node, InitializedTensorSet& constant_inputs); + /** Gets the name of the incoming NodeArg with the specified index for the given node. */ const std::string& GetNodeInputName(const Node& node, int index); @@ -71,9 +86,6 @@ bool GetRepeatedNodeAttributeValues(const Node& node, return false; } -Status ForAllMutableSubgraphs(Graph& main_graph, std::function func); -Status ForAllSubgraphs(const Graph& main_graph, std::function func); - /** Removes the given Node from the Graph and keeps Graph consistent by rebuilding needed connections. We support the removal of the Node as long as the following conditions hold: - There should be no implicit inputs. diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 96684e3ed0..ff7882e974 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -23,23 +23,24 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level) ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level)); + InitializedTensorSet constant_inputs; + // Check if constant folding can be applied on this node. if (!graph_utils::IsSupportedProvider(*node, GetCompatibleExecutionProviders()) || excluded_op_types_.find(node->OpType()) != excluded_op_types_.end() || - // constant folding is not currently supported for nodes that include subgraphs (control flow operators, - // such as If/Loop/Scan, fall into this category). + // constant folding does not support executing a node that includes subgraphs (control flow operators, + // such as If/Loop/Scan, fall into this category). individual nodes in the subgraph will be processed + // by the Recurse call above node->ContainsSubgraph() || // if the node output is in the graph output, we will get a graph with no nodes. // TODO check if this is allowed in ONNX and ORT. graph.IsNodeOutputsInGraphOutputs(*node) || - !graph_utils::AllNodeInputsAreConstant(graph, *node)) { + !graph_utils::AllNodeInputsAreConstant(graph, *node, constant_inputs)) { continue; } // Create execution frame for executing constant nodes. - // NOTE: As we call AllNodeInputsAreConstant we can use the full list of initializers from - // graph.GetAllInitializedTensors() without filtering out overridable (i.e. non-constant) initializers - OptimizerExecutionFrame::Info info({node}, graph.GetAllInitializedTensors()); + OptimizerExecutionFrame::Info info({node}, constant_inputs); std::vector fetch_mlvalue_idxs; for (const auto* node_out : node->OutputDefs()) { diff --git a/onnxruntime/core/optimizer/conv_add_fusion.cc b/onnxruntime/core/optimizer/conv_add_fusion.cc index 94296e908f..64c468e98f 100644 --- a/onnxruntime/core/optimizer/conv_add_fusion.cc +++ b/onnxruntime/core/optimizer/conv_add_fusion.cc @@ -15,19 +15,14 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie const auto& conv_inputs = conv_node.InputDefs(); const auto& add_inputs = add_node.InputDefs(); - const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) { - return Status::OK(); - } + const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name()); + ORT_ENFORCE(conv_W_tensor_proto); - const ONNX_NAMESPACE::TensorProto* add_B_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(add_inputs[1]->Name(), add_B_tensor_proto)) { - return Status::OK(); - } + const auto* add_B_tensor_proto = graph_utils::GetConstantInitializer(graph, add_inputs[1]->Name()); + ORT_ENFORCE(add_B_tensor_proto); // Currently, fusion is only supported for float or double data type. - if (!Initializer::IsSupportedDataType(add_B_tensor_proto) || - conv_W_tensor_proto->dims_size() < 4) { + if (!Initializer::IsSupportedDataType(add_B_tensor_proto) || conv_W_tensor_proto->dims_size() < 4) { return Status::OK(); } @@ -51,11 +46,9 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie } } - const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr; if (conv_inputs.size() == 3) { - if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto)) { - return Status::OK(); - } + const auto* conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name()); + ORT_ENFORCE(conv_B_tensor_proto); if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) || conv_B_tensor_proto->data_type() != add_B_tensor_proto->data_type() || @@ -78,8 +71,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie conv_B->ToProto(&new_conv_B_tensor_proto); // Replace initializers of conv node - graph.RemoveInitializedTensor(conv_inputs[2]->Name()); - graph.AddInitializedTensor(new_conv_B_tensor_proto); + graph_utils::ReplaceInitializer(graph, conv_inputs[2]->Name(), new_conv_B_tensor_proto); } else { NodeArg* add_B_node_arg = graph.GetNodeArg(add_B_tensor_proto->name()); if (add_B_node_arg == nullptr) { @@ -92,8 +84,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie new_conv_B_tensor_proto.clear_dims(); new_conv_B_tensor_proto.add_dims(dim); - graph.RemoveInitializedTensor(add_B_tensor_proto->name()); - graph.AddInitializedTensor(new_conv_B_tensor_proto); + graph_utils::ReplaceInitializer(graph, add_B_tensor_proto->name(), new_conv_B_tensor_proto); // Update shape of NodeArg TensorShapeProto shape; diff --git a/onnxruntime/core/optimizer/conv_bn_fusion.cc b/onnxruntime/core/optimizer/conv_bn_fusion.cc index 65469b365d..9dea7bed86 100644 --- a/onnxruntime/core/optimizer/conv_bn_fusion.cc +++ b/onnxruntime/core/optimizer/conv_bn_fusion.cc @@ -23,31 +23,21 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff // Get initializers of BatchNormalization const auto& bn_inputs = bn_node.InputDefs(); - const ONNX_NAMESPACE::TensorProto* bn_scale_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(bn_inputs[1]->Name(), bn_scale_tensor_proto)) { - return Status::OK(); - } + const auto* bn_scale_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[1]->Name()); + ORT_ENFORCE(bn_scale_tensor_proto); - const ONNX_NAMESPACE::TensorProto* bn_B_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(bn_inputs[2]->Name(), bn_B_tensor_proto)) { - return Status::OK(); - } + const auto* bn_B_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[2]->Name()); + ORT_ENFORCE(bn_B_tensor_proto); - const ONNX_NAMESPACE::TensorProto* bn_mean_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(bn_inputs[3]->Name(), bn_mean_tensor_proto)) { - return Status::OK(); - } + const auto* bn_mean_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[3]->Name()); + ORT_ENFORCE(bn_mean_tensor_proto); - const ONNX_NAMESPACE::TensorProto* bn_var_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(bn_inputs[4]->Name(), bn_var_tensor_proto)) { - return Status::OK(); - } + const auto* bn_var_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[4]->Name()); + ORT_ENFORCE(bn_var_tensor_proto); const auto& conv_inputs = conv_node.InputDefs(); - const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) { - return Status::OK(); - } + const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name()); + ORT_ENFORCE(conv_W_tensor_proto); // Currently, fusion is only supported for float or double data type. if (!Initializer::IsSupportedDataType(bn_scale_tensor_proto) || @@ -76,12 +66,11 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff auto bn_var = std::make_unique(bn_var_tensor_proto); auto conv_W = std::make_unique(conv_W_tensor_proto); - const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr; std::unique_ptr conv_B = nullptr; + const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr; if (conv_inputs.size() == 3) { - if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto)) { - return Status::OK(); - } + conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name()); + ORT_ENFORCE(conv_B_tensor_proto); if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) || conv_B_tensor_proto->dims_size() != 1 || @@ -124,24 +113,23 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff } // Replace initializers of conv node - graph.RemoveInitializedTensor(conv_W_tensor_proto->name()); + graph_utils::ReplaceInitializer(graph, conv_W_tensor_proto->name(), new_conv_W_tensor_proto); + if (conv_inputs.size() == 3) { #ifdef _MSC_VER #pragma warning(push) #pragma warning(disable : 6011) // Not deferencing null pointer. conv_B_tensor_proto is set on line 93 #endif - graph.RemoveInitializedTensor(conv_B_tensor_proto->name()); + graph_utils::ReplaceInitializer(graph, conv_B_tensor_proto->name(), new_conv_B_tensor_proto); #ifdef _MSC_VER #pragma warning(pop) #endif } else { - graph.RemoveInitializedTensor(bn_B_tensor_proto->name()); + graph_utils::ReplaceInitializer(graph, bn_B_tensor_proto->name(), new_conv_B_tensor_proto); conv_node.MutableInputDefs().push_back(bn_B_node_arg); conv_node.MutableInputArgsCount()[2] = 1; } - graph.AddInitializedTensor(new_conv_W_tensor_proto); - graph.AddInitializedTensor(new_conv_B_tensor_proto); // Remove BN node. auto* bn_node_to_remove = graph.GetNode(bn_node.Index()); diff --git a/onnxruntime/core/optimizer/conv_mul_fusion.cc b/onnxruntime/core/optimizer/conv_mul_fusion.cc index 54408f3185..32b4c0bbd4 100644 --- a/onnxruntime/core/optimizer/conv_mul_fusion.cc +++ b/onnxruntime/core/optimizer/conv_mul_fusion.cc @@ -15,15 +15,11 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef const auto& conv_inputs = conv_node.InputDefs(); const auto& mul_inputs = mul_node.InputDefs(); - const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) { - return Status::OK(); - } + const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name()); + ORT_ENFORCE(conv_W_tensor_proto); - const ONNX_NAMESPACE::TensorProto* mul_B_tensor_proto = nullptr; - if (!graph.GetInitializedTensor(mul_inputs[1]->Name(), mul_B_tensor_proto)) { - return Status::OK(); - } + const auto* mul_B_tensor_proto = graph_utils::GetConstantInitializer(graph, mul_inputs[1]->Name()); + ORT_ENFORCE(mul_B_tensor_proto); if (!Initializer::IsSupportedDataType(conv_W_tensor_proto) || !Initializer::IsSupportedDataType(mul_B_tensor_proto) || @@ -61,10 +57,9 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef std::unique_ptr conv_B = nullptr; const bool is_3d = conv_inputs.size() == 3; if (is_3d) { - if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto)) - return Status::OK(); - if (conv_B_tensor_proto == nullptr) - return Status(ONNXRUNTIME, FAIL, "Internal error in ConvMulFusion. conv_B_tensor_proto is NULL"); + conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name()); + ORT_ENFORCE(conv_B_tensor_proto); + if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) || conv_B_tensor_proto->data_type() != mul_B_tensor_proto->data_type() || conv_B_tensor_proto->dims_size() != 1 || @@ -90,14 +85,12 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef conv_W->ToProto(&new_conv_W_tensor_proto); // Replace initializers of conv node - graph.RemoveInitializedTensor(conv_inputs[1]->Name()); - graph.AddInitializedTensor(new_conv_W_tensor_proto); + graph_utils::ReplaceInitializer(graph, conv_inputs[1]->Name(), new_conv_W_tensor_proto); if (is_3d) { ONNX_NAMESPACE::TensorProto new_conv_B_tensor_proto(*conv_B_tensor_proto); conv_B->ToProto(&new_conv_B_tensor_proto); - graph.RemoveInitializedTensor(conv_inputs[2]->Name()); - graph.AddInitializedTensor(new_conv_B_tensor_proto); + graph_utils::ReplaceInitializer(graph, conv_inputs[2]->Name(), new_conv_B_tensor_proto); } // Remove Mul node. diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.cc b/onnxruntime/core/optimizer/unsqueeze_elimination.cc index c35d078546..c031fe143d 100644 --- a/onnxruntime/core/optimizer/unsqueeze_elimination.cc +++ b/onnxruntime/core/optimizer/unsqueeze_elimination.cc @@ -25,11 +25,9 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& // Generate new dims. NodeArg* input_def = node.MutableInputDefs()[0]; - const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr; - graph.GetInitializedTensor(input_def->Name(), tensor_proto); - if (tensor_proto == nullptr) { - return Status::OK(); - } + const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, input_def->Name()); + ORT_ENFORCE(tensor_proto); + std::vector new_dims(axes.size() + tensor_proto->dims().size(), 0); if (new_dims.size() >= std::numeric_limits::max()) { return Status(ONNXRUNTIME, FAIL, "index out of range"); @@ -56,8 +54,13 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& new_tensor_proto.add_dims(new_dims[i]); } } - graph.RemoveInitializedTensor(input_def->Name()); - graph.AddInitializedTensor(new_tensor_proto); + + // TODO: This seems wrong as there's no check whether another node is using the initializer before replacing it. + // Shouldn't we check that or alternatively create an initializer with a different name and let + // Graph::CleanUnusedInitializers remove the original one if nothing else consumes it? + // graph.RemoveInitializedTensor(input_def->Name()); + // graph.AddInitializedTensor(new_tensor_proto); + graph_utils::ReplaceInitializer(graph, input_def->Name(), new_tensor_proto); // Update shape of NodeArg. TensorShapeProto shape; @@ -75,8 +78,8 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& } // namespace onnxruntime bool UnsqueezeElimination::SatisfyCondition(const Graph& graph, const Node& node) const { - // Attempt to remove an Unsqueeze operator only if it gets an initializer as input. - return node.GetInputEdgesCount() == 0 && + // Attempt to remove an Unsqueeze operator only if it gets a constant initializer as input. + return graph_utils::IsConstantInitializer(graph, node.InputDefs()[0]->Name()) && !graph.IsNodeOutputsInGraphOutputs(node); } diff --git a/onnxruntime/test/framework/session_state_test.cc b/onnxruntime/test/framework/session_state_test.cc index b2e422c69e..6064b74630 100644 --- a/onnxruntime/test/framework/session_state_test.cc +++ b/onnxruntime/test/framework/session_state_test.cc @@ -125,7 +125,7 @@ TEST(SessionStateTest, TestInitializerProcessing) { bool found = initialized_tensors.find(idx) != initialized_tensors.cend(); ASSERT_TRUE(found) << "Missing entry for " << entry.first << " in session state initialized tensors"; - if (graph_utils::IsConstantInitializer(graph, entry.first, true)) { + if (graph_utils::IsConstantInitializer(graph, entry.first, false)) { found = const_initialized_tensors.find(idx) != const_initialized_tensors.cend(); ASSERT_TRUE(found) << "Missing entry for " << entry.first << " in session state const initialized tensors"; } diff --git a/onnxruntime/test/framework/test_utils.cc b/onnxruntime/test/framework/test_utils.cc index 5d630b832a..7604e15ce7 100644 --- a/onnxruntime/test/framework/test_utils.cc +++ b/onnxruntime/test/framework/test_utils.cc @@ -42,15 +42,30 @@ IExecutionProvider* TestNnapiExecutionProvider() { } #endif +static void CountOpsInGraphImpl(const Graph& graph, std::map& ops) { + for (auto& node : graph.Nodes()) { + auto pos = ops.find(node.OpType()); + if (pos == ops.end()) { + ops[node.OpType()] = 1; + } else { + ++pos->second; + } + + if (node.ContainsSubgraph()) { + for (auto& subgraph : node.GetSubgraphs()) { + CountOpsInGraphImpl(*subgraph, ops); + } + } + } +} + // Returns a map with the number of occurrences of each operator in the graph. // Helper function to check that the graph transformations have been successfully applied. std::map CountOpsInGraph(const Graph& graph) { - std::map op_to_count; - for (auto& node : graph.Nodes()) { - op_to_count[node.OpType()] = - op_to_count.count(node.OpType()) == 0 ? 1 : ++op_to_count[node.OpType()]; - } - return op_to_count; + std::map ops; + CountOpsInGraphImpl(graph, ops); + + return ops; } } // namespace test diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 32b439a55f..aed983207a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -115,6 +115,80 @@ TEST(GraphTransformationTests, ConstantFolding) { ASSERT_TRUE(op_to_count["Unsqueeze"] == 0); } +TEST(GraphTransformationTests, ConstantFoldingSubgraph) { + TensorProto value_tensor; + value_tensor.add_dims(1); + value_tensor.add_float_data(1.f); + value_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + + TypeProto float_tensor_type; + float_tensor_type.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + float_tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto create_subgraph = [&](GraphProto& graph_proto) { + // create subgraph that has an Add node to add a local and parent graph initializer + Model model("ConstantFoldingSubgraphTest_subgraph"); + auto& graph = model.MainGraph(); + + TensorProto local_constant(value_tensor); + local_constant.set_name("local_constant"); + graph.AddInitializedTensor(local_constant); + + auto& local_constant_arg = graph.GetOrCreateNodeArg("local_constant", &float_tensor_type); + auto& parent_constant_arg = graph.GetOrCreateNodeArg("parent_constant", &float_tensor_type); + graph.AddOuterScopeNodeArg("parent_constant"); + + auto& add_out = graph.GetOrCreateNodeArg("add_out", &float_tensor_type); + graph.AddNode("add", "Add", "Add two inputs.", {&parent_constant_arg, &local_constant_arg}, {&add_out}); + + auto& subgraph_out = graph.GetOrCreateNodeArg("subgraph_out", &float_tensor_type); + graph.AddNode("identity", "Identity", "So Add isn't providing graph output.", {&add_out}, {&subgraph_out}); + + auto status = graph.Resolve(); + ASSERT_TRUE(status.IsOK()) << status; + graph_proto = graph.ToGraphProto(); + }; + + Model model("ConstantFoldingSubgraphTest_main_graph"); + auto& graph = model.MainGraph(); + + // add initializer at parent level + TensorProto parent_value_tensor(value_tensor); + parent_value_tensor.set_name("parent_constant"); + graph.AddInitializedTensor(parent_value_tensor); + + // put the subgraph in an If node + TypeProto if_cond_type; + if_cond_type.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + if_cond_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + auto& if_cond_input = graph.GetOrCreateNodeArg("if_in", &if_cond_type); + auto& if_output = graph.GetOrCreateNodeArg("if_out", &float_tensor_type); + + auto& if_node = graph.AddNode("if", "If", "If node", {&if_cond_input}, {&if_output}); + + GraphProto subgraph; + create_subgraph(subgraph); + + if_node.AddAttribute("then_branch", {subgraph}); + if_node.AddAttribute("else_branch", {subgraph}); + + auto status = graph.Resolve(); + ASSERT_TRUE(status.IsOK()) << status; + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Add"] == 2); // one in each subgraph + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level1); + + status = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1); + ASSERT_TRUE(status.IsOK()) << status; + + op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Add"] == 0) + << "Constant folding should have been able to remove the Add node in both subgraphs"; +} + TEST(GraphTransformationTests, ShapeToInitializer) { string model_uri = MODEL_FOLDER + "shape-add.onnx"; std::shared_ptr model; @@ -308,9 +382,10 @@ TEST(GraphTransformationTests, NegativeFuseConvAddNoBias) { ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2).IsOK()); // Nodes are not fused because the weights to conv/add are not constants (they appear in the graph inputs). + // Unsqueeze is also not eliminated as the initializer that is its input is also not constant std::map op_to_count = CountOpsInGraph(graph); ASSERT_TRUE(op_to_count["Add"] != 0); - ASSERT_TRUE(op_to_count["Unsqueeze"] == 0); + ASSERT_TRUE(op_to_count["Unsqueeze"] != 0); } TEST(GraphTransformationTests, FuseConvAddMul3D) {