[IR] introduce remove_nodes (#2294)

Convenience function to remove a set of nodes.

Author: Johansmm

onnxscript/ir/_convenience/__init__.py

@@ -15,6 +15,7 @@
     "create_value_mapping",
     "replace_nodes_and_values",
     "insert_nodes_in_value",
+    "remove_nodes",
 ]
 
 from typing import Mapping, Sequence, Union
@@ -487,3 +488,65 @@ def insert_nodes_in_value(
         # Insert new nodes if there is a graph
         graph.extend(new_nodes)
         graph.sort()
+
+
+def remove_nodes(nodes: Sequence[_core.Node]) -> None:
+    """Remove a sequence of nodes.
+
+    This allows to delete a list of LINKED nodes (over the same context).
+
+    For example, suppose we have the following graph::
+
+        input -> A := node_A(input) -> B := node_B(A) -> C := node_C(B) -> output
+
+    We want to prune [node_B]::
+
+        >>> from onnxscript import ir
+        >>> input = ir.Input("input")
+        >>> node_A = ir.node("op_A", [input])
+        >>> node_B = ir.node("op_B", node_A.outputs)
+        >>> node_C = ir.node("op_C", node_B.outputs)
+        >>> # Delete node_B
+        >>> remove_nodes([node_B])
+        >>> len(node_A.outputs[0].consumers())
+        1
+        >>> node_A.outputs[0].consumers()[0].op_type
+        'op_C'
+        >>> len(node_C.inputs)
+        1
+        >>> node_C.inputs[0].producer().op_type
+        'op_A'
+        >>> node_B.inputs
+        (None,)
+        >>> len(node_B.outputs)
+        1
+        >>> len(node_B.outputs[0].consumers())
+        0
+
+    Args:
+        nodes: The nodes to remove.
+    """
+    # Search the unique inputs/outputs in new_nodes, keeping the order.
+    inputs, outputs = _find_inputs_outputs(nodes)
+
+    # Sanity check.
+    if len(inputs) != len(outputs):
+        raise ValueError(
+            f"The number of inputs ({inputs}) and outputs ({outputs}) in nodes must match."
+        )
+
+    # Remove nodes, in several steps:
+    # 1. Reconnect the users of outputs with inputs
+    replace_all_uses_with(outputs, inputs)
+    # 2. Detach nodes for their inputs
+    for node in nodes:
+        for i in range(len(node.inputs)):
+            node.replace_input_with(i, None)
+
+    # Update graph if there is one:
+    if (graph := inputs[-1].graph) is not None:
+        # Update graph/function outputs if the node generates output
+        _update_graph_or_function_outputs(graph, outputs, inputs)
+
+        # Drop nodes from graph
+        graph.remove(nodes, safe=True)

onnxscript/ir/_convenience/_init_test.py

@@ -7,7 +7,7 @@
 import unittest
 
 from onnxscript import ir
-from onnxscript.ir._convenience import insert_nodes_in_value
+from onnxscript.ir._convenience import insert_nodes_in_value, remove_nodes
 
 
 def _create_model(model_text: str) -> ir.Model:
@@ -141,6 +141,108 @@ def test_value_error_for_wrong_number_of_points(self):
         with self.assertRaisesRegex(ValueError, "The number of values and outputs"):
             insert_nodes_in_value(ir_model.graph[1].outputs, [node])
 
+    def test_remove_nodes(self):
+        # Main graph
+        input = ir.Input("input")
+        node_A = ir.node("op_A", [input])
+        node_B = ir.node("op_B", node_A.outputs)
+        node_C = ir.node("op_C", node_B.outputs)
+
+        # Delete node_B
+        remove_nodes([node_B])
+        self.assertEqual(len(node_A.outputs[0].consumers()), 1)
+        self.assertEqual(node_A.outputs[0].consumers()[0].op_type, "op_C")
+        self.assertEqual(len(node_C.inputs), 1)
+        self.assertEqual(node_C.inputs[0].producer().op_type, "op_A")
+
+        self.assertEqual((len(node_B.inputs), len(node_B.outputs)), (1, 1))
+        self.assertEqual(node_B.inputs, (None,))
+        self.assertEqual(len(node_B.outputs[0].consumers()), 0)
+
+    def test_remove_nodes_in_graph(self):
+        ir_model = _create_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x) => (float[N] z) {
+                two = Constant<value_float=2.0>()
+                a, b = MergeAndSplit(x, two)
+                z = MergeNode(a, b, two)
+            }
+        """
+        )
+        # Sanity check previous to delete nodes
+        x, two = ir_model.graph.inputs[0], ir_model.graph[0].outputs[0]
+        self.assertEqual(len(x.consumers()), 1)
+        self.assertEqual(len(two.consumers()), 2)
+
+        # Delete 'MergeAndSplit'
+        target_node = ir_model.graph[1]
+        remove_nodes([target_node])
+
+        # Check 'MergeNode' has new inputs
+        a, b, _ = ir_model.graph[-1].inputs
+        self.assertEqual(a.name, "x")
+        self.assertEqual(b.name, "two")
+
+        # Check x/two consumers have been updated
+        self.assertEqual(len(x.consumers()), 1)
+        self.assertEqual(len(two.consumers()), 1)
+
+        # Check nodes have been deleted in the graph
+        self.assertEqual(len(ir_model.graph), 2)
+
+    def test_remove_nodes_in_input(self):
+        ir_model = _create_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x) => (float[N] z) {
+                y = Sigmoid(x)
+                z = Mul(y, y)
+            }
+        """
+        )
+        # Remove the node linked to the input
+        remove_nodes([ir_model.graph[0]])
+        self.assertEqual(len(ir_model.graph), 1)
+        self.assertEqual(ir_model.graph[0].op_type, "Mul")
+        self.assertEqual(ir_model.graph[0].inputs[0].name, "x")
+        self.assertEqual(ir_model.graph[0].inputs[1].name, "x")
+        self.assertEqual(ir_model.graph[0].outputs[0].name, "z")
+
+    def test_remove_nodes_in_output(self):
+        ir_model = _create_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x) => (float[N] z) {
+                y = Mul(x, x)
+                z = Sigmoid(y)
+            }
+        """
+        )
+        # Remove the node linked to the input
+        remove_nodes([ir_model.graph[-1]])
+        self.assertEqual(len(ir_model.graph), 1)
+        self.assertEqual(ir_model.graph[0].op_type, "Mul")
+        self.assertEqual(ir_model.graph[0].outputs[0].name, "y")
+        self.assertEqual(ir_model.graph.outputs[0].name, "y")
+
+    def test_remove_nodes_error_for_wrong_number_of_inputs_and_outputs(self):
+        ir_model = _create_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x) => (float[N] z) {
+                two = Constant<value_float=2.0>()
+                a, b = SplitNode(x)
+                z = MergeNode(a, b, two)
+            }
+        """
+        )
+        with self.assertRaisesRegex(ValueError, "The number of inputs"):
+            remove_nodes([ir_model.graph[0]])
+
+        with self.assertRaisesRegex(ValueError, "The number of inputs"):
+            remove_nodes([ir_model.graph[1]])
+
 
 if __name__ == "__main__":
     unittest.main()

onnxscript/ir/convenience.py

@@ -11,6 +11,7 @@
     "replace_nodes_and_values",
     "create_value_mapping",
     "insert_nodes_in_value",
+    "remove_nodes",
 ]
 
 from onnxscript.ir._convenience import (
@@ -20,6 +21,7 @@
     replace_all_uses_with,
     replace_nodes_and_values,
     insert_nodes_in_value,
+    remove_nodes,
 )
 
 # NOTE: Do not implement any other functions in this module.