onnx-ir 0.0.1__py3-none-any.whl → 0.1.0__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,36 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "AddInitializersToInputsPass",
+    "CheckerPass",
+    "ClearMetadataAndDocStringPass",
+    "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.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/constant_manipulation.py

@@ -0,0 +1,232 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Lift constants to initializers."""
+
+from __future__ import annotations
+
+__all__ = [
+    "AddInitializersToInputsPass",
+    "LiftConstantsToInitializersPass",
+    "LiftSubgraphInitializersToMainGraphPass",
+    "RemoveInitializersFromInputsPass",
+]
+
+import logging
+
+import numpy as np
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class LiftConstantsToInitializersPass(ir.passes.InPlacePass):
+    """Lift constants to initializers.
+
+    Attributes:
+        lift_all_constants: Whether to lift all Constant nodes, including those that does not contain a tensor attribute (e.g. with value_ints etc.)
+            Default to False, where only Constants with the ``value`` attribute are lifted.
+        size_limit: The minimum size of the tensor to be lifted. If the tensor contains
+            number of elements less than size_limit, it will not be lifted. Default is 16.
+    """
+
+    def __init__(self, lift_all_constants: bool = False, size_limit: int = 16):
+        super().__init__()
+        self.lift_all_constants = lift_all_constants
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = 0
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            assert node.graph is not None
+            if node.op_type != "Constant" or node.domain not in ("", "onnx.ai"):
+                continue
+            if node.outputs[0].is_graph_output():
+                logger.debug(
+                    "Constant node '%s' is used as output, so it can't be lifted.", node.name
+                )
+                continue
+            constant_node_attribute = set(node.attributes.keys())
+            if len(constant_node_attribute) != 1:
+                logger.debug(
+                    "Invalid constant node '%s' has more than one attribute", node.name
+                )
+                continue
+
+            attr_name, attr_value = next(iter(node.attributes.items()))
+            initializer_name = node.outputs[0].name
+            assert initializer_name is not None
+            assert isinstance(attr_value, ir.Attr)
+            tensor = self._constant_node_attribute_to_tensor(
+                node, attr_name, attr_value, initializer_name
+            )
+            if tensor is None:
+                # The reason of None is logged in _constant_node_attribute_to_tensor
+                continue
+            # Register an initializer with the tensor value
+            initializer = ir.Value(
+                name=initializer_name,
+                shape=tensor.shape,  # type: ignore[arg-type]
+                type=ir.TensorType(tensor.dtype),
+                const_value=tensor,
+            )
+            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.graph.remove(node, safe=True)
+            count += 1
+            logger.debug(
+                "Converted constant node '%s' to initializer '%s'", node.name, initializer_name
+            )
+        if count:
+            logger.debug("Lifted %s constants to initializers", count)
+        return ir.passes.PassResult(model, modified=bool(count))
+
+    def _constant_node_attribute_to_tensor(
+        self, node, attr_name: str, attr_value: ir.Attr, initializer_name: str
+    ) -> ir.TensorProtocol | None:
+        """Convert constant node attribute to tensor."""
+        if not self.lift_all_constants and attr_name != "value":
+            logger.debug(
+                "Constant node '%s' has non-tensor attribute '%s'", node.name, attr_name
+            )
+            return None
+
+        tensor: ir.TensorProtocol
+        if attr_name == "value":
+            tensor = attr_value.as_tensor()
+        elif attr_name == "value_int":
+            tensor = ir.tensor(
+                attr_value.as_int(), dtype=ir.DataType.INT64, name=initializer_name
+            )
+        elif attr_name == "value_ints":
+            tensor = ir.tensor(
+                attr_value.as_ints(), dtype=ir.DataType.INT64, name=initializer_name
+            )
+        elif attr_name == "value_float":
+            tensor = ir.tensor(
+                attr_value.as_float(), dtype=ir.DataType.FLOAT, name=initializer_name
+            )
+        elif attr_name == "value_floats":
+            tensor = ir.tensor(
+                attr_value.as_floats(), dtype=ir.DataType.FLOAT, name=initializer_name
+            )
+        elif attr_name in ("value_string", "value_strings"):
+            tensor = ir.StringTensor(
+                np.array(attr_value.value, dtype=np.bytes_), name=initializer_name
+            )
+        else:
+            raise ValueError(
+                f"Unsupported constant node '{node.name}' attribute '{attr_name}'"
+            )
+
+        if tensor.size < self.size_limit:
+            logger.debug(
+                "Tensor from node '%s' has less than %s elements",
+                node.name,
+                self.size_limit,
+            )
+            return None
+        return tensor
+
+
+class LiftSubgraphInitializersToMainGraphPass(ir.passes.InPlacePass):
+    """Lift subgraph initializers to main graph.
+
+    This pass lifts the initializers of a subgraph to the main graph.
+    It is used to ensure that the initializers are available in the main graph
+    for further processing or optimization.
+
+    Initializers that are also graph inputs will not be lifted.
+
+    Preconditions:
+        - All initializers in the model must have unique names across the main graph and subgraphs.
+    """
+
+    def requires(self, model: ir.Model) -> None:
+        """Ensure all initializer names are unique."""
+        registered_initializer_names: set[str] = set()
+        duplicated_initializers: list[ir.Value] = []
+        for graph in model.graphs():
+            for initializer in graph.initializers.values():
+                if initializer.name is None:
+                    raise ir.passes.PreconditionError(
+                        f"Initializer name is None. Please ensure all initializers have unique names: {initializer!r}"
+                    )
+                if initializer.name in registered_initializer_names:
+                    duplicated_initializers.append(initializer)
+                else:
+                    registered_initializer_names.add(initializer.name)
+        if duplicated_initializers:
+            raise ir.passes.PreconditionError(
+                "Found duplicated initializers in the model. "
+                "Initializer name must be unique across the main graph and subgraphs. "
+                "Please ensure all initializers have unique names. Duplicated: "
+                f"{duplicated_initializers!r}"
+            )
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = 0
+        for graph in model.graphs():
+            if graph is model.graph:
+                continue
+            for name in tuple(graph.initializers):
+                initializer = graph.initializers[name]
+                if initializer.is_graph_input():
+                    # Skip the ones that are also graph inputs
+                    logger.debug(
+                        "Initializer '%s' is also a graph input, so it can't be lifted",
+                        initializer.name,
+                    )
+                    continue
+                # Remove the initializer from the subgraph
+                graph.initializers.pop(name)
+                model.graph.register_initializer(initializer)
+                count += 1
+                logger.debug(
+                    "Lifted initializer '%s' from subgraph '%s' to main graph",
+                    initializer.name,
+                    graph.name,
+                )
+        return ir.passes.PassResult(model, modified=bool(count))
+
+
+class RemoveInitializersFromInputsPass(ir.passes.InPlacePass):
+    """Remove initializers from inputs.
+
+    This pass finds all graph inputs that have a const_value and removes them from the graph.inputs list.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = 0
+        for graph in model.graphs():
+            initializers = set(graph.initializers.values())
+            new_inputs = []
+            for input_value in graph.inputs:
+                if input_value in initializers:
+                    count += 1
+                else:
+                    new_inputs.append(input_value)
+            graph.inputs.clear()
+            graph.inputs.extend(new_inputs)
+        logger.info("Removed %s initializers from graph inputs", count)
+        return ir.passes.PassResult(model, modified=bool(count))
+
+
+class AddInitializersToInputsPass(ir.passes.InPlacePass):
+    """Add initializers to inputs.
+
+    This pass finds all initializers and adds them to the graph.inputs list if they are not already present.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = 0
+        for graph in model.graphs():
+            inputs_set = set(graph.inputs)
+            for initializer in graph.initializers.values():
+                if initializer not in inputs_set:
+                    graph.inputs.append(initializer)
+                    count += 1
+        logger.info("Added %s initializers to graph inputs", count)
+        return ir.passes.PassResult(model, modified=bool(count))