From 512cfdd9feba54bd7cad6a7b9033ac707327fd3f Mon Sep 17 00:00:00 2001 From: Konstantinos Karanasos Date: Wed, 3 Apr 2019 22:34:20 -0700 Subject: [PATCH] Generalize node removal (#743) Generalize node removal method in graph_utils. This is a higher-level method that keeps the graph consistent so that no Resolve is needed after the removal of a node. The new method supports the removal of nodes with a single input (be it an incoming node or an initializer) and a single output (but allowing multiple output edges of that output). It also takes into account the case that one of the output edges is fed to a subgraph. Also updated the rewrite rules to use this new, less restrictive method, and improved the rules' conditions. Introduced a GraphEdge struct to simplify various methods in graph_utils. --- include/onnxruntime/core/graph/graph.h | 10 +- onnxruntime/core/graph/graph.cc | 21 +- onnxruntime/core/graph/graph_utils.cc | 412 +++++++++++------- onnxruntime/core/graph/graph_utils.h | 17 +- .../core/optimizer/identity_elimination.cc | 8 +- .../core/optimizer/slice_elimination.cc | 13 +- .../core/optimizer/unsqueeze_elimination.cc | 32 +- onnxruntime/test/ir/utils_test.cc | 57 ++- 8 files changed, 378 insertions(+), 192 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 9f48e2c732..8c5e6629b2 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -156,6 +156,11 @@ class Node { return definitions_.implicit_input_defs; } + /** Gets a modifiable collection of the Node's implicit input definitions. */ + std::vector& MutableImplicitInputDefs() noexcept { + return definitions_.implicit_input_defs; + } + /** Gets the Node's output definitions. @remarks requires ConstPointerContainer wrapper to apply const to the NodeArg pointers so access is read-only. */ const ConstPointerContainer> OutputDefs() const noexcept { @@ -744,6 +749,9 @@ class Graph { graph_output_order_ = outputs; } + /** Returns true if this is a subgraph or fase if it is a high-level graph. */ + bool IsSubgraph() const { return parent_graph_ != nullptr; } + /** 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. @@ -905,8 +913,6 @@ class Graph { std::vector CreateNodeArgs(const google::protobuf::RepeatedPtrField& names, const ArgNameToTypeMap& name_to_type_map); - bool IsSubgraph() const { return parent_graph_ != nullptr; } - void AddFunction(const ONNX_NAMESPACE::FunctionProto* func_proto); // GraphProto to store name, version, initializer. diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index e60a042a23..23a3071264 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -869,7 +869,13 @@ void Graph::RemoveEdge(NodeIndex src_node_index, NodeIndex dst_node_index, int s if (nullptr == dst_arg) { ORT_THROW("Invalid destination node arg slot specified when removing edge."); } - + + if (src_arg != dst_arg) { + // The edge ends specified by source and destination arg slot are not referring to same node arg. + // It means there was no edge between these two slots before. + ORT_THROW("Argument type mismatch when removing edge."); + } + nodes_[dst_node_index]->MutableRelationships().input_edges.erase(Node::EdgeEnd(*nodes_[src_node_index], src_arg_slot, dst_arg_slot)); nodes_[src_node_index]->MutableRelationships().output_edges.erase(Node::EdgeEnd(*nodes_[dst_node_index], src_arg_slot, dst_arg_slot)); } @@ -881,6 +887,8 @@ Status Graph::BuildConnections(std::vector& outer_scope_node_args_c // recurse into subgraphs first so we can update any nodes in this graph that are used by those subgraphs if (!resolve_context_.nodes_with_subgraphs.empty()) { + const bool loaded_from_model_file = GraphLoadedFromModelFile(graph_proto_); + for (auto* node : resolve_context_.nodes_with_subgraphs) { for (auto& subgraph : node->MutableSubgraphs()) { std::vector node_args_consumed; @@ -932,6 +940,17 @@ Status Graph::BuildConnections(std::vector& outer_scope_node_args_c AddEdge(output_node.Index(), node->Index(), entry->second.second, input_slot_index); inner_nodes.insert(&output_node); + + // If this Graph was built manually, remove the implicit input from the graph outputs if it is present there + // and not explicitly listed in the ordered graph outputs (as that implies we should leave it as an output). + // If the Graph was loaded from a GraphProto, honor the explicit graph outputs and leave as is. + if (!loaded_from_model_file) { + auto in_ordered_graph_outputs = find(graph_output_order_.cbegin(), graph_output_order_.cend(), node_arg); + if (in_ordered_graph_outputs == graph_output_order_.cend()) { + graph_outputs_.erase(std::remove(graph_outputs_.begin(), graph_outputs_.end(), node_arg), + graph_outputs_.end()); + } + } } } } diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index ae01720d84..f82962174d 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -4,28 +4,264 @@ #include "core/graph/graph_utils.h" #include "core/graph/graph.h" #include "core/framework/tensorprotoutils.h" +#include "core/common/logging/logging.h" namespace onnxruntime { namespace graph_utils { -// local helpers +struct GraphEdge { + NodeIndex src_node; + NodeIndex dst_node; + int src_arg_index; + int dst_arg_index; + std::string arg_name; + + GraphEdge(NodeIndex src_node, NodeIndex dst_node, + int src_arg_index, int dst_arg_index, const std::string& arg_name) : src_node(src_node), + dst_node(dst_node), + src_arg_index(src_arg_index), + dst_arg_index(dst_arg_index), + arg_name(arg_name) {} + + // Constructs a GraphEdge given a node, an edge_end, and a boolean for the edge direction. + static GraphEdge CreateGraphEdge(const Node& node, const Node::EdgeEnd& edge_end, bool is_input_edge) { + return is_input_edge + ? GraphEdge(edge_end.GetNode().Index(), + node.Index(), + edge_end.GetSrcArgIndex(), + edge_end.GetDstArgIndex(), + GetNodeInputName(node, edge_end.GetDstArgIndex())) + : GraphEdge(node.Index(), + edge_end.GetNode().Index(), + edge_end.GetSrcArgIndex(), + edge_end.GetDstArgIndex(), + GetNodeOutputName(node, edge_end.GetSrcArgIndex())); + } +}; + +//--------------------- +//--- local helpers --- +//--------------------- // check if an output edge provides an implicit input to the destination node -static bool OutputEdgeProvidesImplicitInput(const Node::EdgeEnd& output_edge) { +static bool OutputEdgeProvidesImplicitInput(const Graph& graph, const GraphEdge& output_edge) { // we treat the explicit and implicit inputs as sequential, so if the destination arg index of an output edge // is past the valid range for the node's explicit inputs, it is for an implicit input - const auto num_explicit_inputs = output_edge.GetNode().InputDefs().size(); - bool is_implicit_input = output_edge.GetDstArgIndex() >= num_explicit_inputs; + const auto num_explicit_inputs = (*graph.GetNode(output_edge.dst_node)).InputDefs().size(); + bool is_implicit_input = output_edge.dst_arg_index >= num_explicit_inputs; return is_implicit_input; } -static const std::string& GetNodeOutputName(const Node& node, int index) { +/** 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, + 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 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) { + return false; + } + + for (auto& subgraph_node : attr_subgraph_pair.second->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(); + if (!subgraph_node_implicit_inputs.empty()) { + auto subgraph_node_also_consumes_nodearg_as_implicit_input = + std::find_if(subgraph_node_implicit_inputs.cbegin(), subgraph_node_implicit_inputs.cend(), + [&removed_output_name](const NodeArg* input) { + return input != nullptr && input->Name() == removed_output_name; + }); + + if (subgraph_node_also_consumes_nodearg_as_implicit_input != subgraph_node_implicit_inputs.cend()) { + if (!CanUpdateImplicitInputNameInSubgraph(subgraph_node, removed_output_name, new_output_name)) + return false; + } + } + } + } + + return true; +} + +/** Updates removed_output_name with new_output_name in the subgraph input. */ +static void UpdateImplicitInputNameInSubgraph(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; + + 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(); + if (!subgraph_node_implicit_inputs.empty()) { + auto subgraph_node_also_consumes_nodearg_as_implicit_input = + std::find_if(subgraph_node_implicit_inputs.cbegin(), subgraph_node_implicit_inputs.cend(), + [&removed_output_name](const NodeArg* input) { + return input->Name() == removed_output_name; + }); + + if (subgraph_node_also_consumes_nodearg_as_implicit_input != subgraph_node_implicit_inputs.cend()) { + UpdateImplicitInputNameInSubgraph(subgraph_node, removed_output_name, new_output_name); + } + } + + // Need mutable input defs to be able to update the implicit input names + auto& input_args = subgraph_node.MutableInputDefs(); + + if (!input_args.empty()) { + int input_slot_index = -1; + for (const auto* input_arg : input_args) { + ++input_slot_index; + // if the input matches, replace the NodeArg with one using the new name + if (input_arg->Exists() && input_arg->Name() == removed_output_name) { + // sanity check there was no edge for this input. implicit inputs from outer scope do not have edges + ORT_ENFORCE(std::count_if(subgraph_node.InputEdgesBegin(), subgraph_node.InputEdgesEnd(), + [input_slot_index](const Node::EdgeEnd& entry) { + return entry.GetDstArgIndex() == input_slot_index; + }) == 0); + + // Create a new NodeArg with the new name + input_args[input_slot_index] = &attr_subgraph_pair.second->GetOrCreateNodeArg(new_output_name, + input_arg->TypeAsProto()); + } + } + } + } + } +} + +/** Returns a vector of the output GraphEdges of a node. */ +static std::vector GetNodeOutputEdges(const Node& node) { + std::vector output_edges; + for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) { + output_edges.push_back(GraphEdge::CreateGraphEdge(node, *it, false)); + } + + return output_edges; +} + +/** Removes a set of GraphEdges from the graph. */ +static void RemoveGraphEdges(Graph& graph, const std::vector& edges) { + for (const auto& edge_to_remove : edges) { + graph.RemoveEdge(edge_to_remove.src_node, + edge_to_remove.dst_node, + edge_to_remove.src_arg_index, + edge_to_remove.dst_arg_index); + } +} + +/** Given a graph, a list of edges, and a NodeArg name, checks if each of the edges provides an implicit input + to a subgraph. If so, it checks if there is no clash of the given NodeArg name in each of the subgraphs. + This is important when removing a node with this NodeArg as input. */ +bool CanUpdateImplicitInputNameInSubgraphs(Graph& graph, + const std::vector& output_edges, + const std::string& input_arg_name) { + for (const auto& output_edge : output_edges) { + if (OutputEdgeProvidesImplicitInput(graph, output_edge)) { + Node& mutable_output_edge_node = *graph.GetNode(output_edge.dst_node); + if (!CanUpdateImplicitInputNameInSubgraph(mutable_output_edge_node, output_edge.arg_name, input_arg_name)) { + LOGS_DEFAULT(WARNING) << " Implicit input name " << input_arg_name + << " cannot be safely updated in one of the subgraphs."; + return false; + } + } + } + return true; +} + +/** Removes a node with a single incoming node. */ +static bool RemoveNodeWithSingleNodeIn(Graph& graph, Node& node) { + // Store info for input and output edges. + std::vector output_edges = GetNodeOutputEdges(node); + const Node::EdgeEnd& input_edge_end = *node.InputEdgesBegin(); + const GraphEdge input_edge = GraphEdge::CreateGraphEdge(node, input_edge_end, true); + + // Check that the incoming NodeArg can be safely used in the presence of subgraphs. + if (!CanUpdateImplicitInputNameInSubgraphs(graph, output_edges, input_edge.arg_name)) { + return false; + } + + // Remove the output edges of the node and then the node itself (this will remove its input edge too). + RemoveGraphEdges(graph, output_edges); + graph.RemoveNode(node.Index()); + + // Create connections between the incoming node and the outgoing nodes of the node that we removed. + for (const auto& output_edge : output_edges) { + // Take care of subgraph inputs. + if (OutputEdgeProvidesImplicitInput(graph, output_edge)) { + Node& mutable_output_edge_node = *graph.GetNode(output_edge.dst_node); + UpdateImplicitInputNameInSubgraph(mutable_output_edge_node, output_edge.arg_name, input_edge.arg_name); + } + + // Add new edge connecting the input with the output nodes directly. + graph.AddEdge(input_edge.src_node, output_edge.dst_node, input_edge.src_arg_index, output_edge.dst_arg_index); + } + + return true; +} + +/** Remove a node with a single incoming initializer (and no other incoming node). */ +static bool RemoveNodeWithSingleInitializerIn(Graph& graph, Node& node) { + // Store info for input initializer and output edges. + auto* input_def = node.MutableInputDefs()[0]; + std::vector output_edges = GetNodeOutputEdges(node); + + // Check that the incoming NodeArg can be safely used in the presence of subgraphs. + if (!CanUpdateImplicitInputNameInSubgraphs(graph, output_edges, input_def->Name())) { + return false; + } + + // Remove the output edges of the node and then the node itself (this will remove its input edge too). + RemoveGraphEdges(graph, output_edges); + graph.RemoveNode(node.Index()); + + // Add the incoming initializer as input to the outgoing nodes of the node that we removed. + for (auto& output_edge : output_edges) { + // Take care of subgraph inputs. + if (OutputEdgeProvidesImplicitInput(graph, output_edge)) { + Node& mutable_output_edge_node = *graph.GetNode(output_edge.dst_node); + UpdateImplicitInputNameInSubgraph(mutable_output_edge_node, output_edge.arg_name, input_def->Name()); + } + + // Replace outgoing node's input to use the initializer. + auto output_node = graph.GetNode(output_edge.dst_node); + ORT_ENFORCE(output_node, "Outgoing node could not be found."); + + auto dst_arg_idx = output_edge.dst_arg_index; + if (dst_arg_idx < output_node->InputDefs().size()) { + output_node->MutableInputDefs()[output_edge.dst_arg_index] = input_def; + } else if (dst_arg_idx < output_node->InputDefs().size() + output_node->ImplicitInputDefs().size()) { + // If we need to update an implicit input. + output_node->MutableImplicitInputDefs()[dst_arg_idx - output_node->InputDefs().size()] = input_def; + } else { + LOGS_DEFAULT(ERROR) << " Invalid value for input index of node " << output_node->Name(); + return false; + } + } + + return true; +} + +//---------------------------- +//--- end of local helpers --- +//---------------------------- + +const std::string& GetNodeInputName(const Node& node, int index) { + const auto& inputs = node.InputDefs(); + ORT_ENFORCE(index < inputs.size(), "Attempting to get an input that does not exist."); + return inputs[index]->Name(); +} + +const std::string& GetNodeOutputName(const Node& node, int index) { const auto& outputs = node.OutputDefs(); - - // this should never happen as it's internal logic so just use an assert assert(index < outputs.size()); - return outputs[index]->Name(); } @@ -44,11 +280,11 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node, bool IsSupportedProvider(const Node& node, const std::unordered_set& compatible_providers) { - if (!compatible_providers.empty() && + if (!compatible_providers.empty() && compatible_providers.find(node.GetExecutionProviderType()) == compatible_providers.end()) { return false; } - + return true; } @@ -99,7 +335,9 @@ Status ForAllSubgraphs(const Graph& graph, std::function f } bool IsSingleInSingleOutNode(const Node& node) { - return node.GetInputEdgesCount() == 1 && node.GetOutputEdgesCount() == 1; + return node.InputDefs().size() == 1 && + node.ImplicitInputDefs().size() == 0 && + node.OutputDefs().size() == 1; } const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const std::string& attr_name) { @@ -108,132 +346,20 @@ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const s return iter == attrs.end() ? nullptr : &iter->second; } -// check if new_output_name can be used to replace removed_output_name -// if there is an existing NodeArg in a subgraph that implicitly consumes removed_output_name, it is not safe. -static bool CanUpdateImplicitInputNameInSubgraph(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 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) { - return false; - } - - for (auto& subgraph_node : attr_subgraph_pair.second->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(); - if (!subgraph_node_implicit_inputs.empty()) { - auto subgraph_node_also_consumes_nodearg_as_implicit_input = - std::find_if(subgraph_node_implicit_inputs.cbegin(), subgraph_node_implicit_inputs.cend(), - [&removed_output_name](const NodeArg* input) { - return input != nullptr && input->Name() == removed_output_name; - }); - - if (subgraph_node_also_consumes_nodearg_as_implicit_input != subgraph_node_implicit_inputs.cend()) { - if (!CanUpdateImplicitInputNameInSubgraph(subgraph_node, removed_output_name, new_output_name)) - return false; - } - } - } - } - - return true; -} - -static void UpdateImplicitInputNameInSubgraph(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; - - 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(); - if (!subgraph_node_implicit_inputs.empty()) { - auto subgraph_node_also_consumes_nodearg_as_implicit_input = - std::find_if(subgraph_node_implicit_inputs.cbegin(), subgraph_node_implicit_inputs.cend(), - [&removed_output_name](const NodeArg* input) { - return input->Name() == removed_output_name; - }); - - if (subgraph_node_also_consumes_nodearg_as_implicit_input != subgraph_node_implicit_inputs.cend()) { - UpdateImplicitInputNameInSubgraph(subgraph_node, removed_output_name, new_output_name); - } - } - - // Need mutable input defs to be able to update the implicit input names - auto& input_args = subgraph_node.MutableInputDefs(); - - if (!input_args.empty()) { - int input_slot_index = -1; - for (const auto* input_arg : input_args) { - ++input_slot_index; - // if the input matches, replace the NodeArg with one using the new name - if (input_arg->Exists() && input_arg->Name() == removed_output_name) { - // sanity check there was no edge for this input. implicit inputs from outer scope do not have edges - ORT_ENFORCE(std::count_if(subgraph_node.InputEdgesBegin(), subgraph_node.InputEdgesEnd(), - [input_slot_index](const Node::EdgeEnd& entry) { - return entry.GetDstArgIndex() == input_slot_index; - }) == 0); - - // Create a new NodeArg with the new name - input_args[input_slot_index] = &attr_subgraph_pair.second->GetOrCreateNodeArg(new_output_name, - input_arg->TypeAsProto()); - } - } - } - } - } -} - -bool RemoveSingleInSingleOutNode(Graph& graph, Node& node) { - if (!IsSingleInSingleOutNode(node)) { +bool RemoveSingleInputNode(Graph& graph, Node& node) { + // Cannot remove a node with multiple output NodeArgs (multiple output edges is fine), neither + // a node whose output is also a graph output. + if (!IsSingleInSingleOutNode(node) || + graph.IsNodeOutputsInGraphOutputs(node)) { return false; } - // Get input/output edges, nodes, and node args. - const Node::EdgeEnd& input_edge = *node.InputEdgesBegin(); - const Node& input_edge_node = input_edge.GetNode(); - const NodeIndex input_edge_node_index = input_edge_node.Index(); - const int input_edge_dst_arg = input_edge.GetSrcArgIndex(); - - const Node::EdgeEnd& output_edge = *node.OutputEdgesBegin(); - const Node& output_edge_node = output_edge.GetNode(); - const NodeIndex output_edge_node_index = output_edge_node.Index(); - const int output_edge_dst_arg = output_edge.GetDstArgIndex(); - - // check if the output from the node to remove is implicit input to the downstream node. - // if so, it's used in subgraph in that node and we have to check if it's valid to update that subgraph - if (OutputEdgeProvidesImplicitInput(output_edge)) { - // the output we need to wire to the downstream node, and the output name from the node we want to remove - const auto& output_name = GetNodeOutputName(input_edge_node, output_edge.GetSrcArgIndex()); - const auto& removed_output_name = GetNodeOutputName(node, 0); - - Node& mutable_output_edge_node = *graph.GetNode(output_edge_node_index); - - if (CanUpdateImplicitInputNameInSubgraph(mutable_output_edge_node, removed_output_name, output_name)) { - UpdateImplicitInputNameInSubgraph(mutable_output_edge_node, removed_output_name, output_name); - } else { - // we can't safely remove this node - return false; - } + // If the single input comes from another node (initializers are not connected with edges to nodes). + if (node.GetInputEdgesCount() == 1) { + return RemoveNodeWithSingleNodeIn(graph, node); + } else { + return RemoveNodeWithSingleInitializerIn(graph, node); } - - // Remove output edge. - graph.RemoveEdge(node.Index(), output_edge_node_index, output_edge.GetSrcArgIndex(), output_edge_dst_arg); - - // Remove node (this will remove the input edge too). - graph.RemoveNode(node.Index()); - - // Add new edge connecting the input with the output nodes directly. - graph.AddEdge(input_edge_node_index, output_edge_node_index, input_edge_dst_arg, output_edge_dst_arg); - - return true; } bool IsGraphInput(const Graph& graph, const NodeArg* input) { @@ -248,8 +374,8 @@ bool AllNodeInputsAreConstant(const Graph& graph, const Node& node) { const onnx::TensorProto* initializer = nullptr; 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 not - // appear in the graph's inputs (that is the only way to guarantee it will always be constant). + // because it can be overriden 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 (!graph.GetInitializedTensor(input_def->Name(), initializer) || IsGraphInput(graph, input_def)) { return false; } @@ -258,20 +384,10 @@ bool AllNodeInputsAreConstant(const Graph& graph, const Node& node) { } size_t RemoveNodeOutputEdges(Graph& graph, Node& node) { - std::vector> edges_to_remove; - for (auto it = node.OutputEdgesBegin(); it != node.OutputEdgesEnd(); ++it) { - edges_to_remove.emplace_back(std::make_tuple(it->GetNode().Index(), - it->GetSrcArgIndex(), - it->GetDstArgIndex())); - } - for (auto& edge_to_remove : edges_to_remove) { - graph.RemoveEdge(node.Index(), - std::get<0>(edge_to_remove), - std::get<1>(edge_to_remove), - std::get<2>(edge_to_remove)); - } + std::vector output_edges = GetNodeOutputEdges(node); + RemoveGraphEdges(graph, output_edges); - return edges_to_remove.size(); + return output_edges.size(); } } // namespace graph_utils diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index be988af987..15d0e19465 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -9,6 +9,7 @@ namespace onnxruntime { namespace graph_utils { + bool IsSupportedOptypeVersionAndDomain(const Node& node, const std::string& op_type, ONNX_NAMESPACE::OperatorSetVersion version, @@ -20,7 +21,9 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node, bool IsSupportedProvider(const Node& node, const std::unordered_set& compatible_providers); -/** Check whether the node has a single input and a single output. */ +/** Check whether the node has a single input and a single output. The single input can be either the output of + another node or an initializer, but not an implicit input from a parent subgraph. The single output can be + fed to multiple downstream operators, i.e., it can have multiple output edges. */ bool IsSingleInSingleOutNode(const Node& node); /** Returns true if the graph has the given input.*/ @@ -29,6 +32,12 @@ bool IsGraphInput(const Graph& graph, const NodeArg* input); /** Checks if the given node has only constant inputs (initializers). */ bool AllNodeInputsAreConstant(const Graph& graph, const Node& node); +/** Get the name of the incoming NodeArg with the specified index for the given node. */ +const std::string& GetNodeInputName(const Node& node, int index); + +/** Get the name of the outgoing NodeArg with the specified index for the given node. */ +const std::string& GetNodeOutputName(const Node& node, int index); + /** Return the attribute of a Node with a given name. */ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const std::string& attr_name); @@ -49,8 +58,10 @@ bool GetRepeatedNodeAttributeValues(const Node& node, Status ForAllMutableSubgraphs(Graph& main_graph, std::function func); Status ForAllSubgraphs(const Graph& main_graph, std::function func); -/** Remove the given single-input-single-output Node from the Graph. */ -bool RemoveSingleInSingleOutNode(Graph& graph, Node& node); +/** Remove the given single-input Node from the Graph. The single input might be either + another node or an initializer, but not an implicit input. The node should have a single + output but can have multiple output edges. */ +bool RemoveSingleInputNode(Graph& graph, Node& node); /** Remove all output edges from the given Node of the Graph. This should probably be elevated to the Graph API eventually. */ diff --git a/onnxruntime/core/optimizer/identity_elimination.cc b/onnxruntime/core/optimizer/identity_elimination.cc index e039ae0734..da24f53775 100644 --- a/onnxruntime/core/optimizer/identity_elimination.cc +++ b/onnxruntime/core/optimizer/identity_elimination.cc @@ -11,15 +11,17 @@ namespace onnxruntime { Status EliminateIdentity::Apply(Graph& graph, Node& node, bool& modified, bool& deleted) { - if (graph_utils::RemoveSingleInSingleOutNode(graph, node)) { + if (graph_utils::RemoveSingleInputNode(graph, node)) { modified = deleted = true; } return Status::OK(); } -bool EliminateIdentity::SatisfyCondition(const Graph& /*graph*/, const Node& node) { - return node.OpType() == included_op_type_ && graph_utils::IsSingleInSingleOutNode(node); +bool EliminateIdentity::SatisfyCondition(const Graph& graph, const Node& node) { + return node.OpType() == included_op_type_ && + graph_utils::IsSingleInSingleOutNode(node) && + !graph.IsNodeOutputsInGraphOutputs(node); } } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/slice_elimination.cc b/onnxruntime/core/optimizer/slice_elimination.cc index db61efa228..821cf7dc21 100644 --- a/onnxruntime/core/optimizer/slice_elimination.cc +++ b/onnxruntime/core/optimizer/slice_elimination.cc @@ -9,20 +9,17 @@ namespace onnxruntime { Status EliminateSlice::Apply(Graph& graph, Node& node, bool& modified, bool& removed) { - if (graph_utils::RemoveSingleInSingleOutNode(graph, node)) { + if (graph_utils::RemoveSingleInputNode(graph, node)) { removed = modified = true; } return Status::OK(); } -bool EliminateSlice::SatisfyCondition(const Graph& /*graph*/, const Node& node) { - if (node.OpType() != included_op_type_) { - return false; - } - - // At the moment, we eliminate a slice operator only if it has a single input and a single output. - if (!graph_utils::IsSingleInSingleOutNode(node)) { +bool EliminateSlice::SatisfyCondition(const Graph& graph, const Node& node) { + if (node.OpType() != included_op_type_ || + !graph_utils::IsSingleInSingleOutNode(node) || + graph.IsNodeOutputsInGraphOutputs(node)) { return false; } diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.cc b/onnxruntime/core/optimizer/unsqueeze_elimination.cc index 75a7301174..c625a36778 100644 --- a/onnxruntime/core/optimizer/unsqueeze_elimination.cc +++ b/onnxruntime/core/optimizer/unsqueeze_elimination.cc @@ -11,7 +11,7 @@ using namespace ::onnxruntime::common; namespace onnxruntime { Status UnsqueezeElimination::Apply(Graph& graph, Node& node, bool& modified, bool& removed) { - // Get "axes" attribute + // Get "axes" attribute. const ONNX_NAMESPACE::AttributeProto* attr = graph_utils::GetNodeAttribute(node, "axes"); if (attr == nullptr || attr->type() != AttributeProto_AttributeType_INTS) { return Status::OK(); @@ -22,7 +22,7 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, bool& modified, boo axes.push_back(static_cast(attr->ints(i))); } - // Generate new dims + // Generate new dims. NodeArg* input_def = node.MutableInputDefs()[0]; const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr; graph.GetInitializedTensor(input_def->Name(), tensor_proto); @@ -45,7 +45,7 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, bool& modified, boo } } - // Update shape of tensor proto + // Update shape of tensor proto. ONNX_NAMESPACE::TensorProto new_tensor_proto(*tensor_proto); for (int i = 0; i < static_cast(new_dims.size()); i++) { @@ -58,36 +58,18 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, bool& modified, boo graph.RemoveInitializedTensor(input_def->Name()); graph.AddInitializedTensor(new_tensor_proto); - // Update shape of NodeArg + // Update shape of NodeArg. TensorShapeProto shape; for (auto dim : new_dims) { shape.add_dim()->set_dim_value(dim); } input_def->SetShape(shape); - // Replace the input of the nodes following the Unsqueeze node. - const NodeArg* output_def = node.OutputDefs()[0]; - for (auto it = node.OutputNodesBegin(), end = node.OutputNodesEnd(); it != end; ++it) { - auto output_node_idx = (*it).Index(); - auto output_node = graph.GetNode(output_node_idx); - if (!output_node) { - return Status(ONNXRUNTIME, INVALID_ARGUMENT); - } - auto& input_defs = output_node->MutableInputDefs(); - for (auto& def : input_defs) { - if (def == output_def) { - def = input_def; - } - } + // Remove Unsqueeze node. + if (graph_utils::RemoveSingleInputNode(graph, node)) { + removed = modified = true; } - // Remove output edges of the Unsqueeze node. - graph_utils::RemoveNodeOutputEdges(graph, node); - - // Remove the Unsqueeze node. - graph.RemoveNode(node.Index()); - removed = modified = true; - return Status::OK(); } // namespace onnxruntime diff --git a/onnxruntime/test/ir/utils_test.cc b/onnxruntime/test/ir/utils_test.cc index 6f8085fe3b..0cc0d98a0e 100644 --- a/onnxruntime/test/ir/utils_test.cc +++ b/onnxruntime/test/ir/utils_test.cc @@ -207,7 +207,7 @@ static void UpdateSubgraphWhenRemovingNode(bool include_nested = false) { auto& node_to_remove = *graph.GetNode(1); const auto& if_node = *graph.GetNode(2); - bool removed = graph_utils::RemoveSingleInSingleOutNode(graph, node_to_remove); + bool removed = graph_utils::RemoveSingleInputNode(graph, node_to_remove); ASSERT_TRUE(removed); // check subgraph implicit input was updated @@ -238,7 +238,7 @@ static void DontRemoveNodeIfItWillBreakSubgraph(bool test_nested = false) { auto& graph = model.MainGraph(); auto& node_to_remove = *graph.GetNode(1); - bool removed = graph_utils::RemoveSingleInSingleOutNode(graph, node_to_remove); + bool removed = graph_utils::RemoveSingleInputNode(graph, node_to_remove); ASSERT_FALSE(removed); } @@ -250,5 +250,58 @@ TEST(GraphUtils, DontRemoveNodeIfItWillBreakNestedSubgraph) { DontRemoveNodeIfItWillBreakSubgraph(true); } +TEST(GraphUtils, TestMultiEdgeRemovalNodes) { + // Create a graph with 5 Id nodes. The graph structure is as follows: Id0 ( Id1 Id2 ( Id3 Id4 ) ). + // First we remove Id2, which leads to: Id0 ( Id1 Id4 Id5 ). + // Then we remove Id1, which leads to: Id2 Id4 Id5, being fed the initializer. + Model model("MultiEdgeRemovalGraph"); + auto& graph = model.MainGraph(); + + TypeProto float_tensor; + float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto& id_0_in = graph.GetOrCreateNodeArg("id_0_in", &float_tensor); + auto& id_0_out = graph.GetOrCreateNodeArg("id_0_out", &float_tensor); + auto& id_1_out = graph.GetOrCreateNodeArg("id_1_out", &float_tensor); + auto& id_2_out = graph.GetOrCreateNodeArg("id_2_out", &float_tensor); + auto& id_3_out = graph.GetOrCreateNodeArg("id_3_out", &float_tensor); + auto& id_4_out = graph.GetOrCreateNodeArg("id_4_out", &float_tensor); + + std::vector nodes; + nodes.push_back(&graph.AddNode("id_0", "Identity", "Identity node 0", {&id_0_in}, {&id_0_out})); + nodes.push_back(&graph.AddNode("id_1", "Identity", "Identity node 1", {&id_0_out}, {&id_1_out})); + nodes.push_back(&graph.AddNode("id_2", "Identity", "Identity node 2", {&id_0_out}, {&id_2_out})); + nodes.push_back(&graph.AddNode("id_3", "Identity", "Identity node 3", {&id_2_out}, {&id_3_out})); + nodes.push_back(&graph.AddNode("id_4", "Identity", "Identity node 4", {&id_2_out}, {&id_4_out})); + + auto status = graph.Resolve(); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + ASSERT_EQ(graph.NumberOfNodes(), 5); + + // Check inputs/outputs of id_0 and id_2 + ASSERT_EQ(nodes[0]->GetInputEdgesCount(), 0); + ASSERT_EQ(nodes[0]->GetOutputEdgesCount(), 2); + ASSERT_EQ(nodes[2]->GetInputEdgesCount(), 1); + ASSERT_EQ(nodes[2]->GetOutputEdgesCount(), 2); + + // Remove id_2. This leaves id_0 with 3 output edges. id_0 is now incoming node to id_3 and id_4. + ASSERT_TRUE(graph_utils::RemoveSingleInputNode(graph, *nodes[2])); + ASSERT_EQ(graph.NumberOfNodes(), 4); + ASSERT_EQ(nodes[0]->GetOutputEdgesCount(), 3); + ASSERT_EQ(nodes[3]->InputDefs().size(), 1); + ASSERT_TRUE(nodes[3]->InputDefs()[0]->Name() == "id_0_out"); + ASSERT_EQ(nodes[4]->InputDefs().size(), 1); + ASSERT_TRUE(nodes[4]->InputDefs()[0]->Name() == "id_0_out"); + + // Remove id_0 + ASSERT_TRUE(graph_utils::RemoveSingleInputNode(graph, *nodes[0])); + ASSERT_EQ(graph.NumberOfNodes(), 3); + ASSERT_TRUE(nodes[1]->InputDefs()[0]->Name() == "id_0_in"); + ASSERT_TRUE(nodes[3]->InputDefs()[0]->Name() == "id_0_in"); + ASSERT_TRUE(nodes[4]->InputDefs()[0]->Name() == "id_0_in"); +} + } // namespace test } // namespace onnxruntime