onnx-ir 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

onnx_ir/passes/_pass_infra.py

@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
 #
 # This module implements some APIs described in
 # https://pytorch.org/executorch/stable/compiler-custom-compiler-passes.html
@@ -16,10 +16,14 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-from typing import Sequence
+from collections.abc import Sequence
+from typing import Literal, final
 
 __all__ = [
     "PassBase",
+    "Sequential",
+    "InPlacePass",
+    "FunctionalPass",
     "PassManager",
     "PassResult",
     # Errors
@@ -58,7 +62,7 @@ class PassResult:
 
     Attributes:
         model: The transformed model.
-        modified: Whether the model was modified.
+        modified: Whether the resulting model is different from the input model.
     """
 
     model: ir.Model
@@ -68,14 +72,89 @@ class PassResult:
 class PassBase(abc.ABC):
     """Base class for all passes.
 
-    Class attributes:
-        in_place: Whether the pass modifies the model in place.
+    ``in_place`` and ``changes_input`` properties and what they mean:
+
+    +------------+------------------+----------------------------+
+    |            | changes_inputs   | not changes_inputs         |
+    +------------+------------------+----------------------------+
+    | in_place   | in place         | Side-effect-only pass      |
+    +------------+------------------+----------------------------+
+    | not        | destructive      | functional                 |
+    | in_place   |                  |                            |
+    +------------+------------------+----------------------------+
     """
 
-    in_place: bool = True
+    @property
+    @abc.abstractmethod
+    def in_place(self) -> bool:
+        """Whether the pass modifies the model in place and returns it.
+
+        If True, the pass will return the same model object that was passed in.
+        If False, the pass will return a new model object.
+        """
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def changes_input(self) -> bool:
+        """Whether the pass modifies input model."""
+        raise NotImplementedError
+
+    @property
+    def destructive(self) -> bool:
+        """Whether the pass will destroy the input model when ``in_place=False``.
 
-    def __call__(self, model: ir.Model) -> PassResult:
-        return self.call(model)
+        A pass is destructive if it is not in place and it modifies the input model.
+        """
+        return not self.in_place and self.changes_input
+
+    def __call__(self, model_or_result: ir.Model | PassResult, /) -> PassResult:
+        if isinstance(model_or_result, PassResult):
+            model = model_or_result.model
+        else:
+            model = model_or_result
+        # Check preconditions
+        try:
+            self.requires(model)
+        except PreconditionError:
+            raise
+        except Exception as e:
+            raise PreconditionError(
+                f"Pre-condition for pass '{self.__class__.__name__}' failed"
+            ) from e
+
+        result = self.call(model)
+
+        # Check postconditions
+        try:
+            self.ensures(model)
+        except PostconditionError:
+            raise
+        except Exception as e:
+            raise PostconditionError(
+                f"Post-condition for pass '{self.__class__.__name__}' failed"
+            ) from e
+
+        if not isinstance(result, PassResult):
+            raise TypeError(
+                f"The result of the pass '{self.__class__.__name__}' should be type PassResult. "
+                "Please create one with ir.passes.PassResult()."
+            )
+
+        # Checks that the declared in-place property is respected
+        if self.in_place and result.model is not model:
+            raise PassError(
+                f"The pass '{self.__class__.__name__}' is declared in-place, "
+                "but the model returned is *not* the same object as the input model. "
+                "Pass developer: Pass should return the same model object or the in_place property should return False."
+            )
+        if not self.in_place and result.model is model:
+            raise PassError(
+                f"The pass '{self.__class__.__name__}' is declared not in-place, "
+                "but the model returned *is* the same object as the input model. "
+                "Pass developer: Pass should return a new model object or the in_place property should return True."
+            )
+        return result
 
     @abc.abstractmethod
     def call(self, model: ir.Model) -> PassResult:
@@ -97,76 +176,114 @@ class PassBase(abc.ABC):
         del model  # Unused
 
 
-class PassManager:
+class InPlacePass(PassBase):
+    """A pass that modifies the input model in place and returns it."""
+
+    @property
+    @final
+    def in_place(self) -> Literal[True]:
+        """An in-place pass is in place."""
+        return True
+
+    @property
+    @final
+    def changes_input(self) -> Literal[True]:
+        """An in-place pass changes the input model."""
+        return True
+
+
+class FunctionalPass(PassBase):
+    """A pass that returns a new model but does not modify the input model."""
+
+    @property
+    @final
+    def in_place(self) -> Literal[False]:
+        """A functional pass is not in place."""
+        return False
+
+    @property
+    @final
+    def changes_input(self) -> Literal[False]:
+        """A functional pass does not change the input model."""
+        return False
+
+
+class Sequential(PassBase):
+    """Run a sequence of passes in order."""
+
+    def __init__(self, *passes: PassBase):
+        if not passes:
+            raise ValueError("Sequential must take at least one pass")
+        self.passes = passes
+        self._in_place = all(pass_.in_place for pass_ in passes)
+        # The reason changes_inputs is decided by the first pass is that if the first pass is either in-place,
+        # or if it is not designed to be in-place but somehow changes the input (destructive),
+        # this pass sequence will change inputs.
+        self._changes_input = self.passes[0].changes_input or self.passes[0].in_place
+
+    @property
+    def in_place(self) -> bool:
+        return self._in_place
+
+    @property
+    def changes_input(self) -> bool:
+        return self._changes_input
+
+    def call(self, model: ir.Model) -> PassResult:
+        modified = False
+        for i, pass_ in enumerate(self.passes):
+            logger.debug("Running the %s-th pass '%s'", i, pass_)
+            try:
+                pass_result = pass_(model)
+            except Exception as e:
+                prev_pass_names = [str(p) for p in self.passes[:i]]
+                raise PassError(
+                    f"An error occurred when running the '{pass_}' pass after the "
+                    f"following passes: {prev_pass_names}"
+                ) from e
+
+            model = pass_result.model
+            modified = modified or pass_result.modified
+
+        return PassResult(model, modified)
+
+
+class PassManager(Sequential):
     """Pass manager for the IR.
 
-    The PassManager is a callable that runs a sequence of passes on a model.
+    The PassManager is a Pass that runs a sequence of passes on a model.
 
     Attributes:
         passes: The passes to run.
-        check_invariants: Whether to check invariants before and after each pass.
         steps: The number of times to run the passes.
+        early_stop: Whether to stop running the passes if the graph stops changing.
     """
 
     def __init__(
         self,
         passes: Sequence[PassBase],
-        check_invariants: bool = False,
         steps: int = 1,
+        early_stop: bool = True,
     ):
         # TODO(justinchuby): Implement constraints
-        self.passes = list(passes)
-        self.check_invariants = check_invariants
+        super().__init__(*passes)
         self.steps = steps
+        self.early_stop = early_stop
 
-    def __call__(self, model: ir.Model) -> PassResult:
+    def call(self, model: ir.Model) -> PassResult:
         """Run the set of passes `steps` number of times or until the graph stops changing."""
         overall_modified = False
         for step in range(self.steps):
-            step_result = self._run_one_step(model, step)
+            try:
+                # 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
             modified = step_result.modified
             overall_modified = overall_modified or modified
             # If the graph no longer changes, then we can stop running these passes
-            if not modified:
+            if not modified and self.early_stop:
                 logger.info("PassManager: No more graph changes detected after step %s", step)
                 break
         return PassResult(model, overall_modified)
-
-    def _run_one_step(self, model: ir.Model, step: int) -> PassResult:
-        modified = False
-        for i, pass_ in enumerate(self.passes):
-            logger.debug("Running the %s-th pass '%s', (step %s)", i, pass_, step)
-
-            # 1. Check preconditions
-            if self.check_invariants:
-                try:
-                    pass_.requires(model)
-                except Exception as e:
-                    raise PreconditionError(f"Pre-condition failed for {pass_}") from e
-
-            # 2. Run the pass
-            try:
-                pass_result = pass_(model)
-            except Exception as e:
-                prev_pass_names = [str(p) for p in self.passes[:i]]
-                raise PassError(
-                    f"An error occurred when running the '{pass_}' pass after the "
-                    f"following passes: {prev_pass_names} during step {step}"
-                ) from e
-            if not isinstance(pass_result, PassResult):
-                raise TypeError(
-                    f"The result of the pass {pass_} should be type PassResult."
-                    "Please create one with ir.passes.PassResult()."
-                )
-
-            model = pass_result.model
-            modified = modified or pass_result.modified
-
-            # 3. Check postconditions
-            if self.check_invariants:
-                try:
-                    pass_.ensures(model)
-                except Exception as e:
-                    raise PostconditionError(f"Post-condition failed for {pass_}") from e
-        return PassResult(model, modified)

onnx_ir/passes/common/__init__.py

@@ -0,0 +1,40 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "AddInitializersToInputsPass",
+    "CheckerPass",
+    "ClearMetadataAndDocStringPass",
+    "CommonSubexpressionEliminationPass",
+    "InlinePass",
+    "LiftConstantsToInitializersPass",
+    "LiftSubgraphInitializersToMainGraphPass",
+    "RemoveInitializersFromInputsPass",
+    "RemoveUnusedFunctionsPass",
+    "RemoveUnusedNodesPass",
+    "RemoveUnusedOpsetsPass",
+    "ShapeInferencePass",
+    "TopologicalSortPass",
+]
+
+from onnx_ir.passes.common.clear_metadata_and_docstring import (
+    ClearMetadataAndDocStringPass,
+)
+from onnx_ir.passes.common.common_subexpression_elimination import (
+    CommonSubexpressionEliminationPass,
+)
+from onnx_ir.passes.common.constant_manipulation import (
+    AddInitializersToInputsPass,
+    LiftConstantsToInitializersPass,
+    LiftSubgraphInitializersToMainGraphPass,
+    RemoveInitializersFromInputsPass,
+)
+from onnx_ir.passes.common.inliner import InlinePass
+from onnx_ir.passes.common.onnx_checker import CheckerPass
+from onnx_ir.passes.common.shape_inference import ShapeInferencePass
+from onnx_ir.passes.common.topological_sort import TopologicalSortPass
+from onnx_ir.passes.common.unused_removal import (
+    RemoveUnusedFunctionsPass,
+    RemoveUnusedNodesPass,
+    RemoveUnusedOpsetsPass,
+)

onnx_ir/passes/common/_c_api_utils.py

@@ -0,0 +1,76 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for interfacing with onnx C APIs."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Callable, TypeVar
+
+import onnx_ir as ir
+
+if TYPE_CHECKING:
+    import onnx
+
+
+logger = logging.getLogger(__name__)
+# Temporarily remove initializers larger than this size to keep model size down
+# for the onnx.shape_inference call because it needs to serialize the model
+_BIG_TENSOR_SIZE_LIMIT = 1000  # 1KB
+_R = TypeVar("_R")
+
+
+def call_onnx_api(func: Callable[[onnx.ModelProto], _R], model: ir.Model) -> _R:
+    """Call an ONNX C API function by temporarily removing initializers.
+
+    This is necessary because the ONNX C API does not support large models
+    with initializers that have large tensor values. The input model is left
+    unchanged no matter the call succeeds or not.
+
+    Args:
+        func: Partially applied function that takes a model proto and returns anything.
+        model: The IR model to pass to the API function.
+
+    Returns:
+        The resulting ModelProto that contains the result of the API call.
+    """
+    # Store the original initializer values so they can be restored
+    initializer_values = tuple(model.graph.initializers.values())
+    tensors = {v.name: v.const_value for v in initializer_values}
+    original_inputs_len = len(model.graph.inputs)
+
+    # Turn the initializers into inputs and clear the initializers
+    # to limit the model size
+    for initializer in initializer_values:
+        # Make sure the initializer has its shape/type set
+        assert initializer.const_value is not None
+        if initializer.shape is None:
+            initializer.shape = initializer.const_value.shape  # type: ignore[assignment]
+        if initializer.dtype is None:
+            initializer.dtype = initializer.const_value.dtype
+        if initializer not in model.graph.inputs:
+            model.graph.inputs.append(initializer)
+        if initializer.const_value.nbytes > _BIG_TENSOR_SIZE_LIMIT:
+            # Temporarily remove the initializer value to reduce model size
+            # for onnx.shape_inference
+            initializer.const_value = None
+            assert initializer.name is not None
+            model.graph.initializers.pop(initializer.name)
+
+    proto = ir.serde.serialize_model(model)
+
+    try:
+        # Call the ONNX C API function
+        result = func(proto)
+    finally:
+        # Restore the original initializer values so the model is unchanged
+        for initializer in initializer_values:
+            initializer.const_value = tensors[initializer.name]
+            model.graph.register_initializer(initializer)
+
+        # Restore the original inputs
+        inputs = model.graph.inputs[:original_inputs_len]
+        model.graph.inputs.clear()
+        model.graph.inputs.extend(inputs)
+
+    return result

onnx_ir/passes/common/clear_metadata_and_docstring.py

@@ -0,0 +1,60 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Clear all metadata and docstring from the model, graphs, nodes, and functions."""
+
+from __future__ import annotations
+
+__all__ = [
+    "ClearMetadataAndDocStringPass",
+]
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class ClearMetadataAndDocStringPass(ir.passes.InPlacePass):
+    """Clear all metadata and docstring from the model, graphs, nodes, and functions."""
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        # 0. TODO: Should we clean model metadata and docstring?
+
+        # 1. Clean up the graph and the belonged nodes metadata properties
+        modified = self._clear_graph_or_function_metadata_and_docstring(model.graph)
+
+        # 2. Clean up all of the functions metadata properties
+        for function in model.functions.values():
+            modified = (
+                self._clear_graph_or_function_metadata_and_docstring(function) or modified
+            )
+        return ir.passes.PassResult(model, modified=modified)
+
+    def _clear_graph_or_function_metadata_and_docstring(
+        self,
+        graph_or_function: ir.Graph | ir.Function,
+    ) -> bool:
+        """Clear metadata and docstring from the graph or function."""
+        checked_graphs_or_functions: set[ir.Graph | ir.Function] = set()
+        modified = False
+        # Clean up all of the nodes metadata properties
+        for node in ir.traversal.RecursiveGraphIterator(graph_or_function):
+            if node.metadata_props:
+                modified = True
+                logger.debug("Removed metadata from %s nodes", node.name)
+            node.metadata_props.clear()
+            node.doc_string = None
+
+            # Clean up the owning graph/function metadata properties
+            # and doc_string if the graph/function is not already checked
+            assert node.graph is not None
+            if node.graph not in checked_graphs_or_functions and (
+                node.graph.metadata_props or node.graph.doc_string
+            ):
+                modified = True
+                logger.debug("Removed metadata from %s graph/function", node.graph.name)
+                node.graph.metadata_props.clear()
+                node.graph.doc_string = None
+                checked_graphs_or_functions.add(node.graph)
+        return modified

onnx_ir/passes/common/common_subexpression_elimination.py

@@ -0,0 +1,177 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Eliminate common subexpression in ONNX graphs."""
+
+from __future__ import annotations
+
+__all__ = [
+    "CommonSubexpressionEliminationPass",
+]
+
+import logging
+from collections.abc import Sequence
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class CommonSubexpressionEliminationPass(ir.passes.InPlacePass):
+    """Eliminate common subexpression in ONNX graphs."""
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Return the same ir.Model but with CSE applied to the graph."""
+        modified = False
+        graph = model.graph
+
+        modified = _eliminate_common_subexpression(graph, modified)
+
+        return ir.passes.PassResult(
+            model,
+            modified=modified,
+        )
+
+
+def _eliminate_common_subexpression(graph: ir.Graph, modified: bool) -> bool:
+    """Eliminate common subexpression in ONNX graphs."""
+    # node to node identifier, length of outputs, inputs, and attributes
+    existing_node_info_to_the_node: dict[
+        tuple[
+            ir.OperatorIdentifier,
+            int,  # len(outputs)
+            tuple[int, ...],  # input ids
+            tuple[tuple[str, object], ...],  # attributes
+        ],
+        ir.Node,
+    ] = {}
+
+    for node in graph:
+        # Skip control flow ops like Loop and If.
+        control_flow_op: bool = False
+        # Use equality to check if the node is a common subexpression.
+        attributes = {}
+        for k, v in node.attributes.items():
+            # TODO(exporter team): CSE subgraphs.
+            # NOTE: control flow ops like Loop and If won't be CSEd
+            # because attribute: graph won't match.
+            if v.type in (ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS):
+                control_flow_op = True
+                logger.debug("Skipping control flow op %s", node)
+            # The attribute value could be directly taken from the original
+            # protobuf, so we need to make a copy of it.
+            value = v.value
+            if v.type in (
+                ir.AttributeType.INTS,
+                ir.AttributeType.FLOATS,
+                ir.AttributeType.STRINGS,
+            ):
+                # For INT, FLOAT and STRING attributes, we convert them to tuples
+                # to ensure they are hashable.
+                value = tuple(value)
+            attributes[k] = value
+
+        if control_flow_op:
+            # If the node is a control flow op, we skip it.
+            logger.debug("Skipping control flow op %s", node)
+            continue
+
+        if _is_non_deterministic_op(node):
+            # If the node is a non-deterministic op, we skip it.
+            logger.debug("Skipping non-deterministic op %s", node)
+            continue
+
+        node_info = (
+            node.op_identifier(),
+            len(node.outputs),
+            tuple(id(input) for input in node.inputs),
+            tuple(sorted(attributes.items())),
+        )
+        # Check if the node is a common subexpression.
+        if node_info in existing_node_info_to_the_node:
+            # If it is, this node has an existing node with the same
+            # operator, number of outputs, inputs, and attributes.
+            # We replace the node with the existing node.
+            modified = True
+            existing_node = existing_node_info_to_the_node[node_info]
+            _remove_node_and_replace_values(
+                graph,
+                remove_node=node,
+                remove_values=node.outputs,
+                new_values=existing_node.outputs,
+            )
+            logger.debug("Reusing node %s", existing_node)
+        else:
+            # If it is not, add to the mapping.
+            existing_node_info_to_the_node[node_info] = node
+    return modified
+
+
+def _remove_node_and_replace_values(
+    graph: ir.Graph,
+    /,
+    remove_node: ir.Node,
+    remove_values: Sequence[ir.Value],
+    new_values: Sequence[ir.Value],
+) -> None:
+    """Replaces nodes and values in the graph or function.
+
+    Args:
+        graph: The graph to replace nodes and values in.
+        remove_node: The node to remove.
+        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))
+        for idx, graph_output in enumerate(graph.outputs):
+            if graph_output in replacement_mapping:
+                new_value = replacement_mapping[graph_output]
+                if new_value.is_graph_output() or new_value.is_graph_input():
+                    # If the new value is also a graph input/output, we need to
+                    # create a Identity node to preserve the remove_value and
+                    # prevent from changing new_value name.
+                    identity_node = ir.node(
+                        "Identity",
+                        inputs=[new_value],
+                        outputs=[
+                            ir.Value(
+                                name=graph_output.name,
+                                type=graph_output.type,
+                                shape=graph_output.shape,
+                            )
+                        ],
+                    )
+                    # reuse the name of the graph output
+                    graph.outputs[idx] = identity_node.outputs[0]
+                    graph.insert_before(
+                        remove_node,
+                        identity_node,
+                    )
+                else:
+                    # if new_value is not graph output, we just
+                    # update it to use old_value name.
+                    new_value.name = graph_output.name
+                    graph.outputs[idx] = new_value
+
+    graph.remove(remove_node, safe=True)
+
+
+def _is_non_deterministic_op(node: ir.Node) -> bool:
+    non_deterministic_ops = frozenset(
+        {
+            "RandomUniform",
+            "RandomNormal",
+            "RandomUniformLike",
+            "RandomNormalLike",
+            "Multinomial",
+        }
+    )
+    return node.op_type in non_deterministic_ops and _is_onnx_domain(node.domain)
+
+
+def _is_onnx_domain(d: str) -> bool:
+    """Check if the domain is the ONNX domain."""
+    return d == ""