onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/passes/_pass_infra.py

@@ -0,0 +1,351 @@
+# 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
+# for the ONNX IR.
+# The classes {PassResult and PassManager} are derived from
+# https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/fx/passes/infra/pass_base.py#L12
+# and
+# https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/fx/passes/infra/pass_manager.py#L147
+# The original code is licensed under the PyTorch License https://github.com/pytorch/pytorch/blob/main/LICENSE
+
+"""Passes infrastructure for the IR."""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+from collections.abc import Sequence
+from typing import Literal, final
+
+__all__ = [
+    "PassBase",
+    "Sequential",
+    "InPlacePass",
+    "FunctionalPass",
+    "PassManager",
+    "PassResult",
+    "functionalize",
+    # Errors
+    "InvariantError",
+    "PreconditionError",
+    "PostconditionError",
+    "PassError",
+]
+
+import abc
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class InvariantError(Exception):
+    """Raised when an invariant is violated."""
+
+
+class PreconditionError(InvariantError):
+    """Raised when a precondition is violated."""
+
+
+class PostconditionError(InvariantError):
+    """Raised when a postcondition is violated."""
+
+
+class PassError(RuntimeError):
+    """Raised when an error occurs during a pass."""
+
+
+@dataclasses.dataclass
+class PassResult:
+    """Result of a pass.
+
+    Attributes:
+        model: The transformed model.
+        modified: Whether the resulting model is different from the input model.
+    """
+
+    model: ir.Model
+    modified: bool
+
+
+class PassBase(abc.ABC):
+    """Base class for all passes.
+
+    ``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   |                  |                            |
+    +------------+------------------+----------------------------+
+    """
+
+    @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``.
+
+        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(result.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:
+        """The main entry point for the pass."""
+        ...
+
+    def requires(self, model: ir.Model) -> None:
+        """Pre-conditions for the pass.
+
+        This is optional to implement, will be called before call() if run by a pass manager.
+        """
+        del model  # Unused
+
+    def ensures(self, model: ir.Model) -> None:
+        """Post-conditions for the pass.
+
+        This is optional to implement, will be called after call() if run by a pass manager.
+        """
+        del model  # Unused
+
+
+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.
+
+    Example::
+        import onnx_ir as ir
+        import onnx_ir.passes.common as common_passes
+
+        passes = ir.passes.Sequential(
+            common_passes.DeduplicateHashedInitializersPass(size_limit=1024 * 1024),
+            common_passes.CommonSubexpressionEliminationPass(),
+            common_passes.ClearMetadataAndDocStringPass(),
+        )
+        result = passes(model)
+    """
+
+    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 Pass that runs a sequence of passes on a model.
+
+    Example::
+        import onnx_ir as ir
+        import onnx_ir.passes.common as common_passes
+
+        model = ir.load("model.onnx")
+        passes = ir.passes.PassManager(
+            [
+                # Pass managers can be nested
+                ir.passes.PassManager(
+                    [
+                        common_passes.DeduplicateHashedInitializersPass(size_limit=1024 * 1024),
+                        common_passes.CommonSubexpressionEliminationPass(),
+                    ],
+                    steps=2,
+                    early_stop=True,
+                ),
+                common_passes.ClearMetadataAndDocStringPass(),
+            ],
+            steps=2,
+            early_stop=False,
+        )
+
+        # Apply the passes to the model
+        result = passes(model)
+
+    Attributes:
+        passes: The passes to run.
+        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],
+        steps: int = 1,
+        early_stop: bool = True,
+    ):
+        # TODO(justinchuby): Implement constraints
+        super().__init__(*passes)
+        self.steps = steps
+        self.early_stop = early_stop
+
+    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):
+            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 and self.early_stop:
+                logger.info("PassManager: No more graph changes detected after step %s", step)
+                break
+        return PassResult(model, overall_modified)
+
+
+class _FunctionalPassWrapper(FunctionalPass):
+    def __init__(self, inner_pass: PassBase) -> None:
+        self._inner_pass = inner_pass
+
+    def call(self, model: ir.Model) -> PassResult:
+        return self._inner_pass(model.clone())
+
+
+def functionalize(pass_instance: PassBase) -> FunctionalPass:
+    """Produce a functional pass from a given pass.
+
+    A new functional pass is created that clones the input model before running the pass.
+
+    .. versionadded:: 0.1.14
+
+    Args:
+        pass_instance: The pass to convert.
+
+    Returns:
+        A functional pass.
+    """
+    return _FunctionalPassWrapper(pass_instance)

onnx_ir/passes/common/__init__.py

@@ -0,0 +1,54 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "AddDefaultAttributesPass",
+    "AddInitializersToInputsPass",
+    "CheckerPass",
+    "ClearMetadataAndDocStringPass",
+    "CommonSubexpressionEliminationPass",
+    "DeduplicateHashedInitializersPass",
+    "DeduplicateInitializersPass",
+    "IdentityEliminationPass",
+    "InlinePass",
+    "LiftConstantsToInitializersPass",
+    "LiftSubgraphInitializersToMainGraphPass",
+    "NameFixPass",
+    "OutputFixPass",
+    "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.default_attributes import AddDefaultAttributesPass
+from onnx_ir.passes.common.identity_elimination import IdentityEliminationPass
+from onnx_ir.passes.common.initializer_deduplication import (
+    DeduplicateHashedInitializersPass,
+    DeduplicateInitializersPass,
+)
+from onnx_ir.passes.common.inliner import InlinePass
+from onnx_ir.passes.common.naming import NameFixPass
+from onnx_ir.passes.common.onnx_checker import CheckerPass
+from onnx_ir.passes.common.output_fix import OutputFixPass
+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  # noqa: TID251
+
+
+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