fix aliasing for primtorch view meta kernels

[bdhirsh]

Fixes #86284

Stack from ghstack (oldest at bottom):

• ban .sizes() and .strides() calls in derivatives.yaml #86611
• symintify all of derivatives.yaml #86610
• min/max support for SymInt/Floats, finish as_strided/scatter/squeeze() backward symint support #86609
• [test] add cross-ref tests for python meta kernels #86228
• fix aliasing bug in pixel shuffle/unshuffle #86608
• fix type promotion for group_norm composite C++ kernel #86607
• -> fix aliasing for primtorch view meta kernels #86285

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/86285

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 4b7afaa:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mruberry]

The approach looks reasonable to me. What about extending test_python_ref_meta to detect this:

pytorch/test/test_ops.py

Line 174 in 17addb3

def test_python_ref_meta(self, device, dtype, op):

[IvanYashchuk]

Is the only difference here that with alias=True the resulting tensor would return True for tensor._is_view()?

import torch
a = torch.zeros(3, 2, device="meta")
b = torch._prims.TensorMeta(a, shape=a.shape, strides=a.stride())
print(b._is_view()) # False
c = a.as_strided(a.shape, a.stride())
print(c._is_view()) # True
assert a.storage().data_ptr() == b.storage().data_ptr() == c.storage().data_ptr() # It's all 0 for meta

[bdhirsh]

@IvanYashchuk since the output isn't a view of the input, then its metadata can also be incorrect. For example if the input is some strided tensor with a non-zero storage offset, then that won't get propagated to the output:

>>> a = torch.ones(2, 2, device='meta')[1]
>>> a.storage_offset()
2
>>> out_aten = a.as_strided(a.shape, a.stride())
>>> out_aten.storage_offset() # prints 2, the same as the input
2
>>> out_prim = torch._prims.TensorMeta(a, shape=a.shape, strides=a.stride())
>>> out_prim.storage_offset() # should print 2, prints 0!
0
>>>

[bdhirsh]

@mruberry Thanks for the pointer to test_python_ref_meta. It looks like today, that test currently checks the output shapes match, but not strides or storage offsets. The two ways I could imagine updating that test are:

(1) Update the test to check for the _is_view() relationship (this seems a bit weird since it's a private API that's dependent on autograd, but maybe this would be fine?)
(1) updating the test to always check that storage offsets match. This is the "broken" thing that I originally noticed. That would mean that the tests always check that sizes and storage_offset match, but no strides. Does that sound reasonable to you?

FWIW, I'm also adding a cross ref test that tests correctness of all of the python decomps and meta functions in torch/_decomp/decompositions.py and torch/_meta_registrations.py. Some of those decomps call into prims, which is where the problem first showed up. Although those two files probably don't comprehensively call into the prims, so a separate prim test sounds good to me.

[mruberry]

@mruberry Thanks for the pointer to test_python_ref_meta. It looks like today, that test currently checks the output shapes match, but not strides or storage offsets. The two ways I could imagine updating that test are:

(1) Update the test to check for the _is_view() relationship (this seems a bit weird since it's a private API that's dependent on autograd, but maybe this would be fine?)

Maybe TestViewOps in test_view_ops.py can provide some inspiration. See

pytorch/test/test_view_ops.py

Line 96 in cebf08a

def is_view_of(self, base, other):

(1) updating the test to always check that storage offsets match. This is the "broken" thing that I originally noticed. That would mean that the tests always check that sizes and storage_offset match, but no strides. Does that sound reasonable to you?

Stride-testing has been disabled for the moment. I spent a good amount of time trying to emulate PyTorch's striding logic, but PyTorch is pretty inconsistent in how it handles strides itself. See #78050. @ezyang concludes that we'd like to emulate strides, but I'm not sure how we would do this.

FWIW, I'm also adding a cross ref test that tests correctness of all of the python decomps and meta functions in torch/_decomp/decompositions.py and torch/_meta_registrations.py. Some of those decomps call into prims, which is where the problem first showed up. Although those two files probably don't comprehensively call into the prims, so a separate prim test sounds good to me.

There are several existing tests in test_ops.py for Python reference consistency, see

pytorch/test/test_ops.py

Line 339 in cebf08a

def test_python_ref(self, device, dtype, op):

These might be interesting to look at when developing a consistency test for decompositions. Ideally I think we'd like to see all decompositions ported to become Python references.

[ezyang]

@bdhirsh, why don't you just convert these into direct as_strided calls on the input tensor, rather than going through the TensorMeta constructor?

[bdhirsh]

@bdhirsh, why don't you just convert these into direct as_strided calls on the input tensor, rather than going through the TensorMeta constructor?

Sounds good - you're right, seems cleaner

[bdhirsh]

@mruberry Updated the PR - I had some issues trying to test the _is_view() property. Since that's not really the property that I cared about in the first place (storage offset was incorrect), I updated the op info tests to check for storage offset, which I confirmed exercises the problem from #86284. Let me know if you're happy with it

[albanD]

The storage_offset argument is missing. Sounds good otherwise.

[albanD]

All the functions here need a storage_offset=a.storage_offset() argument!

[albanD]

SGTM

[mruberry]

LGTM!