feat: support aten.roll dynamo converter

[zewenli98]

Description

Support aten.roll dynamo converter.

Fixes #2567

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[gs-olive]

Consider adding a case where both shifts and dims are single integers, which is a supported case in the docstring. These may be casted to lists in the operator before the converter ever gets them, but it is still a valid input I believe.

[zewenli98]

Thanks for your review! I found the schema is:

- func: roll(Tensor self, SymInt[1] shifts, int[1] dims=[]) -> Tensor

Does this mean shifts and dims should be a 1d list?

[gs-olive]

roll(tensor, 2, 3) --> roll(tensor, [2], [3])
pool(3) --> pool([3, 3])
To share additional documentation

[gs-olive]

@zewenli98 see the comment here:

https://github.com/pytorch/pytorch/blob/40dbd567e04483c671f9c897171bf9d1e7162b68/torch/csrc/jit/frontend/schema_matching.cpp#L173-L180

[zewenli98]

@gs-olive Thanks for the details!
Unfortunately, when testing shifts=2, dims=0, I got error:

File "<eval_with_key>.0 from /home/zewenl/TensorRT/tests/py/dynamo/conversion/test_roll_aten.py:35 in forward", line 5, in forward
    roll_default = torch.ops.aten.roll.default(x, 2, 0);  x = None
  File "/home/zewenl/.local/lib/python3.10/site-packages/torch/_ops.py", line 571, in __call__
    return self_._op(*args, **(kwargs or {}))
RuntimeError: aten::roll() Expected a value of type 'List[int]' for argument 'shifts' but instead found type 'int'.
Position: 1
Value: 2
Declaration: aten::roll(Tensor self, SymInt[1] shifts, int[1] dims=[]) -> Tensor
 Python error details: TypeError: 'int' object is not iterable

Then, I also tested shifts=(2,), dims=0, it works.

It seems that pytorch requires shifts to be a list.

According to the schema - func: roll(Tensor self, SymInt[1] shifts, int[1] dims=[]) -> Tensor, I guess SymInt[1] and int[1] may have different behaviors?

[gs-olive]

It might be so, yes, though that is strange - thanks for the update.

[zewenli98]

Yes, I'm working on adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor converter, output_size expects List[int] as well:

RuntimeError: aten::adaptive_avg_pool2d() Expected a value of type 'List[int]' for argument 'output_size' but instead found type 'int'.
Position: 1
Value: 3
Declaration: aten::adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor