[IR] introduce slice support on graph

[Johansmm]

Introduce slice support
Close #2302

[justinchuby]

Thanks!

[justinchuby]

Could you also update the following text to match the implementation?

        although accessing nodes at either end of the set is O(1). I.e.
        ``linked_set[0]`` and ``linked_set[-1]`` are O(1).

in the class docstring above.

also https://github.com/Johansmm/onnxscript/blob/862de43e60f51ee4833104a254da9d8db80420a6/onnxscript/ir/_core.py#L2319 on complexity?

There are two type annotations to update: https://github.com/Johansmm/onnxscript/blob/862de43e60f51ee4833104a254da9d8db80420a6/onnxscript/ir/_core.py#L2715 and https://github.com/Johansmm/onnxscript/blob/862de43e60f51ee4833104a254da9d8db80420a6/onnxscript/ir/_core.py#L2285 . Please follow the recommendation below to create typing overloads for the methods

[justinchuby]

Some comments, thank you!

[codecov]

Codecov Report

Attention: Patch coverage is 72.41379% with 8 lines in your changes missing coverage. Please review.

Project coverage is 73.76%. Comparing base (644e30c) to head (e29419f).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/ir/_core.py	60.00%	0 Missing and 6 partials ⚠️
onnxscript/ir/_linked_list.py	75.00%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2307      +/-   ##
==========================================
- Coverage   73.76%   73.76%   -0.01%     
==========================================
  Files         239      239              
  Lines       30906    30930      +24     
  Branches     3494     3503       +9     
==========================================
+ Hits        22799    22815      +16     
  Misses       6907     6907              
- Partials     1200     1208       +8     

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

[Johansmm]

Last push force with suggestions

[justinchuby]

The final thing I wanted to make sure before merge is that accessing nodes on either ends (graph.node[0], graph.node[-1]) is not much slower (better if it turns out to be faster) than before. Would you be able to run a small test?

[Johansmm]

@justinchuby, the new complexity is O(n) for any case.
I made the following python script to evaluate this:

import time

from onnxscript.ir import _linked_list


class _TestElement:
    def __init__(self, value):
        self.value = value

    def __repr__(self) -> str:
        return f"_TestElement({self.value})"


N = 100000
linked_list = _linked_list.DoublyLinkedSet()
elems = [_TestElement(i) for i in range(N)]
linked_list.extend(elems)

time_list_1 = []
idx = -1
for _ in range(1000):
    t = time.time()
    x = tuple(linked_list)[idx]
    t = time.time() - t
    time_list_1.append(t)

time_list_2 = []
for _ in range(1000):
    t = time.time()
    for i in reversed(linked_list):
        x = i
        break
    t = time.time() - t
    time_list_2.append(t)


print(sum(time_list_1) / 1000)
print(sum(time_list_2) / 1000)

Results changing N:

N	time previous	time now
100	0.0	2.5733470916748046e-05
1000	0.0	0.00013534140586853028
10000	0.0	0.0014723334312438964
100000	1.3112068176269532e-05	0.015811803340911864

Therefore:

1. do you want me to reverse the code as long as I handle the integers separately?
2. I leave the code as it is, since it is not common to have models with more than 1000 nodes ?

[justinchuby]

It is in fact quite common to have models with more than 1000 nodes in the case of LLMs. I would prefer it to run faster. Could you do (1) as you suggested, and revert the comment changes? Thanks a lot! I think it makes sense to handle slices in an if branch.

[Johansmm]

@justinchuby, I have updated the code to produce a more optimal response. Now __getitem__ is fully O(n).
Please check it and let me know what you think.
Unfortunately, now __getitem__ returns a list instead of a tuple, because items = items + (item,) has O(n) at each iteration, but items.append(item) has O(1) (at each iteration).
I hope these changes are relevant.

[justinchuby]

Thank you! The changes look very reasonable. Were you able to run the benchmark again with some examples?

[Johansmm]

@justinchuby can you explain what benchmark you are talking about? If you are referring to the IR unit tests, I ran them and they passed without any problem.

[justinchuby]

I was referring to #2307 (comment) this script in your comment

[Copilot]

Pull Request Overview

This PR adds slice support for accessing elements from the graph‐related collections and linked list implementation. Key changes include:

• Addition of parameterized tests for slicing in _linked_list_test.py.
• Implementation of slice overloads and corresponding logic in DoublyLinkedSet in _linked_list.py.
• Similar overload enhancements for getitem in node collections in _core.py.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
onnxscript/ir/_linked_list_test.py	Added parameterized test to validate slice support
onnxscript/ir/_linked_list.py	Updated getitem to support slicing with overloads and custom logic
onnxscript/ir/_core.py	Added overload annotations for slice support in node collection access

[justinchuby]

Sorry after taking another look I wonder if the added complexity is worth it (line 160-162, 166-167). Would it be a good idea to (1) keep the original logic for int index, and (2) use tuple for slices?

[Johansmm]

@justinchuby, I see that you don't agree with my algorithm, so I have gone back to the initial solution. Note that this one is a bit more expensive.
Do not hesitate to ask me for anything else if necessary 👍

[justinchuby]

Thank you! I think it’s a good trade off for simplicity

[justinchuby]

Hi @Johansmm, we are in the process of migrating the IR component to https://github.com/onnx/ir-py so that it becomes an official onnx project. Would you like to help us duplicate your PR to the new repository as well? We are looking to create an initial release within about a week or less.

[Johansmm]

@justinchuby done: onnx/ir-py#25