Fix slice followed by reshape (#237)

Co-authored-by: Denis Vieriu <denisvieriu@mac-3045BF.local>

Author: 2 people

aten/src/ATen/native/mps/operations/View.mm

@@ -540,11 +540,26 @@ bool canSliceViewTensor(const Tensor& src, MPSShape *mpsShape) {
   std::vector<int64_t> src_view_shape = getViewShape(src, mpsShape);
   size_t src_ndim_view = src_view_shape.size();
   if (src_ndim_base == src_ndim_view) {
-    for (const auto i : c10::irange(src_ndim_base)) {
+    for (const auto i: c10::irange(src_ndim_base)) {
       if (src_view_shape[i] > src_base_shape[i]) {
         return false;
       }
     }
+  } else {
+    // Detect slice followed by reshape cases, e.g (1,4800,2) -> (1,4800)
+    bool allDimsEqual = true;
+    auto min_ndim = std::min(src_ndim_base, src_ndim_view);
+    for (const auto i: c10::irange(min_ndim)) {
+      if (src_view_shape[i] > src_base_shape[i]) {
+        return false;
+      }
+      else if (src_view_shape[i] != src_base_shape[i]) {
+        allDimsEqual = false;
+      }
+    }
+    if (allDimsEqual) {
+      return false;
+    }
   }
   return true;
 }

test/test_mps.py

@@ -1623,6 +1623,29 @@ def test_cpu_to_strided_mps_copy(self):
 
         self.assertEqual(a1, a2)
 
+    def test_slice_reshape(self):
+        x = torch.randn([1, 6, 4, 2], dtype=torch.float, device="mps")
+        x_cpu = x.detach().clone().to("cpu")
+
+        x = x[:,3:].view(2, 3, 4, 1)
+        x_cpu = x_cpu[:,3:].view(2, 3, 4, 1)
+        self.assertEqual(x, x_cpu)
+
+        x = x + 2
+        x_cpu = x_cpu + 2
+        self.assertEqual(x, x_cpu)
+
+    def test_slice_reshape_contg_view(self):
+        import torch
+
+        x_mps = torch.randn(1, 4800, 2, device="mps")
+        x_cpu = x_mps.detach().clone().cpu()
+
+        r_mps = x_mps + 2
+        r_cpu = x_cpu + 2
+
+        self.assertEqual(r_mps, r_cpu)
+
     def test_view_slice(self):
         # https://github.com/pytorch/pytorch/issues/83995
         NUM_SAMPLES = 60