[MPS] Fix sliced cast (#138535)
[MPS] Fix sliced cast (#138314)
This fixes internal crash due to the invalid bufer size computation if sliced API is used
Not sure what was the purpose of
```c++
IntArrayRef baseShape;
if (src.is_view()) {
baseShape = src._base().sizes();
} else {
baseShape = getIMPSAllocator()->getBufferShape(src.storage().data());
}
int flattenedShaped = 1;
for (const auto i : c10::irange(baseShape.size())) {
flattenedShaped *= baseShape[i];
}
```
As flattenShaped could be much easier computed as `[srcBuf
lengh]/src.element_size()`, and even if `srcBuf` is padded it's a safe thing to do.
When someone allocated buffer to hold say uint8 and that view-casted it
to float16, attempt to compute `baseShape` returned sizes of original
tensor in its data type, rather than size in new dtypes
Fixes https://github.com/pytorch/pytorch/issues/137800
Pull Request resolved: https://github.com/pytorch/pytorch/pull/138314
Approved by: https://github.com/albanD, https://github.com/DenisVieriu97
(cherry picked from commit de16159e565e7a08294347e31e97ca08a3468227)
Co-authored-by: Nikita Shulga <[email protected]>diff --git a/test/test_mps.py b/test/test_mps.py
index 03dab90..ad9316c 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -10964,6 +10964,12 @@
t1.start()
t2.start()
+ def test_sliced_view_cast(self):
+ # This used to crash on MacOS Sequoia
+ # See https://github.com/pytorch/pytorch/issues/137800
+ x = torch.rand(16, 16, device='mps', dtype=torch.float16)
+ y = x[:, 0:2].view(torch.float32) + 1
+
def test_masked_select(self):
x = torch.randn(3, 4)
x_mps = x.to("mps")
