onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/passes/common/common_subexpression_elimination.py

@@ -0,0 +1,207 @@
+# 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.
+
+    .. versionadded:: 0.1.1
+
+    .. versionchanged:: 0.1.3
+        Constant nodes with values smaller than ``size_limit`` will be CSE'd.
+
+    Attributes:
+        size_limit: The maximum size of the tensor to be csed. If the tensor contains
+            number of elements larger than size_limit, it will not be cse'd. Default is 10.
+
+    """
+
+    def __init__(self, size_limit: int = 10):
+        """Initialize the CommonSubexpressionEliminationPass."""
+        super().__init__()
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Return the same ir.Model but with CSE applied to the graph."""
+        graph = model.graph
+        modified = self._eliminate_common_subexpression(graph)
+
+        return ir.passes.PassResult(
+            model,
+            modified=modified,
+        )
+
+    def _eliminate_common_subexpression(self, graph: ir.Graph) -> bool:
+        """Eliminate common subexpression in ONNX graphs."""
+        modified: bool = False
+        # 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
+            # Skip large tensors to avoid cse weights and bias.
+            large_tensor: 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
+                    break
+                # 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)
+                elif v.type is ir.AttributeType.TENSOR:
+                    if value.size > self.size_limit:
+                        # If the tensor is larger than the size limit, we skip it.
+                        large_tensor = True
+                        break
+                    np_value = value.numpy()
+
+                    value = (np_value.shape, str(np_value.dtype), np_value.tobytes())
+                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 large_tensor:
+                # If the node has a large tensor, we skip it.
+                logger.debug("Skipping large tensor in node %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.
+    """
+    # 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
+
+    # Reconnect the users of the deleted values to use the new values
+    ir.convenience.replace_all_uses_with(remove_values, new_values)
+
+    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 == ""

onnx_ir/passes/common/constant_manipulation.py

@@ -0,0 +1,230 @@
+# 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,
+                # Preserve metadata from Constant value into the onnx model
+                metadata_props=node.outputs[0].metadata_props.copy(),
+            )
+            # Preserve value meta from the Constant output for intermediate analysis
+            initializer.meta.update(node.outputs[0].meta)
+
+            assert node.graph is not None
+            node.graph.register_initializer(initializer)
+            # Replace the constant node with the initializer
+            node.outputs[0].replace_all_uses_with(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.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        count = 0
+        registered_initializer_names: dict[str, int] = {}
+        for graph in model.graphs():
+            if graph is model.graph:
+                continue
+            for name in tuple(graph.initializers):
+                assert name is not None
+                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
+                if initializer.is_graph_output():
+                    logger.debug(
+                        "Initializer '%s' is used as output, so it can't be lifted",
+                        initializer.name,
+                    )
+                    continue
+                # Remove the initializer from the subgraph
+                graph.initializers.pop(name)
+                # To avoid name conflicts, we need to rename the initializer
+                # to a unique name in the main graph
+                new_name = name
+                while new_name in model.graph.initializers:
+                    if name in registered_initializer_names:
+                        registered_initializer_names[name] += 1
+                    else:
+                        registered_initializer_names[name] = 1
+                    new_name = f"{name}_{registered_initializer_names[name]}"
+                initializer.name = new_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))

onnx_ir/passes/common/default_attributes.py

@@ -0,0 +1,99 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Add default attributes to nodes that are missing optional attributes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "AddDefaultAttributesPass",
+]
+
+import logging
+
+import onnx  # noqa: TID251
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def _has_valid_default(attr_def: onnx.defs.OpSchema.Attribute) -> bool:
+    """Check if an attribute definition has a valid default value."""
+    return bool(
+        attr_def.default_value and attr_def.default_value.type != onnx.AttributeProto.UNDEFINED
+    )
+
+
+class AddDefaultAttributesPass(ir.passes.InPlacePass):
+    """Add default values for optional attributes that are not present in nodes.
+
+    This pass iterates through all nodes in the model and for each node:
+    1. Gets the ONNX schema for the operator
+    2. For each optional attribute with a default value in the schema
+    3. If the attribute is not present in the node, adds it with the default value
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Main entry point for the add default attributes pass."""
+        modified = False
+
+        # Process all nodes in the model graph and subgraphs
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            if _add_default_attributes_to_node(node, model.graph.opset_imports):
+                modified = True
+
+        # Process nodes in functions
+        for function in model.functions.values():
+            for node in ir.traversal.RecursiveGraphIterator(function):
+                if _add_default_attributes_to_node(node, model.graph.opset_imports):
+                    modified = True
+
+        if modified:
+            logger.info("AddDefaultAttributes pass modified the model")
+
+        return ir.passes.PassResult(model, modified=modified)
+
+
+def _add_default_attributes_to_node(node: ir.Node, opset_imports: dict[str, int]) -> bool:
+    """Add default attributes to a single node. Returns True if modified."""
+    # Get the operator schema
+    if node.version is not None:
+        opset_version = node.version
+    elif node.domain in opset_imports:
+        opset_version = opset_imports[node.domain]
+    else:
+        logger.warning(
+            "OpSet version for domain '%s' not found. Skipping node %s",
+            node.domain,
+            node,
+        )
+        return False
+
+    try:
+        op_schema = onnx.defs.get_schema(node.op_type, opset_version, domain=node.domain)
+    except onnx.defs.SchemaError:
+        logger.debug(
+            "Schema not found for %s, skipping default attribute addition",
+            node,
+        )
+        return False
+
+    modified = False
+    # Iterate through all attributes in the schema
+    for attr_name, attr_def in op_schema.attributes.items():
+        # Skip if attribute is required or already present in the node
+        if attr_def.required or attr_name in node.attributes:
+            continue
+
+        # Skip if attribute doesn't have a default value
+        if not _has_valid_default(attr_def):
+            continue
+
+        # Create an IR Attr from the ONNX AttributeProto default value
+        default_attr_proto = attr_def.default_value
+        default_attr = ir.serde.deserialize_attribute(default_attr_proto)
+        node.attributes[attr_name] = default_attr
+        logger.debug("Added default attribute '%s' to node %s", attr_name, node)
+        modified = True
+
+    return modified

onnx_ir/passes/common/identity_elimination.py

@@ -0,0 +1,120 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Identity elimination pass for removing redundant Identity nodes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "IdentityEliminationPass",
+]
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def _merge_shapes(shape1: ir.Shape | None, shape2: ir.Shape | None) -> ir.Shape | None:
+    def merge_dims(dim1, dim2):
+        if dim1 == dim2:
+            return dim1
+        if not isinstance(dim1, ir.SymbolicDim):
+            return dim1  # Prefer int value over symbolic dim
+        if not isinstance(dim2, ir.SymbolicDim):
+            return dim2
+        if dim1.value is None:
+            return dim2
+        return dim1
+
+    if shape1 is None:
+        return shape2
+    if shape2 is None:
+        return shape1
+    if len(shape1) != len(shape2):
+        raise ValueError(
+            f"Shapes must have the same rank, got {len(shape1)} and {len(shape2)}."
+        )
+    return ir.Shape([merge_dims(dim1, dim2) for dim1, dim2 in zip(shape1, shape2)])
+
+
+class IdentityEliminationPass(ir.passes.InPlacePass):
+    """Pass for eliminating redundant Identity nodes.
+
+    This pass removes Identity nodes according to the following rules:
+
+    1. For any node of the form `y = Identity(x)`, where `y` is not an output
+       of any graph, replace all uses of `y` with a use of `x`, and remove the node.
+    2. If `y` is an output of a graph, and `x` is not an input of any graph,
+       we can still do the elimination, but the value `x` should be renamed to be `y`.
+    3. If `y` is a graph-output and `x` is a graph-input, we cannot eliminate
+       the node. It should be retained.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Main entry point for the identity elimination pass."""
+        modified = False
+
+        # Use RecursiveGraphIterator to process all nodes in the model graph and subgraphs
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            if self._try_eliminate_identity_node(node):
+                modified = True
+
+        # Process nodes in functions
+        for function in model.functions.values():
+            for node in ir.traversal.RecursiveGraphIterator(function):
+                if self._try_eliminate_identity_node(node):
+                    modified = True
+
+        if modified:
+            logger.info("Identity elimination pass modified the model")
+
+        return ir.passes.PassResult(model, modified=modified)
+
+    def _try_eliminate_identity_node(self, node: ir.Node) -> bool:
+        """Try to eliminate a single identity node. Returns True if modified."""
+        if node.op_type != "Identity" or node.domain != "":
+            return False
+
+        if len(node.inputs) != 1 or len(node.outputs) != 1:
+            # Invalid Identity node, skip
+            return False
+
+        input_value = node.inputs[0]
+        output_value = node.outputs[0]
+
+        if input_value is None:
+            # Cannot eliminate if input is None
+            return False
+
+        # Get the graph that contains this node
+        graph_like = node.graph
+        assert graph_like is not None, "Node must be in a graph"
+
+        output_is_graph_output = output_value.is_graph_output()
+        input_is_graph_input = input_value.is_graph_input()
+
+        # Case 3: Both output is graph output and input is graph input - keep the node
+        if output_is_graph_output and input_is_graph_input:
+            return False
+
+        # Copy over shape/type if the output has more complete information
+        input_value.shape = _merge_shapes(input_value.shape, output_value.shape)
+        if input_value.type is None:
+            input_value.type = output_value.type
+
+        # Case 1 & 2 (merged): Eliminate the identity node
+        # Replace all uses of output with input
+        ir.convenience.replace_all_uses_with(
+            output_value, input_value, replace_graph_outputs=True
+        )
+
+        # If output is a graph output, we need to rename input and update graph outputs
+        if output_is_graph_output:
+            # Update the input value to have the output's name
+            input_value.name = output_value.name
+
+        # Remove the identity node
+        graph_like.remove(node, safe=True)
+        logger.debug("Eliminated identity node: %s", node)
+        return True