onnx-ir 0.0.1__py3-none-any.whl

onnx_ir/_io.py

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Load and save ONNX models."""
+
+from __future__ import annotations
+
+__all__ = ["load", "save"]
+
+import os
+
+import onnx
+
+from onnx_ir import _core, _external_data, serde
+
+
+def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
+    """Load an ONNX model from a file.
+
+    Args:
+        path: The path to the ONNX file.
+        format: The format of the file (e.g. protobuf, textproto, json, etc.).
+            If None, the format is inferred from the file extension.
+
+    Returns:
+        The loaded model.
+    """
+    # Do not use ONNX to load external data because the IR handles external data
+    # by doing memory mapping directly.
+    proto = onnx.load(path, format=format, load_external_data=False)
+    model = serde.deserialize_model(proto)
+    base_dir = os.path.dirname(path)
+    # Set the base directory for external data to the directory of the ONNX file
+    # so that relative paths are resolved correctly.
+    _external_data.set_base_dir(model.graph, base_dir)
+    return model
+
+
+def save(model: _core.Model, path: str | os.PathLike, format: str | None = None) -> None:
+    """Save an ONNX model to a file.
+
+    Args:
+        model: The model to save.
+        path: The path to save the model to.
+        format: The format of the file (e.g. protobuf, textproto, json, etc.).
+            If None, the format is inferred from the file extension.
+    """
+    proto = serde.serialize_model(model)
+    onnx.save(proto, path, format=format)
+    # TODO(justinchuby): Handle external data when the relative path has changed
+    # TODO(justinchuby): Handle off loading external data to disk when saving

onnx_ir/_linked_list.py

@@ -0,0 +1,276 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Mutable list for nodes in a graph with safe mutation properties."""
+
+from __future__ import annotations
+
+from typing import Generic, Iterable, Iterator, Sequence, TypeVar
+
+T = TypeVar("T")
+
+
+class _LinkBox(Generic[T]):
+    """A link in a doubly linked list that has a reference to the actual object in the link.
+
+    The :class:`_LinkBox` is a container for the actual object in the list. It is used to
+    maintain the links between the elements in the linked list. The actual object is stored in the
+    :attr:`value` attribute.
+
+    By using a separate container for the actual object, we can safely remove the object from the
+    list without losing the links. This allows us to remove the object from the list during
+    iteration and place the object into a different list without breaking any chains.
+
+    This is an internal class and should only be initialized by the :class:`DoublyLinkedSet`.
+
+    Attributes:
+        prev: The previous box in the list.
+        next: The next box in the list.
+        erased: A flag to indicate if the box has been removed from the list.
+        owning_list: The :class:`DoublyLinkedSet` to which the box belongs.
+        value: The actual object in the list.
+    """
+
+    __slots__ = ("next", "owning_list", "prev", "value")
+
+    def __init__(self, owner: DoublyLinkedSet[T], value: T | None) -> None:
+        """Create a new link box.
+
+        Args:
+            owner: The linked list to which this box belongs.
+            value: The value to be stored in the link box. When the value is None,
+                the link box is considered erased (default). The root box of the list
+                should be created with a None value.
+        """
+        self.prev: _LinkBox[T] = self
+        self.next: _LinkBox[T] = self
+        self.value: T | None = value
+        self.owning_list: DoublyLinkedSet[T] = owner
+
+    @property
+    def erased(self) -> bool:
+        return self.value is None
+
+    def erase(self) -> None:
+        """Remove the link from the list and detach the value from the box."""
+        if self.value is None:
+            raise ValueError("_LinkBox is already erased")
+        # Update the links
+        prev, next_ = self.prev, self.next
+        prev.next, next_.prev = next_, prev
+        # Detach the value
+        self.value = None
+
+    def __repr__(self) -> str:
+        return f"_LinkBox({self.value!r}, erased={self.erased}, prev={self.prev.value!r}, next={self.next.value!r})"
+
+
+class DoublyLinkedSet(Sequence[T], Generic[T]):
+    """A doubly linked ordered set of nodes.
+
+    The container can be viewed as a set as it does not allow duplicate values. The order of the
+    elements is maintained. One can typically treat it as a doubly linked list with list-like
+    methods implemented.
+
+    Adding and removing elements from the set during iteration is safe. Moving elements
+    from one set to another is also safe.
+
+    During the iteration:
+    - If new elements are inserted after the current node, the iterator will
+        iterate over them as well.
+    - If new elements are inserted before the current node, they will
+        not be iterated over in this iteration.
+    - If the current node is lifted and inserted in a different location,
+        iteration will start from the "next" node at the _original_ location.
+
+    Time complexity:
+        Inserting and removing nodes from the set is O(1). Accessing nodes by index is O(n),
+        although accessing nodes at either end of the set is O(1). I.e.
+        ``linked_set[0]`` and ``linked_set[-1]`` are O(1).
+
+    Values need to be hashable. ``None`` is not a valid value in the set.
+    """
+
+    __slots__ = ("_length", "_root", "_value_ids_to_boxes")
+
+    def __init__(self, values: Iterable[T] | None = None) -> None:
+        # Using the root node simplifies the mutation implementation a lot
+        # The list is circular. The root node is the only node that is not a part of the list values
+        root_ = _LinkBox(self, None)
+        self._root: _LinkBox = root_
+        self._length = 0
+        self._value_ids_to_boxes: dict[int, _LinkBox] = {}
+        if values is not None:
+            self.extend(values)
+
+    def __iter__(self) -> Iterator[T]:
+        """Iterate over the elements in the list.
+
+        - If new elements are inserted after the current node, the iterator will
+            iterate over them as well.
+        - If new elements are inserted before the current node, they will
+            not be iterated over in this iteration.
+        - If the current node is lifted and inserted in a different location,
+            iteration will start from the "next" node at the _original_ location.
+        """
+        box = self._root.next
+        while box is not self._root:
+            if box.owning_list is not self:
+                raise RuntimeError(f"Element {box!r} is not in the list")
+            if not box.erased:
+                assert box.value is not None
+                yield box.value
+            box = box.next
+
+    def __reversed__(self) -> Iterator[T]:
+        """Iterate over the elements in the list in reverse order."""
+        box = self._root.prev
+        while box is not self._root:
+            if not box.erased:
+                assert box.value is not None
+                yield box.value
+            box = box.prev
+
+    def __len__(self) -> int:
+        assert self._length == len(
+            self._value_ids_to_boxes
+        ), "Bug in the implementation: length mismatch"
+        return self._length
+
+    def __getitem__(self, index: int) -> T:
+        """Get the node at the given index.
+
+        Complexity is O(n).
+        """
+        if index >= self._length or index < -self._length:
+            raise IndexError(
+                f"Index out of range: {index} not in range [-{self._length}, {self._length})"
+            )
+        if index < 0:
+            # Look up from the end of the list
+            iterator = reversed(self)
+            item = next(iterator)
+            for _ in range(-index - 1):
+                item = next(iterator)
+        else:
+            iterator = iter(self)  # type: ignore[assignment]
+            item = next(iterator)
+            for _ in range(index):
+                item = next(iterator)
+        return item
+
+    def _insert_one_after(
+        self,
+        box: _LinkBox[T],
+        new_value: T,
+    ) -> _LinkBox[T]:
+        """Insert a new value after the given box.
+
+        All insertion methods should call this method to ensure that the list is updated correctly.
+
+        Example::
+            Before: A  <->  B  <->  C
+                    ^v0     ^v1     ^v2
+            Call: _insert_one_after(B, v3)
+            After:  A  <->  B  <->  new_box  <->  C
+                    ^v0     ^v1       ^v3         ^v2
+
+        Args:
+            box: The box which the new value is to be inserted.
+            new_value: The new value to be inserted.
+        """
+        if new_value is None:
+            raise TypeError(f"{self.__class__.__name__} does not support None values")
+        if box.value is new_value:
+            # Do nothing if the new value is the same as the old value
+            return box
+        if box.owning_list is not self:
+            raise ValueError(f"Value {box.value!r} is not in the list")
+
+        if (new_value_id := id(new_value)) in self._value_ids_to_boxes:
+            # If the value is already in the list, remove it first
+            self.remove(new_value)
+
+        # Create a new _LinkBox for the new value
+        new_box = _LinkBox(self, new_value)
+        # original_box <=> original_next
+        # becomes
+        # original_box <=> new_box <=> original_next
+        original_next = box.next
+        box.next = new_box
+        new_box.prev = box
+        new_box.next = original_next
+        original_next.prev = new_box
+
+        # Be sure to update the length and mapping
+        self._length += 1
+        self._value_ids_to_boxes[new_value_id] = new_box
+
+        return new_box
+
+    def _insert_many_after(
+        self,
+        box: _LinkBox[T],
+        new_values: Iterable[T],
+    ):
+        """Insert multiple new values after the given box."""
+        insertion_point = box
+        for new_value in new_values:
+            insertion_point = self._insert_one_after(insertion_point, new_value)
+
+    def remove(self, value: T) -> None:
+        """Remove a node from the list."""
+        if (value_id := id(value)) not in self._value_ids_to_boxes:
+            raise ValueError(f"Value {value!r} is not in the list")
+        box = self._value_ids_to_boxes[value_id]
+        # Remove the link box and detach the value from the box
+        box.erase()
+
+        # Be sure to update the length and mapping
+        self._length -= 1
+        del self._value_ids_to_boxes[value_id]
+
+    def append(self, value: T) -> None:
+        """Append a node to the list."""
+        _ = self._insert_one_after(self._root.prev, value)
+
+    def extend(
+        self,
+        values: Iterable[T],
+    ) -> None:
+        for value in values:
+            self.append(value)
+
+    def insert_after(
+        self,
+        value: T,
+        new_values: Iterable[T],
+    ) -> None:
+        """Insert new nodes after the given node.
+
+        Args:
+            value: The value after which the new values are to be inserted.
+            new_values: The new values to be inserted.
+        """
+        if (value_id := id(value)) not in self._value_ids_to_boxes:
+            raise ValueError(f"Value {value!r} is not in the list")
+        insertion_point = self._value_ids_to_boxes[value_id]
+        return self._insert_many_after(insertion_point, new_values)
+
+    def insert_before(
+        self,
+        value: T,
+        new_values: Iterable[T],
+    ) -> None:
+        """Insert new nodes before the given node.
+
+        Args:
+            value: The value before which the new values are to be inserted.
+            new_values: The new values to be inserted.
+        """
+        if (value_id := id(value)) not in self._value_ids_to_boxes:
+            raise ValueError(f"Value {value!r} is not in the list")
+        insertion_point = self._value_ids_to_boxes[value_id].prev
+        return self._insert_many_after(insertion_point, new_values)
+
+    def __repr__(self) -> str:
+        return f"DoublyLinkedSet({list(self)})"

onnx_ir/_metadata.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Class for storing metadata about the IR objects."""
+
+from __future__ import annotations
+
+import collections
+from typing import Any, Mapping
+
+
+class MetadataStore(collections.UserDict):
+    """Class for storing metadata about the IR objects.
+
+    Metadata is stored as key-value pairs. The keys are strings and the values
+    can be any Python object.
+
+    The metadata store also supports marking keys as invalid. This is useful
+    when a pass wants to mark a key that needs to be recomputed.
+    """
+
+    def __init__(self, data: Mapping[str, Any] | None = None, /) -> None:
+        super().__init__(data)
+        self._invalid_keys: set[str] = set()
+
+    def __setitem__(self, key: str, item: Any) -> None:
+        self.data[key] = item
+        self._invalid_keys.discard(key)
+
+    def invalidate(self, key: str) -> None:
+        self._invalid_keys.add(key)
+
+    def is_valid(self, key: str) -> bool:
+        """Returns whether the value is valid.
+
+        Note that default values (None) are not necessarily invalid. For example,
+        a shape that is unknown (None) may be still valid if shape inference has
+        determined that the shape is unknown.
+
+        Whether a value is valid is solely determined by the user that sets the value.
+        """
+        return key not in self._invalid_keys
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.data!r}, invalid_keys={self._invalid_keys!r})"

onnx_ir/_name_authority.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Auxiliary class for managing names in the IR."""
+
+from __future__ import annotations
+
+from onnx_ir import _core
+
+
+class NameAuthority:
+    """Class for giving names to values and nodes in the IR.
+
+    The names are generated in the format ``val_{value_counter}`` for values and
+    ``node_{op_type}_{node_counter}`` for nodes. The counter is incremented each time
+    a new value or node is named.
+
+    This class keeps tracks of the names it has generated and existing names
+    in the graph to prevent producing duplicated names.
+
+    .. note::
+        Once a name is tracked, it will not be made available even if the node/value
+        is removed from the graph. It is possible to improve this behavior by keeping
+        track of the names that are no longer used, but it is not implemented yet.
+
+    However, if a value/node is already named when added to the graph,
+    the name authority will not change its name.
+    It is the responsibility of the user to ensure that the names are unique
+    (typically by running a name-fixing pass on the graph).
+
+    TODO(justichuby): Describe the pass when we have a reference implementation.
+    """
+
+    def __init__(self):
+        self._value_counter = 0
+        self._node_counter = 0
+        self._value_names: set[str] = set()
+        self._node_names: set[str] = set()
+
+    def _unique_value_name(self) -> str:
+        """Generate a unique name for a value."""
+        while True:
+            name = f"val_{self._value_counter}"
+            self._value_counter += 1
+            if name not in self._value_names:
+                return name
+
+    def _unique_node_name(self, op_type: str) -> str:
+        """Generate a unique name for a node."""
+        while True:
+            name = f"node_{op_type}_{self._node_counter}"
+            self._node_counter += 1
+            if name not in self._node_names:
+                return name
+
+    def register_or_name_value(self, value: _core.Value) -> None:
+        # TODO(justinchuby): Record names of the initializers and graph inputs
+        if value.name is None:
+            value.name = self._unique_value_name()
+        # If the name is already specified, we do not change it because keeping
+        # track of the used names can be costly when nodes can be removed from the graph:
+        # How do we know if a name is no longer used? We cannot reserve unused names
+        # because users may want to use them.
+        self._value_names.add(value.name)
+
+    def register_or_name_node(self, node: _core.Node) -> None:
+        if node.name is None:
+            node.name = self._unique_node_name(node.op_type)
+        # If the name is already specified, we do not change it because keeping
+        # track of the used names can be costly when nodes can be removed from the graph:
+        # How do we know if a name is no longer used? We cannot reserve unused names
+        # because users may want to use them.
+        self._node_names.add(node.name)