[MPS] Fix views with 3 or more sliced dimensions (#95762) (#95871)

kulinseth · DenisVieriu97 · web-flow · commit 82b078ba643d · 2023-03-02T12:27:46.000-08:00
Fixes #95482 Pull Request resolved: #95762 Approved by: https://github.com/razarmehr Co-authored-by: Denis Vieriu <dvieriu@apple.com>
diff --git a/aten/src/ATen/native/mps/operations/View.mm b/aten/src/ATen/native/mps/operations/View.mm
@@ -510,7 +510,6 @@ bool canSliceViewTensor(const Tensor& src, MPSShape *mpsShape) {
   MPSNDArrayDescriptor *srcTensorNDArrayDesc = nil;
   MPSNDArray *srcTensorNDArray = nil;
   id<MTLCommandBuffer> commandBuffer = getCurrentMPSStream()->commandBuffer();
-
   int64_t base_idx = 0;
 
   std::vector<int64_t> src_base_shape_vec;
@@ -544,20 +543,20 @@ bool canSliceViewTensor(const Tensor& src, MPSShape *mpsShape) {
   }
 
   int64_t sliceOffset = src.storage_offset() / view_numel;
-  // There are cases where both dimensions of a view can shrink
-  // E.g: x = torch.randn((3,6))[1, 1:3]
-  int64_t nextSliceOffset = 0;
-  bool sliceNextDim = (firstDimToSlice < (src_base_shape.size() - 1)) &&
-                      (src_view_shape[firstDimToSlice + 1] != src_base_shape[firstDimToSlice + 1]);
-
-  [srcTensorNDArrayDesc sliceDimension:src_ndim_base - 1 - firstDimToSlice withSubrange:{static_cast<NSUInteger>(sliceOffset), static_cast<NSUInteger>(src.sizes()[firstDimToSlice])}];
-  if (sliceNextDim) {
-    if (firstDimToSlice + 1 == src_base_shape.size() - 1) {
-      nextSliceOffset = src.storage_offset() % src_base_shape[src_base_shape.size() - 1];
-    } else {
-      nextSliceOffset = (src.storage_offset() % view_numel) / (view_numel / src_base_shape[firstDimToSlice + 1]);
+  [srcTensorNDArrayDesc sliceDimension:src_ndim_base - 1 - firstDimToSlice
+                          withSubrange:{static_cast<NSUInteger>(sliceOffset), static_cast<NSUInteger>(src.sizes()[firstDimToSlice])}];
+
+  // Slice any remaining dimensions
+  for (const auto crtSliceOffset: c10::irange(firstDimToSlice + 1, src_base_shape.size())) {
+    if (src_view_shape[crtSliceOffset] != src_base_shape[crtSliceOffset]) {
+      if (crtSliceOffset == src_base_shape.size() - 1) {
+        sliceOffset = src.storage_offset() % src_base_shape[src_base_shape.size() - 1];
+      } else {
+        sliceOffset = (src.storage_offset() % view_numel) / (view_numel / src_base_shape[crtSliceOffset]);
+      }
+      [srcTensorNDArrayDesc sliceDimension:src_ndim_base - 1 - crtSliceOffset
+                              withSubrange:{static_cast<NSUInteger>(sliceOffset), static_cast<NSUInteger>(src.sizes()[crtSliceOffset])}];
     }
-    [srcTensorNDArrayDesc sliceDimension:src_ndim_base - 2 - firstDimToSlice withSubrange:{static_cast<NSUInteger>(nextSliceOffset), static_cast<NSUInteger>(src.sizes()[firstDimToSlice+1])}];
   }
   srcTensorNDArrayView = [srcTensorNDArray arrayViewWithCommandBuffer:commandBuffer
                                                            descriptor:srcTensorNDArrayDesc
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -1884,6 +1884,15 @@ def helper(shape):
         helper([3, 4, 18, 22])
         helper([3, 4, 18, 22, 150])
 
+    def test_contiguous_slice_3d(self):
+        x = torch.randn(2, 3, 3, device="mps")
+        x_cpu = x.detach().clone().cpu()
+        x = x[:1]
+        x_cpu = x_cpu[:1]
+        out = x[:, 0:1, 0:1] * x[:, 1:2, 1:2]
+        out_cpu = x_cpu[:, 0:1, 0:1] * x_cpu[:, 1:2, 1:2]
+        self.assertEqual(out, out_cpu)
+
     def test_view_slice(self):
         # https://github.com/pytorch/pytorch/issues/83995
         NUM_SAMPLES = 60
