[MPS] Skip gather/blit calls in case of strided output (#94260)
Skip gather/blit calls in case of strided output - this prevents:
- allocating additional memory for the output
- additional transpose for both the input and output
Fixes:
```
x = torch.rand((256,10), device='mps')
x = x.permute(1,0)
x.exp()
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94260
Approved by: https://github.com/razarmehr
diff --git a/test/test_mps.py b/test/test_mps.py
index 9ecaa30..3b3bf9e 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -251,6 +251,17 @@
input = torch.tensor([-0.1, 3.0, -0.9]).to('mps')
output = torch.exp(input).to('cpu')
+ def test_exp_strided_output(self):
+ x = torch.rand((256, 10), device='mps')
+ x_cpu = x.to("cpu")
+
+ x = x.permute(1, 0)
+ x_cpu = x_cpu.permute(1, 0)
+
+ res = x.exp()
+ res_cpu = x_cpu.exp()
+ self.assertEqual(res, res_cpu)
+
def _testLeakyRelu(self, np_features, negative_slope, device):
cpu_x = torch.from_numpy(np_features).requires_grad_()
mps_x = torch.from_numpy(np_features).to('mps').requires_grad_()
