Add stride operation primitive

[riccardofelluga]

🚀 Feature

I would like to have Thunder manage stride information to allow for tensor manipulations.

In particular I think the following points need to be discussed:

• Where does the stride information go? Should it be part of TensorProxy?
• What can a stride manipulation primitive look like? Are there any particular things we need to be careful about?

Motivation

This will enable us to add new operators that reshape the tensor using the stride like torch.as_strided or torch.Tensor.unfold.

[mruberry]

Hey @riccardofelluga!

fyi @jjsjann123, too

We actually have a stride order primitive today:

lightning-thunder/thunder/core/prims.py

Line 158 in 393828f

STRIDE_ORDER = auto()

We don't have plans to implement support for setting exact strides, just to reorder them (like the stride order primitive does). What were you thinking of?

[riccardofelluga]

Hey! So what I was thinking was more about setting exact strides.

However, @IvanYashchuk pointed out to me that this discussion has been held once before. That discussion led to the following conclusion:

[...] we did identify that having a stride semantics in the graph is useful, but it's very hard to do it right and efficiently, without putting unnecessary constraints on backends.

Additionally, there is a comment on jax that I think summarizes enough the possible complexity of having such a feature.

In conclusion, I think it's fair to say that unless there is a stronger need for it, it's better to put this idea aside for now