GraphIO

justinchuby · justinchuby · commit da95c66aa858 · 2025-05-07T15:42:08.000-07:00
diff --git a/onnxscript/ir/_core.py b/onnxscript/ir/_core.py
@@ -1757,6 +1757,9 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
 
     __slots__ = (
         "_const_value",
+        "_graph_initializer_of",
+        "_graph_input_of",
+        "_graph_output_of",
         "_index",
         "_metadata",
         "_metadata_props",
@@ -1808,6 +1811,13 @@ def __init__(
         self._uses: dict[Usage, None] = {}
         self.doc_string = doc_string
 
+        # The graph this value belongs to. It is set *only* when the value is added as
+        # a graph input, graph output, or initializer.
+        # The three properties can only be set by the Graph class (GraphIO).
+        self._graph_initializer_of: Graph | None = None
+        self._graph_input_of: Graph | None = None
+        self._graph_output_of: Graph | None = None
+
     def __repr__(self) -> str:
         value_name = self.name if self.name else "anonymous:" + str(id(self))
         type_text = f", type={self.type!r}" if self.type is not None else ""
@@ -1986,15 +1996,17 @@ def metadata_props(self) -> dict[str, str]:
             self._metadata_props = {}
         return self._metadata_props
 
+    def is_initializer(self) -> bool:
+        """Whether the value is an initializer."""
+        return self._graph_initializer_of is not None
+
+    def is_graph_input(self) -> bool:
+        """Whether the value is an input of a graph."""
+        return self._graph_input_of is not None
+
     def is_graph_output(self) -> bool:
         """Whether the value is an output of a graph."""
-        if (producer := self.producer()) is None:
-            return False
-        if (graph := producer.graph) is None:
-            return False
-        # Cannot use `in` because __eq__ may be defined by subclasses, even though
-        # it is not recommended
-        return any(output is self for output in graph.outputs)
+        return self._graph_output_of is not None
 
 
 def Input(
diff --git a/onnxscript/ir/_tracked_lists.py b/onnxscript/ir/_tracked_lists.py
@@ -5,185 +5,92 @@
 from __future__ import annotations
 
 import collections
-import contextlib
-from typing import TYPE_CHECKING, NoReturn
+from typing import TYPE_CHECKING, Iterable, Literal, SupportsIndex
 
 if TYPE_CHECKING:
     from onnxscript.ir import _core
 
 
-@contextlib.contextmanager
-def _update_value_usages(node: _core.Node, inputs: list[_core.Value | None]):
-    """Temporarily unset usages of the inputs."""
-    for i, item in enumerate(inputs):
-        if item is not None:
-            item._remove_usage(node, i)  # pylint: disable=protected-access
-    try:
-        # Caller will modify the inputs
-        yield
-    finally:
-        for i, item in enumerate(inputs):
-            if item is not None:
-                item._add_usage(node, i)  # pylint: disable=protected-access
+class GraphIO(collections.UserList[_core.Value]):
+    """The inputs and outputs of a Graph."""
 
-
-class NodeInputs(collections.UserList[_core.Value | None]):
-    def __init__(self, node: _core.Node, initlist=None):
+    def __init__(self, graph: _core.Graph, typ: Literal["input", "output"], initlist=None):
         super().__init__(initlist)
-        self._node = node
-
-    def prepend(self, item: _core.Value | None) -> None:
-        """Add a new input to the node."""
-        self.insert(0, item)
-
-    def append(self, item: _core.Value | None) -> None:
-        """Add a new input to the node."""
-        index = len(self.data)
-        if item is not None:
-            item._add_usage(self._node, index)  # pylint: disable=protected-access
-        self.data.append(item)
-
-    def extend(self, other) -> None:
-        for item in other:
-            self.append(item)
-
-    def insert(self, i: int, item: _core.Value | None) -> None:
-        with _update_value_usages(self._node, self.data):
-            self.data.insert(i, item)
+        self._graph = graph
+        assert typ in {"intput", "output"}
+        self._typ = typ
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph_input_of is not None and value._graph_input_of is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an input of a different graph: {value._graph_input_of!r}"
+            )
+        if value._graph_output_of is not None and value._graph_output_of is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an output of a different graph: {value._graph_output_of!r}"
+            )
 
-    def pop(self, i: int = -1) -> _core.Value | None:
-        item = self.data[i]
-        if i == -1:
-            # Remove the last item. No usages need to be updated
-            if item is not None:
-                item._remove_usage(self._node, i)
-            return self.data.pop()
-        # Otherwise we need to update usages
-        with _update_value_usages(self._node, self.data):
-            result = self.data.pop(i)
+        if self._typ == "input":
+            value._graph_input_of = self._graph
+        else:
+            value._graph_output_of = self._graph
 
-        return result
-
-    def clear(self) -> None:
-        """Clear the list."""
-        for i, item in enumerate(self.data):
-            if item is not None:
-                item._remove_usage(self._node, i)
-        self.data.clear()
-
-    def __setitem__(self, i: int, item: _core.Value | None) -> None:
-        """Replace an input to the node."""
-        if i < -len(self.data) or i >= len(self.data):
-            raise ValueError(f"index out of range: {i}")
-        if i < 0:
-            i += len(self.data)
-        assert i >= 0
-        old_input = self.data[i]
-        if old_input is not None:
-            old_input._remove_usage(self._node, i)  # pylint: disable=protected-access
-        if item is not None:
-            item._add_usage(self._node, i)  # pylint: disable=protected-access
-        self.data[i] = item
-
-    def unsupported(self, *_args, **_kwargs):
-        raise NotImplementedError("Method is not supported")
-
-    __lt__ = unsupported
-    __le__ = unsupported
-    __gt__ = unsupported
-    __ge__ = unsupported
-    __add__ = unsupported
-    __radd__ = unsupported
-    __iadd__ = unsupported
-    __mul__ = unsupported
-    reverse = unsupported
-    sort = unsupported
-
-
-class NodeOutputs(collections.UserList[_core.Value]):
-    def __init__(self, node: _core.Node, initlist=None):
-        super().__init__(initlist)
-        self._node = node
+    def _unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        if self._typ == "input":
+            value._graph_input_of = None
+        else:
+            value._graph_output_of = None
 
     def append(self, item: _core.Value) -> None:
-        """Add a new output to the node."""
-        if item.producer() is not None and item.producer() is not self._node:
-            raise NotImplementedError(
-                f"Output already has a producer that is not this node ({item.producer()}). "
-                "An output value can be owned by only one node throughout its lifetime. "
-                "Instead, create a new value and assign it to the output. Replace all usages of the old value with the new one."
-            )
-        item._producer = self._node  # pylint: disable=protected-access
-        item._index = len(self.data)  # pylint: disable=protected-access
-        self.data.append(item)
+        """Add a new input to the graph."""
+        super().append(item)
+        self._set_graph(item)
 
     def extend(self, other) -> None:
+        """Extend the list of inputs or outputs."""
+        super().extend(other)
         for item in other:
-            self.append(item)
+            self._set_graph(item)
+
+    def insert(self, i: int, item: _core.Value) -> None:
+        """Insert an input/output to the graph."""
+        super().insert(i, item)
+        self._set_graph(item)
 
-    def pop(self, i: int = -1) -> _core.Value | None:
-        item = self.data[i]
-        if i == -1:
-            # Remove the last item. No usages need to be updated
-            if item.uses():
-                raise ValueError(
-                    f"Cannot remove output {item} because it is still used by other nodes."
-                )
-            return self.data.pop()
+    def pop(self, i: int = -1) -> _core.Value:
+        """Remove an input/output from the graph."""
+        value = super().pop(i)
+        self._unset_graph(value)
+        return value
 
-        for j, output in enumerate(self.data):
-            output._index = j  # pylint: disable=protected-access
+    def remove(self, item: _core.Value) -> None:
+        """Remove an input/output from the graph."""
+        super().remove(item)
+        self._unset_graph(item)
 
     def clear(self) -> None:
         """Clear the list."""
-        for _ in range(len(self.data)):
-            self.pop()
-
-    def __setitem__(self, i: int, item: _core.Value) -> None:
-        if item is None:
-            raise NotImplementedError(
-                "An output cannot be None. To remove a trailing output, use pop(). "
-                "To remove an output in the middle, set its name to an empty string instead."
-            )
-        self.pop(i)
-        self.insert(i, item)
-
-    def insert(self, i: int, item: _core.Value) -> None:
-        """Replace an output to the node."""
-        if i < -len(self.data) or i >= len(self.data):
-            raise ValueError(f"index out of range: {i}")
-        if i < 0:
-            i += len(self.data)
-        assert i >= 0
-        if item.producer() is not None and item.producer() is not self._node:
-            raise NotImplementedError(
-                f"Output already has a producer that is not this node ({item.producer()}). "
-                "An output value can be owned by only one node throughout its lifetime. "
-                "Instead, create a new value and assign it to the output. Replace all usages of the old value with the new one."
-            )
-
-        item._producer = self._node  # pylint: disable=protected-access
-
-        # Update the index of the item being replaced
-        self.data.insert(i, item)
-        for j, output in enumerate(self.data):
-            output._index = j  # pylint: disable=protected-access
-
-
-    def unsupported(self, *_args, **_kwargs):
-        raise NotImplementedError("Method is not supported")
-
-    __lt__ = unsupported
-    __le__ = unsupported
-    __gt__ = unsupported
-    __ge__ = unsupported
-    __add__ = unsupported
-    __radd__ = unsupported
-    __iadd__ = unsupported
-    __mul__ = unsupported
-    # NOTE: We don't support insertion and removal of items in the middle of the list
-    # because in ONNX outputs are positional and index dependent. To remove an output in
-    # the middle, set its name to an empty empty string instead.
-    reverse = unsupported
-    sort = unsupported
-    remove = unsupported
+        for value in self.data:
+            self._unset_graph(value)
+        super().clear()
+
+    def __setitem__(self, i, item) -> None:
+        """Replace an input/output to the node."""
+        if isinstance(item, Iterable) and isinstance(i, slice):
+            # Modify a slice of the list
+            for value in self.data[i]:
+                self._unset_graph(value)
+            for value in item:
+                self._set_graph(value)
+            super().__setitem__(i, item)
+            return
+        elif isinstance(item, _core.Value) and isinstance(i, SupportsIndex):
+            # Replace a single item
+            self._unset_graph(self.data[i])
+            self._set_graph(item)
+            super().__setitem__(i, item)
+            return
+
+        raise TypeError(f"Invalid types for __setitem__: {type(i)} and {type(item)}")
