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

@@ -156,6 +156,11 @@ class Node {
     return definitions_.implicit_input_defs;
   }
 
+  /** Gets a modifiable collection of the Node's implicit input definitions. */
+  std::vector<NodeArg*>& 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<std::vector<NodeArg*>> 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<NodeArg*> CreateNodeArgs(const google::protobuf::RepeatedPtrField<std::string>& 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.

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<std::string>& 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<std::string> node_args_consumed;
@@ -932,6 +940,17 @@ Status Graph::BuildConnections(std::vector<std::string>& 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());
+              }
+            }
           }
         }
       }

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<GraphEdge> GetNodeOutputEdges(const Node& node) {
+  std::vector<GraphEdge> 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<GraphEdge>& 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<GraphEdge>& 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<GraphEdge> 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<GraphEdge> 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<std::string>& 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<Status(const Graph&)> 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<std::tuple<NodeIndex, int, int>> 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<GraphEdge> output_edges = GetNodeOutputEdges(node);
+  RemoveGraphEdges(graph, output_edges);
 
-  return edges_to_remove.size();
+  return output_edges.size();
 }
 
 }  // namespace graph_utils

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<std::string>& 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<Status(Graph&)> func);
 Status ForAllSubgraphs(const Graph& main_graph, std::function<Status(const Graph&)> 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. */

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

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;
   }
 

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<int64_t>(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<int>(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
 

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<Node*> 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