onnx-ir 0.1.15__py3-none-any.whl

onnx_ir/passes/common/initializer_deduplication.py

@@ -0,0 +1,179 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Pass for removing duplicated initializer tensors from a graph."""
+
+from __future__ import annotations
+
+__all__ = ["DeduplicateInitializersPass", "DeduplicateHashedInitializersPass"]
+
+
+import hashlib
+import logging
+
+import numpy as np
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def _should_skip_initializer(initializer: ir.Value, size_limit: int) -> bool:
+    """Check if the initializer should be skipped for deduplication."""
+    if initializer.is_graph_input() or initializer.is_graph_output():
+        # Skip graph inputs and outputs
+        logger.warning(
+            "Skipped deduplication of initializer '%s' as it is a graph input or output",
+            initializer.name,
+        )
+        return True
+
+    const_val = initializer.const_value
+    if const_val is None:
+        # Skip if initializer has no constant value
+        logger.warning(
+            "Skipped deduplication of initializer '%s' as it has no constant value. The model may contain invalid initializers",
+            initializer.name,
+        )
+        return True
+
+    if const_val.size > size_limit:
+        # Skip if the initializer is larger than the size limit
+        logger.debug(
+            "Skipped initializer '%s' as it exceeds the size limit of %d elements",
+            initializer.name,
+            size_limit,
+        )
+        return True
+    return False
+
+
+def _tobytes(val):
+    """StringTensor does not support tobytes. Use 'string_data' instead.
+
+    However, 'string_data' yields a list of bytes which cannot be hashed, i.e.,
+    cannot be used to index into a dict. To generate keys for identifying
+    tensors in initializer deduplication the following converts the list of
+    bytes to an array of fixed-length strings which can be flattened into a
+    bytes-string. This, together with the tensor shape, is sufficient for
+    identifying tensors for deduplication, but it differs from the
+    representation used for serializing tensors (that is string_data) by adding
+    padding bytes so that each string occupies the same number of consecutive
+    bytes in the flattened .tobytes representation.
+    """
+    if val.dtype.is_string():
+        return np.array(val.string_data()).tobytes()
+    return val.tobytes()
+
+
+class DeduplicateInitializersPass(ir.passes.InPlacePass):
+    """Remove duplicated initializer tensors from the main graph and all subgraphs.
+
+    This pass detects initializers with identical shape, dtype, and content,
+    and replaces all duplicate references with a canonical one.
+
+    Initializers are deduplicated within each graph. To deduplicate initializers
+    in the model globally (across graphs), use :class:`~onnx_ir.passes.common.LiftSubgraphInitializersToMainGraphPass`
+    to lift the initializers to the main graph first before running pass.
+
+    .. versionadded:: 0.1.3
+    .. versionchanged:: 0.1.7
+        This pass now deduplicates initializers in subgraphs as well.
+    """
+
+    def __init__(self, size_limit: int = 1024):
+        super().__init__()
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        modified = False
+
+        for graph in model.graphs():
+            initializers: dict[tuple[ir.DataType, tuple[int, ...], bytes], ir.Value] = {}
+            for initializer in tuple(graph.initializers.values()):
+                if _should_skip_initializer(initializer, self.size_limit):
+                    continue
+
+                const_val = initializer.const_value
+                assert const_val is not None
+
+                key = (const_val.dtype, tuple(const_val.shape), _tobytes(const_val))
+                if key in initializers:
+                    modified = True
+                    initializer_to_keep = initializers[key]  # type: ignore[index]
+                    initializer.replace_all_uses_with(initializer_to_keep)
+                    assert initializer.name is not None
+                    graph.initializers.pop(initializer.name)
+                    logger.info(
+                        "Replaced initializer '%s' with existing initializer '%s'",
+                        initializer.name,
+                        initializer_to_keep.name,
+                    )
+                else:
+                    initializers[key] = initializer  # type: ignore[index]
+
+        return ir.passes.PassResult(model=model, modified=modified)
+
+
+class DeduplicateHashedInitializersPass(ir.passes.InPlacePass):
+    """Remove duplicated initializer tensors (using a hashed method) from the graph.
+
+    This pass detects initializers with identical shape, dtype, and hashed content,
+    and replaces all duplicate references with a canonical one.
+
+    This pass should have a lower peak memory usage than :class:`DeduplicateInitializersPass`
+    as it does not store the full tensor data in memory, but instead uses a hash of the tensor data.
+
+    .. versionadded:: 0.1.7
+    """
+
+    def __init__(self, size_limit: int = 4 * 1024 * 1024 * 1024):
+        super().__init__()
+        # 4 GB default size limit for deduplication
+        self.size_limit = size_limit
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        modified = False
+
+        for graph in model.graphs():
+            initializers: dict[tuple[ir.DataType, tuple[int, ...], str], ir.Value] = {}
+
+            for initializer in tuple(graph.initializers.values()):
+                if _should_skip_initializer(initializer, self.size_limit):
+                    continue
+
+                const_val = initializer.const_value
+                assert const_val is not None
+
+                # Hash tensor data to avoid storing large amounts of data in memory
+                hashed = hashlib.sha512()
+                tensor_data = const_val.numpy()
+                hashed.update(tensor_data)
+                tensor_digest = hashed.hexdigest()
+
+                tensor_dims = tuple(const_val.shape.numpy())
+
+                key = (const_val.dtype, tensor_dims, tensor_digest)
+
+                if key in initializers:
+                    if _tobytes(initializers[key].const_value) != _tobytes(const_val):
+                        logger.warning(
+                            "Initializer deduplication failed: "
+                            "hashes match but values differ with values %s and %s",
+                            initializers[key],
+                            initializer,
+                        )
+                        continue
+                    modified = True
+                    initializer_to_keep = initializers[key]  # type: ignore[index]
+                    initializer.replace_all_uses_with(initializer_to_keep)
+                    assert initializer.name is not None
+                    graph.initializers.pop(initializer.name)
+                    logger.info(
+                        "Replaced initializer '%s' with existing initializer '%s'",
+                        initializer.name,
+                        initializer_to_keep.name,
+                    )
+                else:
+                    initializers[key] = initializer  # type: ignore[index]
+
+        return ir.passes.PassResult(model=model, modified=modified)

onnx_ir/passes/common/inliner.py

@@ -0,0 +1,223 @@
+# 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
+from onnx_ir import _cloner
+
+# 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
+
+
+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]
+
+        def rename(node: ir.Node) -> None:
+            """Rename node/values in inlined node to ensure uniqueness in the inlined context."""
+            node_name = node.name or "node"
+            node.name = _make_unique_name(node_name, new_call_stack, self.used_node_names)
+            for output in node.outputs:
+                if output is not None:
+                    output_name = output.name or "val"
+                    output.name = _make_unique_name(
+                        output_name, new_call_stack, self.used_value_names
+                    )
+            # Update context in case the new node is itself a call node that will be inlined.
+            self.node_context[node] = new_call_stack
+
+        cloner = _cloner.Cloner(
+            attr_map=attributes,
+            value_map=value_map,
+            metadata_props=node.metadata_props,
+            post_process=rename,
+            resolve_ref_attrs=True,
+        )
+
+        # 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 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/naming.py

@@ -0,0 +1,280 @@
+# Copyright (c) ONNX Project Contributors
+# SPDX-License-Identifier: Apache-2.0
+"""Name fix pass for ensuring unique names for all values and nodes."""
+
+from __future__ import annotations
+
+__all__ = [
+    "NameFixPass",
+    "NameGenerator",
+    "SimpleNameGenerator",
+]
+
+import collections
+import logging
+from typing import Protocol
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class NameGenerator(Protocol):
+    def generate_node_name(self, node: ir.Node) -> str:
+        """Generate a preferred name for a node."""
+        ...
+
+    def generate_value_name(self, value: ir.Value) -> str:
+        """Generate a preferred name for a value."""
+        ...
+
+
+class SimpleNameGenerator(NameGenerator):
+    """Base class for name generation functions."""
+
+    def generate_node_name(self, node: ir.Node) -> str:
+        """Generate a preferred name for a node."""
+        return node.name or "node"
+
+    def generate_value_name(self, value: ir.Value) -> str:
+        """Generate a preferred name for a value."""
+        return value.name or "v"
+
+
+class NameFixPass(ir.passes.InPlacePass):
+    """Pass for fixing names to ensure all values and nodes have unique names.
+
+    This pass ensures that:
+    1. Graph inputs and outputs have unique names (take precedence)
+    2. All intermediate values have unique names (assign names to unnamed values)
+    3. All values in subgraphs have unique names within their graph and parent graphs
+    4. All nodes have unique names within their graph
+
+    The pass maintains global uniqueness across the entire model.
+
+    You can customize the name generation functions for nodes and values by passing
+    a subclass of :class:`NameGenerator`.
+
+    For example, you can use a custom naming scheme like this::
+
+        class CustomNameGenerator:
+            def custom_node_name(node: ir.Node) -> str:
+                return f"custom_node_{node.op_type}"
+
+            def custom_value_name(value: ir.Value) -> str:
+                return f"custom_value_{value.type}"
+
+        name_fix_pass = NameFixPass(name_generator=CustomNameGenerator())
+
+    .. versionadded:: 0.1.6
+    """
+
+    def __init__(
+        self,
+        name_generator: NameGenerator | None = None,
+    ) -> None:
+        """Initialize the NameFixPass with custom name generation functions.
+
+        Args:
+            name_generator (NameGenerator, optional): An instance of a subclass of
+                :class:`NameGenerator` to customize name generation for nodes and values.
+                If not provided, defaults to a basic implementation that uses
+                the node's or value's existing name or a generic name like "node" or "v".
+        """
+        super().__init__()
+        self._name_generator = name_generator or SimpleNameGenerator()
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        # Process the main graph
+        modified = self._fix_graph_names(model.graph)
+
+        # Process functions
+        for function in model.functions.values():
+            modified = self._fix_graph_names(function) or modified
+
+        return ir.passes.PassResult(model, modified=modified)
+
+    def _fix_graph_names(self, graph_like: ir.Graph | ir.Function) -> bool:
+        """Fix names in a graph and return whether modifications were made."""
+        modified = False
+
+        # Set to track which values have been assigned names
+        seen_values: set[ir.Value] = set()
+
+        # The first set is a dummy placeholder so that there is always a [-1] scope for access
+        # (even though we don't write to it)
+        scoped_used_value_names: list[set[str]] = [set()]
+        scoped_used_node_names: list[set[str]] = [set()]
+
+        # Counters for generating unique names (using list to pass by reference)
+        value_counter: collections.Counter[str] = collections.Counter()
+        node_counter: collections.Counter[str] = collections.Counter()
+
+        def enter_graph(graph_like) -> None:
+            """Callback for entering a subgraph."""
+            # Initialize new scopes with all names from the parent scope
+            scoped_used_value_names.append(set(scoped_used_value_names[-1]))
+            scoped_used_node_names.append(set())
+
+            nonlocal modified
+
+            # Step 1: Fix graph input names first (they have precedence)
+            for input_value in graph_like.inputs:
+                if self._process_value(
+                    input_value, scoped_used_value_names[-1], seen_values, value_counter
+                ):
+                    modified = True
+
+            # Step 2: Fix graph output names (they have precedence)
+            for output_value in graph_like.outputs:
+                if self._process_value(
+                    output_value, scoped_used_value_names[-1], seen_values, value_counter
+                ):
+                    modified = True
+
+            if isinstance(graph_like, ir.Graph):
+                # For graphs, also fix initializers
+                for initializer in graph_like.initializers.values():
+                    if self._process_value(
+                        initializer, scoped_used_value_names[-1], seen_values, value_counter
+                    ):
+                        modified = True
+
+        def exit_graph(_) -> None:
+            """Callback for exiting a subgraph."""
+            # Pop the current scope
+            scoped_used_value_names.pop()
+            scoped_used_node_names.pop()
+
+        # Step 3: Process all nodes and their values
+        for node in ir.traversal.RecursiveGraphIterator(
+            graph_like, enter_graph=enter_graph, exit_graph=exit_graph
+        ):
+            # Fix node name
+            if not node.name:
+                if self._assign_node_name(node, scoped_used_node_names[-1], node_counter):
+                    modified = True
+            else:
+                if self._fix_duplicate_node_name(
+                    node, scoped_used_node_names[-1], node_counter
+                ):
+                    modified = True
+
+            # Fix input value names (only if not already processed)
+            for input_value in node.inputs:
+                if input_value is not None:
+                    if self._process_value(
+                        input_value, scoped_used_value_names[-1], seen_values, value_counter
+                    ):
+                        modified = True
+
+            # Fix output value names (only if not already processed)
+            for output_value in node.outputs:
+                if self._process_value(
+                    output_value, scoped_used_value_names[-1], seen_values, value_counter
+                ):
+                    modified = True
+
+        return modified
+
+    def _process_value(
+        self,
+        value: ir.Value,
+        used_value_names: set[str],
+        seen_values: set[ir.Value],
+        value_counter: collections.Counter[str],
+    ) -> bool:
+        """Process a value only if it hasn't been processed before."""
+        if value in seen_values:
+            return False
+
+        modified = False
+
+        if not value.name:
+            modified = self._assign_value_name(value, used_value_names, value_counter)
+        else:
+            modified = self._fix_duplicate_value_name(value, used_value_names, value_counter)
+            # initializers dictionary is updated automatically when the Value is renamed
+
+        # Record the final name for this value
+        assert value.name is not None
+        seen_values.add(value)
+        return modified
+
+    def _assign_value_name(
+        self, value: ir.Value, used_names: set[str], counter: collections.Counter[str]
+    ) -> bool:
+        """Assign a name to an unnamed value. Returns True if modified."""
+        assert not value.name, (
+            "value should not have a name already if function is called correctly"
+        )
+
+        preferred_name = self._name_generator.generate_value_name(value)
+        value.name = _find_and_record_next_unique_name(preferred_name, used_names, counter)
+        logger.debug("Assigned name %s to unnamed value", value.name)
+        return True
+
+    def _assign_node_name(
+        self, node: ir.Node, used_names: set[str], counter: collections.Counter[str]
+    ) -> bool:
+        """Assign a name to an unnamed node. Returns True if modified."""
+        assert not node.name, (
+            "node should not have a name already if function is called correctly"
+        )
+
+        preferred_name = self._name_generator.generate_node_name(node)
+        node.name = _find_and_record_next_unique_name(preferred_name, used_names, counter)
+        logger.debug("Assigned name %s to unnamed node", node.name)
+        return True
+
+    def _fix_duplicate_value_name(
+        self, value: ir.Value, used_names: set[str], counter: collections.Counter[str]
+    ) -> bool:
+        """Fix a value's name if it conflicts with existing names. Returns True if modified."""
+        original_name = value.name
+
+        assert original_name, (
+            "value should have a name already if function is called correctly"
+        )
+
+        if original_name not in used_names:
+            # Name is unique, just record it
+            used_names.add(original_name)
+            return False
+
+        # If name is already used, make it unique
+        base_name = self._name_generator.generate_value_name(value)
+        value.name = _find_and_record_next_unique_name(base_name, used_names, counter)
+        logger.debug("Renamed value from %s to %s for uniqueness", original_name, value.name)
+        return True
+
+    def _fix_duplicate_node_name(
+        self, node: ir.Node, used_names: set[str], counter: collections.Counter[str]
+    ) -> bool:
+        """Fix a node's name if it conflicts with existing names. Returns True if modified."""
+        original_name = node.name
+
+        assert original_name, "node should have a name already if function is called correctly"
+
+        if original_name not in used_names:
+            # Name is unique, just record it
+            used_names.add(original_name)
+            return False
+
+        # If name is already used, make it unique
+        base_name = self._name_generator.generate_node_name(node)
+        node.name = _find_and_record_next_unique_name(base_name, used_names, counter)
+        logger.debug("Renamed node from %s to %s for uniqueness", original_name, node.name)
+        return True
+
+
+def _find_and_record_next_unique_name(
+    preferred_name: str, used_names: set[str], counter: collections.Counter[str]
+) -> str:
+    """Generate a unique name based on the preferred name and current counter."""
+    new_name = preferred_name
+    while new_name in used_names:
+        counter[preferred_name] += 1
+        new_name = f"{preferred_name}_{counter[preferred_name]}"
+    used_names.add(new_name)
+    return new_name