Fixed a bug in interpolate uint8 AVX2 on non-contig input (pytorch#101136)

Description:
- Fixed a bug in interpolate uint8 AVX2 on non-contig input
- Added tests

Pull Request resolved: pytorch#101136
Approved by: https://github.com/NicolasHug

Author: vfdev-5

aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h

@@ -266,6 +266,7 @@ void ImagingResampleVertical(
   auto xout = unpacked_output.size(2);
   auto yout = unpacked_output.size(1);
   const auto num_channels = unpacked_input.size(0);
+  TORCH_INTERNAL_ASSERT(num_channels == unpacked_output.size(0));
 
   auto xout_stride = xout * num_channels;
   for (const auto yy : c10::irange(yout)) {
@@ -301,23 +302,34 @@ void ImagingResampleVertical(
 // weights, but when aa=False they could be optimized further.
 template <typename scale_type, class F>
 void upsample_avx_bilinear_uint8(
-    const at::Tensor& input,
+    const at::Tensor& input_,
     const at::Tensor& output,
     bool align_corners,
     const scale_type& scales,
     bool antialias) {
-  auto batch_size = input.size(0);
-  auto num_channels = input.size(1);
-  auto xin = input.size(3);
-  auto yin = input.size(2);
+  auto batch_size = input_.size(0);
+  auto num_channels = input_.size(1);
+  auto xin = input_.size(3);
+  auto yin = input_.size(2);
   auto xout = output.size(3);
   auto yout = output.size(2);
 
   if (xin == xout && yin == yout) {
-    output.copy_(input);
+    output.copy_(input_);
     return;
   }
 
+  at::Tensor input = input_;
+  if (!(input.is_contiguous() || input.is_contiguous(at::MemoryFormat::ChannelsLast))) {
+    // If input is not contiguous with memory format channels first or channels last,
+    // we explicitly convert the input to contiguous channels last memory format.
+    // This simplifies the rest of the code and let us assume that the format is only contiguous channels first or channels last,
+    // Most tensors going through this `if` block won't need to go through unpacking, but those having C < 3 may
+    // have to (this means 2 copies are made). We could avoid the extra copy by handling non-contiguous input
+    // directly within unpack_rgb() and pack_rgb(), but initial attempts showed that this is fairly complex.
+    input = input.contiguous(at::MemoryFormat::ChannelsLast);
+  }
+
   auto need_horizontal = xout != xin;
   auto need_vertical = yout != yin;
 

test/test_nn.py

@@ -9674,18 +9674,35 @@ def test_upsamplingBilinear2d_aa_correctness(self, device, memory_format):
     @parametrize_test("num_channels", [3, 5])
     @parametrize_test("output_size", [32, 600])
     @parametrize_test("check_as_unsqueezed_3d_tensor", [True, False])
+    @parametrize_test("non_contig", [False, "sliced", "restrided"])
+    @parametrize_test("batch_size", [1, 5])
     def test_upsamplingBiLinear2d_consistency(
-        self, device, memory_format, antialias, align_corners, num_channels, output_size, check_as_unsqueezed_3d_tensor
+        self,
+        device,
+        memory_format,
+        antialias,
+        align_corners,
+        num_channels,
+        output_size,
+        check_as_unsqueezed_3d_tensor,
+        non_contig,
+        batch_size,
     ):
         if torch.device(device).type == "cuda":
             raise SkipTest("CUDA implementation is not yet supporting uint8")
 
         mode = "bilinear"
-        # Check if Max Abs Error between resized input_uint8 and resized input_float is smaller than a tolerated value, e.g. 1.0
-        input_ui8 = torch.randint(0, 256, size=(1, num_channels, 400, 400), dtype=torch.uint8, device=device)
+        # Check if Max Abs Error between resized input_uint8 and resized input_float is
+        # smaller than a tolerated value, e.g. 1.0
+        input_ui8 = torch.randint(0, 256, size=(batch_size, num_channels, 400, 400), dtype=torch.uint8, device=device)
         input_ui8 = input_ui8.contiguous(memory_format=memory_format)
 
-        if check_as_unsqueezed_3d_tensor:
+        if non_contig == "sliced":
+            input_ui8 = input_ui8[:, :, 10:-10, 10:-10]
+        elif non_contig == "restrided":
+            input_ui8 = input_ui8[:, :, ::2, ::2]
+
+        if batch_size == 1 and check_as_unsqueezed_3d_tensor:
             input_ui8 = input_ui8[0, ...]
             input_ui8 = input_ui8[None, ...]
 
@@ -9698,15 +9715,16 @@ def test_upsamplingBiLinear2d_consistency(
             input_ui8, size=(output_size, output_size), mode=mode, align_corners=align_corners, antialias=antialias
         )
 
+        if non_contig is False:
+            self.assertTrue(input_ui8.is_contiguous(memory_format=memory_format))
+
         # FIXME if-clause shows the current behaviour which is definitely unexpected.
         # Ideally we want to fix it such that both the ui8 and f32 outputs are also channels_last
         # See for more details: https://github.com/pytorch/pytorch/pull/100373
-        if check_as_unsqueezed_3d_tensor and memory_format == torch.channels_last:
-            self.assertTrue(input_ui8.is_contiguous(memory_format=torch.channels_last))
+        if batch_size == 1 and check_as_unsqueezed_3d_tensor and memory_format == torch.channels_last:
             self.assertTrue(output_ui8.is_contiguous())
             self.assertTrue(output_f32.is_contiguous())
         else:
-            self.assertTrue(input_ui8.is_contiguous(memory_format=memory_format))
             self.assertTrue(output_ui8.is_contiguous(memory_format=memory_format))
             self.assertTrue(output_f32.is_contiguous(memory_format=memory_format))