fixes

alexdremov · alexdremov · commit dc8bab3a7fe4 · 2023-08-26T16:33:50.000+03:00
diff --git a/aten/src/ATen/native/mps/operations/PixelShuffle.mm b/aten/src/ATen/native/mps/operations/PixelShuffle.mm
@@ -21,7 +21,6 @@ static Tensor pixel_shuffle_helper(const Tensor& self, int64_t factor, bool upsc
     check_pixel_unshuffle_shapes(self, factor);
   }
 
-  MPSGraphCache* cache_ = MPSGraphCache::getInstance();
   MPSStream* stream = getCurrentMPSStream();
 
   const int64_t c = self.size(-3);
@@ -47,43 +46,32 @@ static Tensor pixel_shuffle_helper(const Tensor& self, int64_t factor, bool upsc
   @autoreleasepool {
     string key = (upscale ? "pixel_shuffle_" : "pixel_unshuffle_") + getTensorsStringKey({self}) + "_factor_" +
         std::to_string(factor);
-    CachedGraph* cachedGraph = cache_->LookUpAs<CachedGraph>(key);
-    if (!cachedGraph) {
-      MPSCachedGraph* tmpCachedGraph = cache_->CreateCachedGraph(key, ^MPSCachedGraph*() {
-        CachedGraph* newCachedGraph = nil;
-
-        @autoreleasepool {
-          MPSGraph* mpsGraph = make_mps_graph();
-          newCachedGraph = new CachedGraph(mpsGraph);
-
-          const auto ndims = self.ndimension();
-          MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
-          MPSGraphTensor* outputTensor;
-          if (upscale) {
-            outputTensor = [mpsGraph depthToSpace2DTensor:inputTensor
-                                                widthAxis:ndims - 1
-                                               heightAxis:ndims - 2
-                                                depthAxis:ndims - 3
-                                                blockSize:factor
-                                     usePixelShuffleOrder:YES
-                                                     name:nil];
-          } else {
-            outputTensor = [mpsGraph spaceToDepth2DTensor:inputTensor
-                                                widthAxis:ndims - 1
-                                               heightAxis:ndims - 2
-                                                depthAxis:ndims - 3
-                                                blockSize:factor
-                                     usePixelShuffleOrder:YES
-                                                     name:nil];
-          }
-
-          newCachedGraph->inputTensor_ = inputTensor;
-          newCachedGraph->outputTensor_ = outputTensor;
-        }
-        return newCachedGraph;
-      });
-      cachedGraph = static_cast<CachedGraph*>(tmpCachedGraph);
-    }
+    CachedGraph* cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
+      const auto ndims = self.ndimension();
+      MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self);
+      MPSGraphTensor* outputTensor = nullptr;
+      if (upscale) {
+        outputTensor = [mpsGraph depthToSpace2DTensor:inputTensor
+                                            widthAxis:ndims - 1
+                                           heightAxis:ndims - 2
+                                            depthAxis:ndims - 3
+                                            blockSize:factor
+                                 usePixelShuffleOrder:YES
+                                                 name:nil];
+      } else {
+        outputTensor = [mpsGraph spaceToDepth2DTensor:inputTensor
+                                            widthAxis:ndims - 1
+                                           heightAxis:ndims - 2
+                                            depthAxis:ndims - 3
+                                            blockSize:factor
+                                 usePixelShuffleOrder:YES
+                                                 name:nil];
+      }
+
+      newCachedGraph->inputTensor_ = inputTensor;
+      newCachedGraph->outputTensor_ = outputTensor;
+      return newCachedGraph;
+    });
 
     Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self);
     Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output);
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -949,11 +949,38 @@ def leak_gpu0():
             leak_gpu0()
 
 
-# These tests were taken from test/test_nn.py
 class TestPixelShuffle(TestCaseMPS):
     def test_pixel_shuffle_unshuffle(self):
         def _test_pixel_shuffle_unshuffle_helper(num_input_dims, valid_channels_dim=True,
-                                                 upscale_factor=None):
+                                                 upscale_factor=None, is_contiguous=True):
+
+            def generate_input():
+                # If valid_channels_dim=False, add 1 to make channels dim indivisible by upscale_factor ** 2.
+                channels = random.randint(1, 4) * upscale_factor ** 2 + (0 if valid_channels_dim else 1)
+                height = random.randint(5, 10)
+                width = random.randint(5, 10)
+
+                if num_input_dims == 1:
+                    input = torch.rand(channels, requires_grad=True, device='mps')
+                    assert is_contiguous
+                elif num_input_dims == 2:
+                    input = torch.rand(width, height, requires_grad=True, device='mps').T
+                    if is_contiguous:
+                        input = input.contiguous()
+                else:
+                    batch_sizes = [random.randint(1, 3) for _ in range(num_input_dims - 3)]
+                    input = torch.rand(*batch_sizes, channels, width, height, requires_grad=True, device='mps')
+                    input = input.transpose(-1, -2)
+                    if is_contiguous:
+                        input = input.contiguous()
+
+                if not is_contiguous and len(input.reshape(-1)) > 0:
+                    assert not input.is_contiguous()
+
+                input = input.detach().clone()
+                input.requires_grad = True
+                return input
+
             # Function to imperatively ensure pixels are shuffled to the correct locations.
             # Used to validate the batch operations in pixel_shuffle.
             def _verify_pixel_shuffle(input, output, upscale_factor):
@@ -967,18 +994,7 @@ def _verify_pixel_shuffle(input, output, upscale_factor):
                             self.assertEqual(output[..., c, h, w], input[..., channel_idx, height_idx, weight_idx])
 
             upscale_factor = random.randint(2, 5) if upscale_factor is None else upscale_factor
-            # If valid_channels_dim=False, add 1 to make channels dim indivisible by upscale_factor ** 2.
-            channels = random.randint(1, 4) * upscale_factor ** 2 + (0 if valid_channels_dim else 1)
-            height = random.randint(5, 10)
-            width = random.randint(5, 10)
-
-            if num_input_dims == 1:
-                input = torch.rand(channels, requires_grad=True, device='mps')
-            elif num_input_dims == 2:
-                input = torch.rand(height, width, requires_grad=True, device='mps')
-            else:
-                batch_sizes = [random.randint(1, 3) for _ in range(num_input_dims - 3)]
-                input = torch.rand(*batch_sizes, channels, height, width, requires_grad=True, device='mps')
+            input = generate_input()
 
             ps = nn.PixelShuffle(upscale_factor)
             pus = nn.PixelUnshuffle(downscale_factor=upscale_factor)
@@ -1018,14 +1034,22 @@ def _test_pixel_unshuffle_error_case_helper(num_input_dims, valid_height_dim=Tru
         def _test_pixel_shuffle_unshuffle_for_input_dims(num_input_dims):
             # For 1D - 2D, this is an error case.
             # For 3D - 5D, this is a success case for pixel_shuffle + pixel_unshuffle.
-            _test_pixel_shuffle_unshuffle_helper(num_input_dims=num_input_dims)
-
-            # Error cases for pixel_shuffle.
-            _test_pixel_shuffle_unshuffle_helper(num_input_dims=num_input_dims, valid_channels_dim=False)
-            _test_pixel_shuffle_unshuffle_helper(num_input_dims=num_input_dims, upscale_factor=0)
-            _test_pixel_shuffle_unshuffle_helper(num_input_dims=num_input_dims, upscale_factor=-2)
-
-            # Error cases for pixel_unshuffle.
+            is_contiguous_check = [True, False] if num_input_dims > 1 else [True]
+            for is_contiguous in is_contiguous_check:
+                _test_pixel_shuffle_unshuffle_helper(
+                    num_input_dims=num_input_dims, is_contiguous=is_contiguous
+                )
+                _test_pixel_shuffle_unshuffle_helper(
+                    num_input_dims=num_input_dims, valid_channels_dim=False, is_contiguous=is_contiguous
+                )
+                _test_pixel_shuffle_unshuffle_helper(
+                    num_input_dims=num_input_dims, upscale_factor=0, is_contiguous=is_contiguous
+                )
+                _test_pixel_shuffle_unshuffle_helper(
+                    num_input_dims=num_input_dims, upscale_factor=-2, is_contiguous=is_contiguous
+                )
+
+                # Error cases for pixel_unshuffle.
             _test_pixel_unshuffle_error_case_helper(num_input_dims=num_input_dims, valid_height_dim=False)
             _test_pixel_unshuffle_error_case_helper(num_input_dims=num_input_dims, valid_width_dim=False)
             _test_pixel_unshuffle_error_case_helper(num_input_dims=num_input_dims, downscale_factor=0)
