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

onnx_ir/passes/common/constant_manipulation.py

@@ -0,0 +1,217 @@
+# 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.
+    """
+
+    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):
+                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)
+                # To avoid name conflicts, we need to rename the initializer
+                # to a unique name in the main graph
+                if name in registered_initializer_names:
+                    name_count = registered_initializer_names[name]
+                    initializer.name = f"{name}_{name_count}"
+                    registered_initializer_names[name] = name_count + 1
+                else:
+                    assert initializer.name is not None
+                    registered_initializer_names[initializer.name] = 1
+                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/inliner.py

@@ -0,0 +1,332 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Implementation of an inliner for onnx_ir."""
+
+from __future__ import annotations
+
+import dataclasses
+
+__all__ = ["InlinePass", "InlinePassResult"]
+
+from collections import defaultdict
+from collections.abc import Iterable, Mapping, Sequence
+
+import onnx_ir as ir
+import onnx_ir.convenience as _ir_convenience
+
+# A replacement for a node specifies a list of nodes that replaces the original node,
+# and a list of values that replaces the original node's outputs.
+
+NodeReplacement = tuple[Sequence[ir.Node], Sequence[ir.Value]]
+
+# A call stack is a list of identifiers of call sites, where the first element is the
+# outermost call site, and the last element is the innermost call site. This is used
+# primarily for generating unique names for values in the inlined functions.
+CallSiteId = str
+CallStack = list[CallSiteId]
+
+
+def _make_unique_name(name: str, callstack: CallStack, used_names: set[str]) -> str:  # pylint: disable=unused-argument
+    """Generate a unique name from a name, calling-context, and set of used names.
+
+    If there is a name clash, we add a numeric suffix to the name to make
+    it unique. We use the same strategy to make node names unique.
+
+    TODO: We can use the callstack in generating a name for a value X in a function
+    that is inlined into a graph. This is not yet implemented. Using the full callstack
+    leads to very long and hard to read names. Some investigation is needed to find
+    a good naming strategy that will produce useful names for debugging.
+    """
+    candidate = name
+    i = 1
+    while candidate in used_names:
+        i += 1
+        candidate = f"{name}_{i}"
+    used_names.add(candidate)
+    return candidate
+
+
+class _CopyReplace:
+    """Utilities for creating a copy of IR objects with substitutions for attributes/input values."""
+
+    def __init__(
+        self,
+        inliner: InlinePass,
+        attr_map: Mapping[str, ir.Attr],
+        value_map: dict[ir.Value, ir.Value | None],
+        metadata_props: dict[str, str],
+        call_stack: CallStack,
+    ) -> None:
+        self._inliner = inliner
+        self._value_map = value_map
+        self._attr_map = attr_map
+        self._metadata_props = metadata_props
+        self._call_stack = call_stack
+
+    def clone_value(self, value: ir.Value) -> ir.Value | None:
+        if value in self._value_map:
+            return self._value_map[value]
+        # If the value is not in the value map, it must be a graph input.
+        assert value.producer() is None, f"Value {value} has no entry in the value map"
+        new_value = ir.Value(
+            name=value.name,
+            type=value.type,
+            shape=value.shape,
+            doc_string=value.doc_string,
+            const_value=value.const_value,
+        )
+        self._value_map[value] = new_value
+        return new_value
+
+    def clone_optional_value(self, value: ir.Value | None) -> ir.Value | None:
+        if value is None:
+            return None
+        return self.clone_value(value)
+
+    def clone_attr(self, key: str, attr: ir.Attr) -> ir.Attr | None:
+        if not attr.is_ref():
+            if attr.type == ir.AttributeType.GRAPH:
+                graph = self.clone_graph(attr.as_graph())
+                return ir.Attr(key, ir.AttributeType.GRAPH, graph, doc_string=attr.doc_string)
+            elif attr.type == ir.AttributeType.GRAPHS:
+                graphs = [self.clone_graph(graph) for graph in attr.as_graphs()]
+                return ir.Attr(
+                    key, ir.AttributeType.GRAPHS, graphs, doc_string=attr.doc_string
+                )
+            return attr
+        assert attr.is_ref()
+        ref_attr_name = attr.ref_attr_name
+        assert ref_attr_name is not None, "Reference attribute must have a name"
+        if ref_attr_name in self._attr_map:
+            ref_attr = self._attr_map[ref_attr_name]
+            if not ref_attr.is_ref():
+                return ir.Attr(
+                    key, ref_attr.type, ref_attr.value, doc_string=ref_attr.doc_string
+                )
+            assert ref_attr.ref_attr_name is not None
+            return ir.RefAttr(
+                key, ref_attr.ref_attr_name, ref_attr.type, doc_string=ref_attr.doc_string
+            )
+        # Note that if a function has an attribute-parameter X, and a call (node) to the function
+        # has no attribute X, all references to X in nodes inside the function body will be
+        # removed. This is just the ONNX representation of optional-attributes.
+        return None
+
+    def clone_node(self, node: ir.Node) -> ir.Node:
+        new_inputs = [self.clone_optional_value(input) for input in node.inputs]
+        new_attributes = [
+            new_value
+            for key, value in node.attributes.items()
+            if (new_value := self.clone_attr(key, value)) is not None
+        ]
+        new_name = node.name
+        if new_name is not None:
+            new_name = _make_unique_name(
+                new_name, self._call_stack, self._inliner.used_node_names
+            )
+
+        new_metadata = {**self._metadata_props, **node.metadata_props}
+        # TODO: For now, node metadata overrides callnode metadata if there is a conflict.
+        # Do we need to preserve both?
+
+        new_node = ir.Node(
+            node.domain,
+            node.op_type,
+            new_inputs,
+            new_attributes,
+            overload=node.overload,
+            num_outputs=len(node.outputs),
+            graph=None,
+            name=new_name,
+            doc_string=node.doc_string,  # type: ignore
+            metadata_props=new_metadata,
+        )
+        new_outputs = new_node.outputs
+        for i, output in enumerate(node.outputs):
+            self._value_map[output] = new_outputs[i]
+            old_name = output.name if output.name is not None else f"output_{i}"
+            new_outputs[i].name = _make_unique_name(
+                old_name, self._call_stack, self._inliner.used_value_names
+            )
+
+        self._inliner.node_context[new_node] = self._call_stack
+
+        return new_node
+
+    def clone_graph(self, graph: ir.Graph) -> ir.Graph:
+        input_values = [self.clone_value(v) for v in graph.inputs]
+        nodes = [self.clone_node(node) for node in graph]
+        initializers = [self.clone_value(init) for init in graph.initializers.values()]
+        output_values = [
+            self.clone_value(v) for v in graph.outputs
+        ]  # Looks up already cloned values
+
+        return ir.Graph(
+            input_values,  # type: ignore
+            output_values,  # type: ignore
+            nodes=nodes,
+            initializers=initializers,  # type: ignore
+            doc_string=graph.doc_string,
+            opset_imports=graph.opset_imports,
+            name=graph.name,
+            metadata_props=graph.metadata_props,
+        )
+
+
+def _abbreviate(
+    function_ids: Iterable[ir.OperatorIdentifier],
+) -> dict[ir.OperatorIdentifier, str]:
+    """Create a short unambiguous abbreviation for all function ids."""
+
+    def id_abbreviation(id: ir.OperatorIdentifier) -> str:
+        """Create a short unambiguous abbreviation for a function id."""
+        domain, name, overload = id
+        # Omit the domain, if it remains unambiguous after omitting it.
+        if any(x[0] != domain and x[1] == name and x[2] == overload for x in function_ids):
+            short_domain = domain + "_"
+        else:
+            short_domain = ""
+        if overload != "":
+            return short_domain + name + "_" + overload
+        return short_domain + name
+
+    return {id: id_abbreviation(id) for id in function_ids}
+
+
+@dataclasses.dataclass
+class InlinePassResult(ir.passes.PassResult):
+    id_count: dict[ir.OperatorIdentifier, int]
+
+
+class InlinePass(ir.passes.InPlacePass):
+    """Inline model local functions to the main graph and clear function definitions."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._functions: dict[ir.OperatorIdentifier, ir.Function] = {}
+        self._function_id_abbreviations: dict[ir.OperatorIdentifier, str] = {}
+        self._opset_imports: dict[str, int] = {}
+        self.used_value_names: set[str] = set()
+        self.used_node_names: set[str] = set()
+        self.node_context: dict[ir.Node, CallStack] = {}
+
+    def _reset(self, model: ir.Model) -> None:
+        self._functions = model.functions
+        self._function_id_abbreviations = _abbreviate(self._functions.keys())
+        self._opset_imports = model.opset_imports
+        self.used_value_names = set()
+        self.used_node_names = set()
+        self.node_context = {}
+
+    def call(self, model: ir.Model) -> InlinePassResult:
+        self._reset(model)
+        id_count = self._inline_calls_in(model.graph)
+        model.functions.clear()
+        return InlinePassResult(model, modified=bool(id_count), id_count=id_count)
+
+    def _instantiate_call(self, node: ir.Node, call_site_id: CallSiteId) -> NodeReplacement:
+        id = node.op_identifier()
+        function = self._functions[id]
+
+        # check opset compatibility and update the opset imports
+        for key, value in function.opset_imports.items():
+            if key not in self._opset_imports:
+                self._opset_imports[key] = value
+            elif self._opset_imports[key] != value:
+                raise ValueError(
+                    f"Opset mismatch: {key} {self._opset_imports[key]} != {value}"
+                )
+
+        # Identify substitutions for both inputs and attributes of the function:
+        attributes: Mapping[str, ir.Attr] = node.attributes
+        default_attr_values = {
+            attr.name: attr
+            for attr in function.attributes.values()
+            if attr.name not in attributes and attr.value is not None
+        }
+        if default_attr_values:
+            attributes = {**attributes, **default_attr_values}
+        if any(
+            attr.type in {ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS}
+            for attr in attributes.values()
+        ):
+            raise ValueError(
+                "Inliner does not support graph attribute parameters to functions"
+            )
+
+        if len(node.inputs) > len(function.inputs):
+            raise ValueError(f"Input mismatch: {len(node.inputs)} > {len(function.inputs)}")
+        value_map = {}
+        for i, input in enumerate(node.inputs):
+            value_map[function.inputs[i]] = input
+        for i in range(len(node.inputs), len(function.inputs)):
+            value_map[function.inputs[i]] = None
+
+        # Identify call-stack for node, used to generate unique names.
+        call_stack = self.node_context.get(node, [])
+        new_call_stack = [*call_stack, call_site_id]
+
+        cloner = _CopyReplace(self, attributes, value_map, node.metadata_props, new_call_stack)
+
+        # iterate over the nodes in the function, creating a copy of each node
+        # and replacing inputs with the corresponding values in the value map.
+        # Update the value map with the new values.
+
+        nodes = [cloner.clone_node(node) for node in function]
+        output_values = [value_map[output] for output in function.outputs]
+        return nodes, output_values  # type: ignore
+
+    def _inline_calls_in(self, graph: ir.Graph) -> dict[ir.OperatorIdentifier, int]:
+        for input in graph.inputs:
+            if input.name is not None:
+                self.used_value_names.add(input.name)
+        for initializer in graph.initializers:
+            self.used_value_names.add(initializer)
+
+        # Pre-processing:
+        # * Count the number of times each function is called in the graph.
+        #   This is used for disambiguating names of values in the inlined functions.
+        # * And identify names of values that are used in the graph.
+        id_count: dict[ir.OperatorIdentifier, int] = defaultdict(int)
+        for node in graph:
+            if node.name:
+                self.used_node_names.add(node.name)
+            id = node.op_identifier()
+            if id in self._functions:
+                id_count[id] += 1
+            for output in node.outputs:
+                if output.name is not None:
+                    self.used_value_names.add(output.name)
+        next_id: dict[ir.OperatorIdentifier, int] = defaultdict(int)
+        for node in graph:
+            id = node.op_identifier()
+            if id in self._functions:
+                # If there are multiple calls to same function, we use a prefix to disambiguate
+                # the different call-sites:
+                if id_count[id] > 1:
+                    call_site_prefix = f"_{next_id[id]}"
+                    next_id[id] += 1
+                else:
+                    call_site_prefix = ""
+                call_site = node.name or (
+                    self._function_id_abbreviations[id] + call_site_prefix
+                )
+                nodes, values = self._instantiate_call(node, call_site)
+                _ir_convenience.replace_nodes_and_values(
+                    graph,
+                    insertion_point=node,
+                    old_nodes=[node],
+                    new_nodes=nodes,
+                    old_values=node.outputs,
+                    new_values=values,
+                )
+            else:
+                for attr in node.attributes.values():
+                    if not isinstance(attr, ir.Attr):
+                        continue
+                    if attr.type == ir.AttributeType.GRAPH:
+                        self._inline_calls_in(attr.as_graph())
+                    elif attr.type == ir.AttributeType.GRAPHS:
+                        for g in attr.as_graphs():
+                            self._inline_calls_in(g)
+        return id_count

onnx_ir/passes/common/onnx_checker.py

@@ -0,0 +1,57 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Passes for debugging purposes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "CheckerPass",
+]
+
+from typing import Literal
+
+import onnx
+
+import onnx_ir as ir
+from onnx_ir.passes.common import _c_api_utils
+
+
+class CheckerPass(ir.passes.PassBase):
+    """Run onnx checker on the model."""
+
+    @property
+    def in_place(self) -> Literal[True]:
+        """This pass does not create a new model."""
+        return True
+
+    @property
+    def changes_input(self) -> Literal[False]:
+        """This pass does not change the input model."""
+        return False
+
+    def __init__(
+        self,
+        full_check: bool = False,
+        skip_opset_compatibility_check: bool = False,
+        check_custom_domain: bool = False,
+    ):
+        super().__init__()
+        self.full_check = full_check
+        self.skip_opset_compatibility_check = skip_opset_compatibility_check
+        self.check_custom_domain = check_custom_domain
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        """Run the onnx checker on the model."""
+
+        def _partial_check_model(proto: onnx.ModelProto) -> None:
+            """Partial function to check the model."""
+            onnx.checker.check_model(
+                proto,
+                full_check=self.full_check,
+                skip_opset_compatibility_check=self.skip_opset_compatibility_check,
+                check_custom_domain=self.check_custom_domain,
+            )
+
+        _c_api_utils.call_onnx_api(func=_partial_check_model, model=model)
+        # The model is not modified
+        return ir.passes.PassResult(model, False)