Track stride and make full_like respect permuted axes

[shino16]

Meant to fixes #2680. This work is experimental.

Change 1

torch.full_like(a)'s behavior is changed to align with eager mode when a's axes are permuted. In such cases, full_like(a) in Thunder will now return a view whose axes are permuted likewise.

import torch, thunder

x = torch.randn(2, 5, 6, device="cuda")[:, :, ::2].permute(1, 2, 0)
y = thunder.jit(lambda x: torch.ones_like(x))(x)
print(x.stride())                   # (30, 2, 10)
print(torch.ones_like(x).stride())  # (15, 1, 5)
print(y.stride())                   # main: (6, 2, 1), This PR: (3, 1, 15)

To achieve this in Thunder, clang.full_like(a) creates a view with axes transposed accordingly. This happens only when we need axis permutation.

The same applies to zeros_like and ones_like which rely on full_like.

Change 2

TensorProxy (and FutureTensorProxy) now tracks its stride, because full_like needs to know how its axes were permuted.

Why we need this

In PR #2600 we introduced Inductor for fallback. Inductor needs FakeTensors as fake inputs, and such FakeTensors come from Dynamo, which correctly tracks strides. This leads to discrepancies when full_like in Thunder-compiled submodules creates a tensor with different strides, causing AssertionError (see #2680).

The intention behind this PR is to see how big the changes will be if we track tensor strides. Alternatively, we could stop supporting full_like and leave it to Inductor fallback, or delay Inductor until we get the real input tensors.

TODO

• Make stride tracking complete

[shino16]

Now #2680 does not raise AssertionError, but gives a segmentation fault instead.

python thunder/benchmarks/benchmark_inference.py --input-length 2048 --output-length 512 --mode thunder --num-iterations 10

[70c7086d6c87:10825:0:10825] Caught signal 11 (Segmentation fault: address not mapped to object at address 0x8)
==== backtrace (tid:  10825) ====
 0  /opt/hpcx/ucx/lib/libucs.so.0(ucs_handle_error+0x2e4) [0x7f81a82a5774]
 1  /opt/hpcx/ucx/lib/libucs.so.0(+0x3796a) [0x7f81a82a596a]
 2  /opt/hpcx/ucx/lib/libucs.so.0(+0x37ba8) [0x7f81a82a5ba8]
 3  /usr/lib/x86_64-linux-gnu/libc.so.6(+0x45330) [0x7f81e0162330]
 4  python() [0x551398]
 5  python() [0x5a51d5]
 6  python() [0x552162]
 7  python(_PyObject_MakeTpCall+0x75) [0x5492f5]
 8  python(_PyEval_EvalFrameDefault+0xadf) [0x5d68bf]
 9  python() [0x54ce3d]
10  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
11  python() [0x54ce3d]
12  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
13  python(_PyObject_Call_Prepend+0xc2) [0x54ab42]
14  python() [0x5a30c8]
15  python(_PyObject_MakeTpCall+0x75) [0x5492f5]
16  python(_PyEval_EvalFrameDefault+0xadf) [0x5d68bf]
17  python() [0x54ce3d]
18  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
19  python() [0x54ce3d]
20  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
21  python() [0x54ce3d]
22  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
23  python(_PyObject_Call_Prepend+0xc2) [0x54ab42]
24  python() [0x5a30c8]
25  python(PyObject_Call+0x6c) [0x54b47c]
26  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
27  python() [0x54ce3d]
28  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
29  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
30  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
31  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
32  python() [0x54ce3d]
33  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
34  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
35  python() [0x54ce3d]
36  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
37  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
38  python(_PyObject_Call_Prepend+0xc2) [0x54ab42]
39  python() [0x5a30c8]
40  python(_PyObject_MakeTpCall+0x75) [0x5492f5]
41  python(_PyEval_EvalFrameDefault+0xadf) [0x5d68bf]
42  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(dynamo_eval_custom_code+0x1f1) [0x7f81d398bd61]
43  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603a7) [0x7f81d398d3a7]
44  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
45  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(dynamo_eval_custom_code+0x1f1) [0x7f81d398bd61]
46  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603a7) [0x7f81d398d3a7]
47  python(PyObject_Call+0x115) [0x54b525]
48  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
49  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
50  python() [0x54cf04]
51  python(PyObject_Call+0x115) [0x54b525]
52  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
53  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
54  python() [0x54cf04]
55  python(PyObject_Call+0x115) [0x54b525]
56  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
57  /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so(+0x8603f9) [0x7f81d398d3f9]
58  python(_PyObject_Call_Prepend+0x18a) [0x54ac0a]
59  python() [0x5a30c8]
60  python(PyObject_Call+0x6c) [0x54b47c]
61  python(_PyEval_EvalFrameDefault+0x4cb0) [0x5daa90]
=================================
[1]    10825 segmentation fault (core dumped)  python thunder/benchmarks/benchmark_inference.py --input-length 2048  512