[IR] Mutable IR implementation #1344

justinchuby · 2024-04-04T04:50:03Z

Implement mutation methods on Graph and Function and a DoublyLinkedList to support safe mutation during iteration. Nodes in different graphs can be moved to other graphs safely during iteration.

Methods implemented

- __getitem__
- __len__
- __iter__
- __reversed__
- append
- extend
- remove
- insert_after
- insert_before
- sort (declared, not implemented)

The mutation methods are inspired by the pytorch FX graph methods.

Safe Iterators

The behavior of the iterators is:

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.

A node cannot be added to a graph if it belongs to another graph. It needs to be removed first. The user is responsible for removing nodes from other graphs before adding the same node to a new graph.

Linked list implementation

The doubly linked list implementation is inspired by the pytorch FX graph: They are both doubly linked lists with a dummy root node. Whereas in FX graph contains the root node and Nodes are the links; we create the DoublyLinkedOrderedSet class to decouple the list implementation form the graph. Additionally, we created the LinkedBox class as a data struct to hold the prev/next pointers for the Nodes to allow nodes to move across graphs. This is in contrast to FX restriction of all nodes must belong to the same graph in their lifecycle. By allowing the nodes to move to different graphs, we are able to flatten/unflatten graphs without copying.

DoublyLinkedOrderedSet is unit tested.

Next steps

Naming authority
Graph input/output/initializers mutation

onnxscript/ir/_core.py

onnxscript/ir/_protocols.py

+        Args:
+            node: The node to put before this node.
+        """
+        ...


onnxscript/ir/_core.py

codecov · 2024-04-04T04:55:20Z

Codecov Report

Attention: Patch coverage is 79.70085% with 95 lines in your changes are missing coverage. Please review.

Project coverage is 65.32%. Comparing base (edb68ed) to head (ae4838c).

Files	Patch %	Lines
onnxscript/ir/_core.py	32.55%	56 Missing and 2 partials ⚠️
onnxscript/ir/_protocols.py	48.64%	11 Missing and 8 partials ⚠️
onnxscript/ir/_linked_list.py	85.32%	8 Missing and 8 partials ⚠️
onnxscript/ir/_linked_list_test.py	99.15%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1344      +/-   ##
==========================================
+ Coverage   64.91%   65.32%   +0.40%     
==========================================
  Files         140      142       +2     
  Lines       19974    20414     +440     
  Branches     3349     3464     +115     
==========================================
+ Hits        12967    13336     +369     
- Misses       6469     6521      +52     
- Partials      538      557      +19

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

github-advanced-security

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.

github-actions · 2024-04-04T05:07:03Z

Test Results

23 files ± 0 23 suites ±0 19m 26s ⏱️ + 2m 25s
8 183 tests + 39 6 353 ✅ + 39 1 829 💤 ±0 1 ❌ ±0
21 356 runs +429 17 371 ✅ +429 3 984 💤 ±0 1 ❌ ±0

For more details on these failures, see this check.

Results for commit ae4838c. ± Comparison against base commit edb68ed.

♻️ This comment has been updated with latest results.

onnxscript/ir/_linked_list.py

onnxscript/ir/_core.py

onnxscript/ir/_linked_list.py

onnxscript/ir/_protocols.py

+
+    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol]) -> None:
+        """Insert new nodes before the given node."""
+        ...


onnxscript/ir/_linked_list.py

justinchuby · 2024-04-05T05:13:11Z

Modification in double iteration on the same graph is safe, but moving nodes from one to another is not. How should we handle/guard against it? Use boxes?

onnxscript/ir/_linked_list.py

onnxscript/ir/_core.py

onnxscript/ir/_linked_list_test.py

onnxscript/ir/_linked_list.py

justinchuby · 2024-04-06T14:39:02Z

I would like to bring #1349 into this PR. Created as a separate PR for review.

onnxscript/ir/_core.py

titaiwangms · 2024-04-08T16:50:48Z

onnxscript/ir/_core.py

+            node, new_nodes, property_modifier=self._add_graph_reference_to_node
+        )
+
+    def insert_before(self, node: Node, new_nodes: Iterable[Node]) -> None:


I see torch using

with graph.inserting_after(node): ...

Cause they have the need of creating nodes on the fly. I guess we don't have this need? Other than that, I think the imple here is simpler.

onnxscript/ir/_core.py

justinchuby · 2024-04-08T18:11:53Z

I decided to remove the read-only protocols. I realized they are overly complicated as many object fields refer to them and it is very hard to stay consistent.

onnxscript/ir/_protocols.py

+    def __getitem__(self, index: int) -> NodeProtocol: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[NodeProtocol]: ...
+    def __reversed__(self) -> Iterator[NodeProtocol]: ...


onnxscript/ir/_protocols.py


+    def __getitem__(self, index: int) -> NodeProtocol: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[NodeProtocol]: ...


onnxscript/ir/_protocols.py

+    def __getitem__(self, index: int) -> NodeProtocol: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[NodeProtocol]: ...
+    def __reversed__(self) -> Iterator[NodeProtocol]: ...


onnxscript/ir/_core.py

Use a dictionary to map the values to the boxes in the doubly linked list implementation to 1. Remove the private pointer from Node to LinkedBox 2. Simplify membership checks 3. Ensure elements in the list are unique Also changed: 1. The node <-> graph relationships are now managed by Graph directly and not by a property_modifier function. Meaning the user is responsible to remove nodes from other graphs before adding the same node to a new graph. This helps separate the concern of the graph from the linked list, allowing us to treat the linked list as a normal list. The list also doesn't need to handle exclusivity, meaning the same value will not be removed from the original list when being added to a new list and can appear in multiple lists, which is what we expect how lists behave.

onnxscript/ir/_core.py

onnxscript/ir/_protocols.py

justinchuby · 2024-04-09T01:31:45Z

Thanks all for reviewing. All comments are addressed. I will create follow-ups if there are more comments.

onnxscript/ir/_protocols.py

+    # Mutation methods
+    def append(self, node: NodeProtocol, /) -> None:
+        """Append a node to the graph."""
+        ...


onnxscript/ir/_protocols.py

+
+    def extend(self, nodes: Iterable[NodeProtocol], /) -> None:
+        """Extend the graph with the given nodes."""
+        ...


onnxscript/ir/_protocols.py

+
+    def remove(self, node: NodeProtocol, /) -> None:
+        """Remove a node from the graph."""
+        ...


onnxscript/ir/_protocols.py

+
+    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+        """Insert new nodes after the given node."""
+        ...


onnxscript/ir/_protocols.py

+    # End Block
+    def append(self, node: NodeProtocol, /) -> None:
+        """Append a node to the function."""
+        ...


onnxscript/ir/_protocols.py

+
+    def extend(self, nodes: Iterable[NodeProtocol], /) -> None:
+        """Extend the function with the given nodes."""
+        ...


onnxscript/ir/_protocols.py

+
+    def remove(self, node: NodeProtocol, /) -> None:
+        """Remove a node from the function."""
+        ...


onnxscript/ir/_protocols.py

+
+    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+        """Insert new nodes after the given node."""
+        ...


justinchuby added the topic: IR Intermediate representation label Apr 4, 2024

justinchuby commented Apr 4, 2024

View reviewed changes

onnxscript/ir/_core.py Show resolved Hide resolved

justinchuby commented Apr 4, 2024

View reviewed changes

onnxscript/ir/_core.py Outdated Show resolved Hide resolved

justinchuby commented Apr 4, 2024

View reviewed changes

onnxscript/ir/_core.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 4, 2024

View reviewed changes

justinchuby commented Apr 4, 2024

View reviewed changes

onnxscript/ir/_core.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 4, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Fixed Show fixed Hide fixed

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

BowenBao reviewed Apr 4, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Fixed Show fixed Hide fixed

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Fixed Show fixed Hide fixed

justinchuby mentioned this pull request Apr 5, 2024

IR Roadmap #1324

Open

67 tasks

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Fixed Show fixed Hide fixed

justinchuby changed the title ~~[IR] Mutable implementation~~ [IR] Mutable IR implementation Apr 5, 2024

justinchuby commented Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Show resolved Hide resolved

justinchuby commented Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Outdated Show resolved Hide resolved

justinchuby commented Apr 5, 2024

View reviewed changes

onnxscript/ir/_core.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

github-advanced-security bot found potential problems Apr 5, 2024

View reviewed changes

onnxscript/ir/_linked_list_test.py Fixed Show fixed Hide fixed

justinchuby added 5 commits April 5, 2024 22:02

[IR] Mutable implementation

0e116e4

snap

5fa4957

snap

19f7c08

DoublyLinkedList

f090a37

snap

a37b318

BowenBao reviewed Apr 6, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Outdated Show resolved Hide resolved

BowenBao reviewed Apr 6, 2024

View reviewed changes

onnxscript/ir/_linked_list.py Show resolved Hide resolved

justinchuby added 2 commits April 6, 2024 14:09

Improve / refactor

0ba5665

docs

6ab22d2

titaiwangms self-requested a review April 8, 2024 01:01

titaiwangms reviewed Apr 8, 2024

View reviewed changes

justinchuby added 2 commits April 8, 2024 18:15

Update protocols

ad230b4

Add copyright notice

c92ace1

github-advanced-security bot found potential problems Apr 8, 2024

View reviewed changes

gramalingam reviewed Apr 8, 2024

View reviewed changes

onnxscript/ir/_core.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 8, 2024

View reviewed changes

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

justinchuby and others added 4 commits April 8, 2024 22:59

Update to use id() as keys

88a862d

Merge branch 'main' into justinchu/ir-mutation

c88b3f5

Support single nodes

63e9250

clean up

8037993

github-advanced-security bot found potential problems Apr 8, 2024

View reviewed changes

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

BowenBao approved these changes Apr 9, 2024

View reviewed changes

fix protocol

e19764d

github-advanced-security bot found potential problems Apr 9, 2024

View reviewed changes

onnxscript/ir/_protocols.py Fixed Show fixed Hide fixed

onnxscript/ir/_protocols.py Fixed Show fixed Hide fixed

onnxscript/ir/_protocols.py Fixed Show fixed Hide fixed

onnxscript/ir/_protocols.py Fixed Show fixed Hide fixed

onnxscript/ir/_protocols.py Fixed Show fixed Hide fixed

fix

5b61642

fix

dbc0db5

github-advanced-security bot found potential problems Apr 9, 2024

View reviewed changes

fix arg positions

ae4838c

github-advanced-security bot found potential problems Apr 9, 2024

View reviewed changes

justinchuby merged commit b66c6d1 into main Apr 9, 2024
31 of 37 checks passed

justinchuby deleted the justinchu/ir-mutation branch April 9, 2024 01:51

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[IR] Mutable IR implementation #1344

[IR] Mutable IR implementation #1344

justinchuby commented Apr 4, 2024 •

edited

Loading

codecov bot commented Apr 4, 2024 •

edited

Loading

github-advanced-security bot left a comment

github-actions bot commented Apr 4, 2024 •

edited

Loading

justinchuby commented Apr 5, 2024

justinchuby commented Apr 6, 2024

titaiwangms Apr 8, 2024

justinchuby commented Apr 8, 2024

justinchuby commented Apr 9, 2024

[IR] Mutable IR implementation #1344

[IR] Mutable IR implementation #1344

Conversation

justinchuby commented Apr 4, 2024 • edited Loading

Methods implemented

Safe Iterators

Linked list implementation

Next steps

codecov bot commented Apr 4, 2024 • edited Loading

Codecov Report

github-advanced-security bot left a comment

Choose a reason for hiding this comment

github-actions bot commented Apr 4, 2024 • edited Loading

Test Results

justinchuby commented Apr 5, 2024

justinchuby commented Apr 6, 2024

titaiwangms Apr 8, 2024

Choose a reason for hiding this comment

justinchuby commented Apr 8, 2024

justinchuby commented Apr 9, 2024

justinchuby commented Apr 4, 2024 •

edited

Loading

codecov bot commented Apr 4, 2024 •

edited

Loading

github-actions bot commented Apr 4, 2024 •

edited

Loading