[Passes] Consolidate DCE passes into common passes (#2143)

Consolidate DCE passes into common passes (unused_removal) for them to
be available for pass users. Refactored usage. Added a pass to remove
unused opset imports.

Author: justinchuby

onnxscript/ir/passes/_pass_infra.py

@@ -249,7 +249,8 @@ def call(self, model: ir.Model) -> PassResult:
         overall_modified = False
         for step in range(self.steps):
             try:
-                step_result = super().__call__(model)
+                # Call the call method of Sequential
+                step_result = super().call(model)
             except Exception as e:
                 raise PassError(f"An error occurred at step {step}") from e
             model = step_result.model

onnxscript/ir/passes/common/unused_removal.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+__all__ = [
+    "RemoveUnusedNodesPass",
+    "RemoveUnusedFunctionsPass",
+    "RemoveUnusedOpsetsPass",
+]
+
+import logging
+
+import onnx
+
+from onnxscript import ir
+
+logger = logging.getLogger(__name__)
+
+
+def _remove_unused_optional_outputs(
+    node: ir.Node, graph_outputs: frozenset[ir.Value], onnx_opset_version: int
+) -> None:
+    try:
+        if node.domain not in {"", "onnx.ai"}:
+            return
+        op_schema = onnx.defs.get_schema(node.op_type, onnx_opset_version, domain=node.domain)
+    except Exception:  # pylint: disable=broad-exception-caught
+        logger.info(
+            "Failed to get schema for %s, skipping optional output removal",
+            node,
+            stack_info=True,
+        )
+        return
+
+    if node.op_type == "BatchNormalization":
+        # BatchNormalization op has 3 outputs: Y, running_mean, running_var
+        # If running_mean and running_var are not used, remove them, and the training_mode attribute
+        def is_used_output(i: int) -> bool:
+            if i < len(node.outputs):
+                val = node.outputs[i]
+                return val in graph_outputs or bool(val.uses())
+            return False
+
+        if is_used_output(1) or is_used_output(2):
+            return
+        if len(node.outputs) > 1:
+            node.outputs[1].name = ""
+        if len(node.outputs) > 2:
+            node.outputs[2].name = ""
+        node.attributes.pop("training_mode", None)
+        return
+
+    optional_info = []
+    for o in op_schema.outputs:
+        # Current ops do not have optional outputs if they have variable number of outputs
+        if o.option == onnx.defs.OpSchema.FormalParameterOption.Variadic:
+            return
+        optional_info.append(o.option == onnx.defs.OpSchema.FormalParameterOption.Optional)
+    # If no optional outputs in spec, skip delete operations
+    if len([o == 1 for o in optional_info]) == 0:
+        return
+
+    for i, out in enumerate(node.outputs):
+        if out not in graph_outputs and (not out.uses()) and optional_info[i] is True:
+            out.name = ""
+
+
+def _remove_unused_nodes_in_graph_like(function_or_graph: ir.Function | ir.Graph) -> int:
+    graph_outputs = frozenset(function_or_graph.outputs)
+    onnx_opset_version = function_or_graph.opset_imports.get("", None)
+    count = 0
+    for node in reversed(function_or_graph):
+        removable = True
+        for output in node.outputs:
+            if output in graph_outputs or output.uses():
+                removable = False
+                break
+        if removable:
+            function_or_graph.remove(node, safe=True)
+            count += 1
+        else:
+            if onnx_opset_version is not None:
+                _remove_unused_optional_outputs(node, graph_outputs, onnx_opset_version)
+            for attr in node.attributes.values():
+                if not isinstance(attr, ir.Attr):
+                    continue
+                if attr.type == ir.AttributeType.GRAPH:
+                    count += _remove_unused_nodes_in_graph_like(attr.as_graph())
+                elif attr.type == ir.AttributeType.GRAPHS:
+                    for graph in attr.as_graphs():
+                        count += _remove_unused_nodes_in_graph_like(graph)
+    return count
+
+
+class RemoveUnusedNodesPass(ir.passes.InPlacePass):
+    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
+        for init in list(initializers.values()):
+            if not (init in graph_outputs or init.uses()):
+                assert init.name is not None
+                del initializers[init.name]
+                count += 1
+        for function in model.functions.values():
+            count += _remove_unused_nodes_in_graph_like(function)
+        if count:
+            logger.info("Removed %s unused nodes", count)
+        return ir.passes.PassResult(model, modified=bool(count))
+
+
+class RemoveUnusedFunctionsPass(ir.passes.InPlacePass):
+    def __init__(self):
+        super().__init__()
+        self._used: set[ir.OperatorIdentifier] | None = None
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        self._used = set()
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            self._call_node(model, node)
+
+        # Update the model to remove unused functions
+        unused = set(model.functions) - self._used
+        if not unused:
+            logger.info("No unused functions to remove")
+            return ir.passes.PassResult(model, modified=False)
+
+        for op_identifier in unused:
+            del model.functions[op_identifier]
+
+        logger.info("Removed %s unused functions", len(unused))
+        logger.debug("Functions left: %s", list(model.functions))
+        logger.debug("Functions removed: %s", unused)
+
+        self._used = None
+        return ir.passes.PassResult(model, modified=bool(unused))
+
+    def _call_function(self, model: ir.Model, function: ir.Function) -> None:
+        assert self._used is not None
+        if function.identifier() in self._used:
+            # The function and its nodes are already recorded as used
+            return
+        self._used.add(function.identifier())
+        for node in ir.traversal.RecursiveGraphIterator(function):
+            self._call_node(model, node)
+
+    def _call_node(self, model: ir.Model, node: ir.Node) -> None:
+        op_identifier = node.op_identifier()
+        if op_identifier not in model.functions:
+            return
+        self._call_function(model, model.functions[op_identifier])
+
+
+class RemoveUnusedOpsetsPass(ir.passes.InPlacePass):
+    """Remove unused opset imports from the model and functions.
+
+    Attributes:
+        process_functions: Whether to process functions in the model. If True, the pass will
+            remove unused opset imports from functions as well. If False, only the main graph
+            will be processed.
+    """
+
+    def __init__(self, process_functions: bool = True):
+        super().__init__()
+        self.process_functions = process_functions
+
+    def _process_graph_like(
+        self, graph_like: ir.Graph | ir.Function, used_domains: set[str]
+    ) -> bool:
+        for node in ir.traversal.RecursiveGraphIterator(graph_like):
+            used_domains.add(node.domain)
+        unused = set(graph_like.opset_imports) - used_domains
+        for domain in unused:
+            del graph_like.opset_imports[domain]
+        return bool(unused)
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        # Record domains of all functions
+        used_domains = set()
+        for function in model.functions.values():
+            used_domains.add(function.domain)
+        modified = self._process_graph_like(model.graph, used_domains=used_domains)
+
+        if self.process_functions:
+            for function in model.functions.values():
+                modified |= self._process_graph_like(function, used_domains=set())
+
+        return ir.passes.PassResult(model, modified=modified)

onnxscript/optimizer/_remove_unused_test.py onnxscript/ir/passes/common/unused_removal_test.py

@@ -25,7 +25,7 @@ def remove_unused_nodes(self, model: onnx.ModelProto):
     def test_remove_unused_nodes(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[N] x) => (float[N] z) {
                 two = Constant <value_float=2.0> ()
                 four = Add(two, two)
@@ -40,7 +40,7 @@ def test_remove_unused_nodes(self):
     def test_remove_unused_initializers(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[N] x) => (float[N] z)
             <float two = {2.0}> {
                 four = Add(two, two)
@@ -57,7 +57,7 @@ def test_remove_unused_initializers(self):
     def test_partially_used_nodes(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[N] x) => (float[M] z) {
                 w1, w2, w3 = Split (x)
                 z = Mul(w3, w3)
@@ -71,7 +71,7 @@ def test_partially_used_nodes(self):
     def test_remove_unused_optional_outputs_maxpool(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 1, 5, 5] x) => (float[1, 1, 5, 5] z) {
                 z, indices = MaxPool <pads = [2, 2, 2, 2], kernel_shape = [5, 5]> (x)
             }
@@ -88,7 +88,7 @@ def test_remove_unused_optional_outputs_maxpool(self):
     def test_remove_unused_optional_outputs_dropout_in_function(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17, "pkg.custom": 1]>
+            <ir_version: 10, opset_import: [ "" : 17, "pkg.custom": 1]>
             agraph (float[1, 1, 5, 5] x) => (float[1, 1, 5, 5] z)
             {
                 z = pkg.custom.afunction (x)
@@ -113,7 +113,7 @@ def test_remove_unused_optional_outputs_dropout_in_function(self):
     def test_remove_used_optional_outputs_maxpool(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 1, 5, 5] x) => (float[1, 1, 5, 5] y, float[1, 1, 5, 5] z) {
                 y, z = MaxPool <pads = [2, 2, 2, 2], kernel_shape = [5, 5]> (x)
             }
@@ -130,7 +130,7 @@ def test_remove_used_optional_outputs_maxpool(self):
     def test_remove_multiple_unused_optional_outputs_layernorm(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 3, 5, 5] x) => (float[1, 3, 5, 5] z) {
                 scale = Constant <value_ints=[3]> ()
                 B = Constant <value_ints=[3]> ()
@@ -149,7 +149,7 @@ def test_remove_multiple_unused_optional_outputs_layernorm(self):
     def test_remove_trailing_unused_optional_outputs_layernorm(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 3, 5, 5] x) => (float[1, 3, 5, 5] z, float[1, 3, 5, 5] mean) {
                 scale = Constant <value_ints=[3]> ()
                 B = Constant <value_ints=[3]> ()
@@ -168,7 +168,7 @@ def test_remove_trailing_unused_optional_outputs_layernorm(self):
     def test_avoid_remove_non_trailing_unused_optional_outputs_layernorm(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 3, 5, 5] x) => (float[1, 3, 5, 5] z, float[1, 3, 5, 5] InvStdDev) {
                 scale = Constant <value_ints=[3]> ()
                 B = Constant <value_ints=[3]> ()
@@ -187,7 +187,7 @@ def test_avoid_remove_non_trailing_unused_optional_outputs_layernorm(self):
     def test_remove_trailing_unused_optional_outputs_batchnorm(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 3, 5, 5] x, float[3] scale, float[3] B) => (float[1, 3, 5, 5] z) {
                 z, mean_out, var_out = BatchNormalization <training_mode=1> (x, scale, B, mean, var)
             }
@@ -204,7 +204,7 @@ def test_remove_trailing_unused_optional_outputs_batchnorm(self):
     def test_avoid_remove_used_optional_outputs_batchnorm(self):
         model = onnx.parser.parse_model(
             """
-            <ir_version: 7, opset_import: [ "" : 17]>
+            <ir_version: 10, opset_import: [ "" : 17]>
             agraph (float[1, 3, 5, 5] x, float[3] scale, float[3] B) => (float[1, 3, 5, 5] z, float[3] mean_out, float[3] var_out) {
                 z, mean_out, var_out = BatchNormalization <training_mode=1> (x, scale, B, mean, var)
             }

onnxscript/optimizer/__init__.py

@@ -14,13 +14,13 @@
 
 import onnx
 
+import onnxscript.ir.passes.common.unused_removal
 import onnxscript.optimizer._constant_folding as constant_folding
 import onnxscript.optimizer._legacy._optimizer as legacy_optimizer
 import onnxscript.optimizer._legacy.constant_folding as legacy_constant_folding
 from onnxscript import ir
 from onnxscript.optimizer._inliner import inline
 from onnxscript.optimizer._optimizer import optimize_ir
-from onnxscript.optimizer._remove_unused import remove_unused_nodes
 
 basic_constant_propagation = constant_folding.basic_constant_propagation
 fold_constants_ir = constant_folding.fold_constants
@@ -40,3 +40,31 @@ def fold_constants(model: ir.Model | onnx.ModelProto, *args, **kwargs) -> bool:
         return constant_folding.fold_constants(model, *args, **kwargs)
     else:
         return legacy_constant_folding.fold_constants(model, *args, **kwargs)
+
+
+def remove_unused_nodes(model: ir.Model | onnx.ModelProto) -> None:
+    """Removes unused nodes from a model inplace."""
+    if isinstance(model, ir.Model):
+        onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass()(model)
+    else:
+        model_ir = ir.serde.deserialize_model(model)
+        model_ir = onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass()(
+            model_ir
+        ).model
+        new_proto = ir.serde.serialize_model(model_ir)
+        model.Clear()
+        model.CopyFrom(new_proto)
+
+
+def remove_unused_functions(model: ir.Model | onnx.ModelProto) -> None:
+    """Removes unused functions from a model inplace."""
+    if isinstance(model, ir.Model):
+        onnxscript.ir.passes.common.unused_removal.RemoveUnusedFunctionsPass()(model)
+    else:
+        model_ir = ir.serde.deserialize_model(model)
+        model_ir = onnxscript.ir.passes.common.unused_removal.RemoveUnusedFunctionsPass()(
+            model_ir
+        ).model
+        new_proto = ir.serde.serialize_model(model_ir)
+        model.Clear()
+        model.CopyFrom(new_proto)

onnxscript/optimizer/_legacy/_optimizer.py

@@ -8,15 +8,14 @@
 import onnx
 import onnx.shape_inference
 
+import onnxscript.optimizer
 from onnxscript import rewriter
 from onnxscript.optimizer._legacy._simple_function_folding import (
     inline_functions_with_unused_outputs,
     inline_simple_functions,
 )
 from onnxscript.optimizer._legacy.constant_folding import fold_constants
 from onnxscript.optimizer._optimizer import _DEFAULT_REWRITE_RULES
-from onnxscript.optimizer._remove_unused import remove_unused_nodes
-from onnxscript.optimizer._remove_unused_function import remove_unused_functions
 
 logger = logging.getLogger(__name__)
 
@@ -71,9 +70,9 @@ def optimize(
             model, external_data_folder, onnx_shape_inference=onnx_shape_inference
         )
 
-        remove_unused_nodes(model)
+        onnxscript.optimizer.remove_unused_nodes(model)
         inline_simple_functions(model)
-        model = remove_unused_functions(model)
+        onnxscript.optimizer.remove_unused_functions(model)
         inline_functions_with_unused_outputs(model)
         # NOTE: This is general rewrite rules
         model = rewriter.rewrite(model, pattern_rewrite_rules=_DEFAULT_REWRITE_RULES)