[pass] Remove unused initialized inputs in DCE (#2212)

Fix #2211

This pull request enhances the functionality of the
`RemoveUnusedNodesPass` class and its associated methods by introducing
an option to remove unused initialized inputs. It also updates the
corresponding tests to validate this new behavior. The changes improve
the flexibility of the unused node removal process and ensure the model
input signature remains consistent unless explicitly modified.

### Enhancements to `RemoveUnusedNodesPass`:

* Added a new `remove_initialized_inputs` attribute to the
`RemoveUnusedNodesPass` class, allowing the removal of unused
initialized inputs when enabled. This change modifies the model input
signature if unused inputs are removed
(`onnxscript/ir/passes/common/unused_removal.py`).

---------

Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: leshabirukov

onnxscript/ir/passes/common/unused_removal.py

@@ -93,10 +93,27 @@ def _remove_unused_nodes_in_graph_like(function_or_graph: ir.Function | ir.Graph
 
 
 class RemoveUnusedNodesPass(ir.passes.InPlacePass):
+    """Pass for removing unused nodes and initializers.
+
+    Attributes:
+        remove_initialized_inputs: When an unused initializer is simultaneously a graph input,
+            remove that input as well. Note that this will change the model input signature.
+    """
+
+    def __init__(self, remove_initialized_inputs: bool = False):
+        super().__init__()
+        self.remove_initialized_inputs = remove_initialized_inputs
+
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         count = _remove_unused_nodes_in_graph_like(model.graph)
         graph_outputs = frozenset(model.graph.outputs)
         initializers = model.graph.initializers
+        if self.remove_initialized_inputs:
+            graph_inputs = model.graph.inputs
+            for i, inp in reversed(list(enumerate(graph_inputs))):
+                if inp.name in initializers and not (inp in graph_outputs or inp.uses()):
+                    del graph_inputs[i]
+                    count += 1
         for init in list(initializers.values()):
             if not (init in graph_outputs or init.uses()):
                 assert init.name is not None

onnxscript/ir/passes/common/unused_removal_test.py

@@ -13,13 +13,15 @@
 class RemoveUnusedTest(unittest.TestCase):
     using_ir: bool
 
-    def remove_unused_nodes(self, model: onnx.ModelProto):
+    def remove_unused_nodes(
+        self, model: onnx.ModelProto, remove_initialized_inputs: bool = False
+    ):
         if self.using_ir:
             model_ir = ir.serde.deserialize_model(model)
-            onnxscript.optimizer.remove_unused_nodes(model_ir)
+            onnxscript.optimizer.remove_unused_nodes(model_ir, remove_initialized_inputs)
             model = ir.serde.serialize_model(model_ir)
             return model
-        onnxscript.optimizer.remove_unused_nodes(model)
+        onnxscript.optimizer.remove_unused_nodes(model, remove_initialized_inputs)
         return model
 
     def test_remove_unused_nodes(self):
@@ -54,6 +56,59 @@ def test_remove_unused_initializers(self):
         self.assertEqual(model.graph.node[0].op_type, "Mul")
         self.assertEqual(len(model.graph.initializer), 0)
 
+    def test_unused_initialized_inputs_are_removed_when_requested(self):
+        # https://github.com/microsoft/onnxscript/issues/2211
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x, float[N] two) => (float[N] z)
+            <float two = {2.0,2.0}> {
+                four = Add(two, two)
+                z = Mul(x, x)
+            }
+        """
+        )
+        model = self.remove_unused_nodes(model, remove_initialized_inputs=True)
+        self.assertEqual(len(model.graph.node), 1)
+        self.assertEqual(model.graph.node[0].op_type, "Mul")
+        self.assertEqual(len(model.graph.input), 1)
+
+    def test_unused_initialized_inputs_are_kept_by_default(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x, float[N] two) => (float[N] z)
+            <float two = {2.0,2.0}> {
+                four = Add(two, two)
+                z = Mul(x, x)
+            }
+        """
+        )
+        model = self.remove_unused_nodes(model)
+        self.assertEqual(len(model.graph.node), 1)
+        self.assertEqual(model.graph.node[0].op_type, "Mul")
+        self.assertEqual(len(model.graph.input), 2)
+
+    @parameterized.parameterized.expand([True, False])
+    def test_unused_inputs_are_not_removed(self, remove_initialized_inputs: bool):
+        # preserve inputs as part of interface
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 10, opset_import: [ "" : 17]>
+            agraph (float[N] x, float[N] two) => (float[N] z)
+            {
+                four = Add(two, two)
+                z = Mul(x, x)
+            }
+        """
+        )
+        model = self.remove_unused_nodes(
+            model, remove_initialized_inputs=remove_initialized_inputs
+        )
+        self.assertEqual(len(model.graph.node), 1)
+        self.assertEqual(model.graph.node[0].op_type, "Mul")
+        self.assertEqual(len(model.graph.input), 2)
+
     def test_partially_used_nodes(self):
         model = onnx.parser.parse_model(
             """

onnxscript/optimizer/__init__.py

@@ -112,15 +112,19 @@ def fold_constants(
         return result
 
 
-def remove_unused_nodes(model: ir.Model | onnx.ModelProto) -> None:
+def remove_unused_nodes(
+    model: ir.Model | onnx.ModelProto, remove_initialized_inputs: bool = False
+) -> None:
     """Removes unused nodes from a model inplace."""
     if isinstance(model, ir.Model):
-        onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass()(model)
+        onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass(
+            remove_initialized_inputs=remove_initialized_inputs
+        )(model)
     else:
         model_ir = ir.serde.deserialize_model(model)
-        model_ir = onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass()(
-            model_ir
-        ).model
+        model_ir = onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass(
+            remove_initialized_inputs=remove_initialized_inputs
+        )(model_ir).model
         new_proto = ir.serde.serialize_model(model_ir)
         model.Clear()
         model.CopyFrom(new_proto)