Implement replace_all_uses_with on Value (#235)

Given it is a fundamental operation in graph transformation, I added it
as a method on `Value` so that it is more convenient and discoverable.
The ir.convenience.replace_all_uses_with function is preserved to handle
batch processing cases on multiple values.

Also added a `replace_graph_outputs` option to allow users to replace
graph outputs as well. Users are still responsible to assign the output
name to the replacement value if they want to preserve the signature.

### BC breaking

An error is raised when `replace_graph_outputs` is False && when the
value to replace is a graph output.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/_convenience/__init__.py

@@ -280,6 +280,7 @@ def convert_attributes(
 def replace_all_uses_with(
     values: _protocols.ValueProtocol | Sequence[_protocols.ValueProtocol],
     replacements: _protocols.ValueProtocol | Sequence[_protocols.ValueProtocol],
+    replace_graph_outputs: bool = False,
 ) -> None:
     """Replace all uses of the given values with the replacements.
 
@@ -318,9 +319,22 @@ def replace_all_uses_with(
         replaced are part of the graph outputs. Be sure to remove the old nodes
         from the graph using ``graph.remove()`` if they are no longer needed.
 
+    .. versionadded:: 0.1.12
+        The ``replace_graph_outputs`` parameter is added.
+
+    .. versionadded:: 0.1.12
+        ValueError is raised when ``replace_graph_outputs`` is False && when the value to
+        replace is a graph output.
+
     Args:
         values: The value or values to be replaced.
         replacements: The new value or values to use as inputs.
+        replace_graph_outputs: If True, graph outputs that reference the values
+            being replaced will also be updated to reference the replacements.
+
+    Raises:
+        ValueError: When ``replace_graph_outputs`` is False && when the value to
+            replace is a graph output.
     """
     if not isinstance(values, Sequence):
         values = (values,)
@@ -329,8 +343,7 @@ def replace_all_uses_with(
     if len(values) != len(replacements):
         raise ValueError("The number of values and replacements must match.")
     for value, replacement in zip(values, replacements):
-        for user_node, index in tuple(value.uses()):
-            user_node.replace_input_with(index, replacement)
+        value.replace_all_uses_with(replacement, replace_graph_outputs=replace_graph_outputs)
 
 
 def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
@@ -408,12 +421,7 @@ def replace_nodes_and_values(
         new_value.name = old_value.name if old_value.name is not None else new_value.name
 
     # Reconnect the users of the deleted values to use the new values
-    replace_all_uses_with(old_values, new_values)
-    # Update graph/function outputs if the node generates output
-    replacement_mapping = dict(zip(old_values, new_values))
-    for idx, graph_or_function_output in enumerate(graph_or_function.outputs):
-        if graph_or_function_output in replacement_mapping:
-            graph_or_function.outputs[idx] = replacement_mapping[graph_or_function_output]
+    replace_all_uses_with(old_values, new_values, replace_graph_outputs=True)
 
     # insert new nodes after the index node
     graph_or_function.insert_after(insertion_point, new_nodes)

src/onnx_ir/_core.py

@@ -2448,6 +2448,55 @@ def is_initializer(self) -> bool:
         """Whether the value is an initializer of a graph."""
         return self._is_initializer
 
+    def replace_all_uses_with(
+        self, replacement: Value, /, replace_graph_outputs: bool = False
+    ) -> None:
+        """Replace all uses of this value with another value.
+
+        If the value is an output of a graph and ``replace_graph_outputs`` is ``True``,
+        the graph output will also be replaced. Be careful when a value appears multiple times
+        in the graph outputs - this is invalid. An identity node will need to be added on each
+        duplicated outputs to ensure a valid ONNX graph.
+
+        You may also want to assign the name of this value to the replacement value
+        to maintain the name when it is a graph output.
+
+        To replace usage of a sequence of values with another sequence of values, consider using
+        :func:`onnx_ir.convenience.replace_all_uses_with`.
+
+        .. versionadded:: 0.1.12
+
+        Args:
+            replacement: The value to replace all uses with.
+            replace_graph_outputs: If True, graph outputs that reference this value
+                will also be updated to reference the replacement.
+
+        Raises:
+            ValueError: When ``replace_graph_outputs`` is False && when the value to
+                replace is a graph output.
+        """
+        # NOTE: Why we don't replace the value name when the value is an output:
+        # When the replacement value is already an output of the graph, renaming it
+        # to the name of this value will cause name conflicts. It is better to let
+        # the user handle the renaming explicitly and insert identity nodes if needed.
+        if self.is_graph_output():
+            graph = self.graph
+            assert graph is not None
+
+            if not replace_graph_outputs:
+                raise ValueError(
+                    f"{self!r} is an output of graph {graph.name!r}. "
+                    "Set replace_graph_outputs=True or replace the graph output frist before "
+                    "calling replace_all_uses_with."
+                )
+
+            for i, output in enumerate(graph.outputs):
+                if output is self:
+                    graph.outputs[i] = replacement
+
+        for user_node, index in self.uses():
+            user_node.replace_input_with(index, replacement)
+
 
 @deprecated("Input is deprecated since 0.1.9. Use ir.val(...) instead.")
 def Input(

src/onnx_ir/_protocols.py

@@ -203,6 +203,17 @@ def is_graph_output(self) -> bool:
         """Whether this value is an output of a graph."""
         ...
 
+    def replace_all_uses_with(
+        self, new_value: ValueProtocol | None, replace_graph_outputs: bool = False
+    ) -> None:
+        """Replace all uses of this value with the given new value.
+
+        Args:
+            new_value: The new value to replace this value with.
+            replace_graph_outputs: Whether to replace graph outputs that use this value.
+        """
+        ...
+
 
 @typing.runtime_checkable
 class NodeProtocol(Protocol):

src/onnx_ir/passes/common/common_subexpression_elimination.py

@@ -150,8 +150,6 @@ def _remove_node_and_replace_values(
         remove_values: The values to replace.
         new_values: The values to replace with.
     """
-    # Reconnect the users of the deleted values to use the new values
-    ir.convenience.replace_all_uses_with(remove_values, new_values)
     # Update graph/function outputs if the node generates output
     if any(remove_value.is_graph_output() for remove_value in remove_values):
         replacement_mapping = dict(zip(remove_values, new_values))
@@ -185,6 +183,9 @@ def _remove_node_and_replace_values(
                     new_value.name = graph_output.name
                     graph.outputs[idx] = new_value
 
+    # Reconnect the users of the deleted values to use the new values
+    ir.convenience.replace_all_uses_with(remove_values, new_values)
+
     graph.remove(remove_node, safe=True)
 
 

src/onnx_ir/passes/common/constant_manipulation.py

@@ -78,7 +78,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             assert node.graph is not None
             node.graph.register_initializer(initializer)
             # Replace the constant node with the initializer
-            ir.convenience.replace_all_uses_with(node.outputs[0], initializer)
+            node.outputs[0].replace_all_uses_with(initializer)
             node.graph.remove(node, safe=True)
             count += 1
             logger.debug(

src/onnx_ir/passes/common/identity_elimination.py

@@ -105,20 +105,14 @@ def _try_eliminate_identity_node(self, node: ir.Node) -> bool:
 
         # Case 1 & 2 (merged): Eliminate the identity node
         # Replace all uses of output with input
-        ir.convenience.replace_all_uses_with(output_value, input_value)
+        ir.convenience.replace_all_uses_with(
+            output_value, input_value, replace_graph_outputs=True
+        )
 
         # If output is a graph output, we need to rename input and update graph outputs
         if output_is_graph_output:
-            # Store the original output name
-            original_output_name = output_value.name
-
             # Update the input value to have the output's name
-            input_value.name = original_output_name
-
-            # Update graph outputs to point to the input value
-            for idx, graph_output in enumerate(graph_like.outputs):
-                if graph_output is output_value:
-                    graph_like.outputs[idx] = input_value
+            input_value.name = output_value.name
 
         # Remove the identity node
         graph_like.remove(node, safe=True)

src/onnx_ir/passes/common/initializer_deduplication.py

@@ -100,7 +100,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 if key in initializers:
                     modified = True
                     initializer_to_keep = initializers[key]  # type: ignore[index]
-                    ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
+                    initializer.replace_all_uses_with(initializer_to_keep)
                     assert initializer.name is not None
                     graph.initializers.pop(initializer.name)
                     logger.info(
@@ -165,7 +165,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                         continue
                     modified = True
                     initializer_to_keep = initializers[key]  # type: ignore[index]
-                    ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
+                    initializer.replace_all_uses_with(initializer_to_keep)
                     assert initializer.name is not None
                     graph.initializers.pop(initializer.name)
                     logger.info(