[MPS][BE] Combine two upsample_kernel_out_template into one (pytorch#148211)

- First, by stopp inverting sizes and strides, i.e. passing them as is, but reading them in inverse order in the shader as 1st stride of 4D tensor is one used for batches, 2nd for channels and 3rd and 4th for spatial coordinates
- Pass `scales` as float2 even in linear tensor
Above  allows one to collide two flavors `upsample_kernel_out_template` into one
Pull Request resolved: pytorch#148211
Approved by: https://github.com/dcci
ghstack dependencies: pytorch#148154, pytorch#148187

Author: malfet

aten/src/ATen/native/mps/kernels/UpSample.metal

@@ -123,8 +123,8 @@ scalar_t upsample_get_value_bounded(
   int access_y = max(min(y, dim.y - 1), 0L);
   int access_x = max(min(x, dim.x - 1), 0L);
   return data
-      [n * strides.w + c * strides.z + access_y * strides.y +
-       access_x * strides.x];
+      [n * strides.x + c * strides.y + access_y * strides.z +
+       access_x * strides.w];
 }
 
 template <typename scalar_t>
@@ -136,7 +136,7 @@ scalar_t upsample_get_value_bounded(
     long c,
     long x) {
   int access_x = max(min(x, dim - 1), 0L);
-  return data[n * strides.z + c * strides.y + access_x * strides.x];
+  return data[n * strides.x + c * strides.y + access_x * strides.z];
 }
 
 template <typename scalar_t>
@@ -153,8 +153,8 @@ void upsample_increment_value_bounded(
   int access_x = max(min(x, dim.x - 1), 0L);
   AtomicType<scalar_t>::atomic_add(
       data,
-      n * strides.w + c * strides.z + access_y * strides.y +
-          access_x * strides.x,
+      n * strides.x + c * strides.y + access_y * strides.z +
+          access_x * strides.w,
       value);
 }
 
@@ -200,24 +200,24 @@ kernel void upsample_linear1d(
     constant ulong3& output_strides [[buffer(3)]],
     constant long3& input_sizes [[buffer(4)]],
     constant long3& output_sizes [[buffer(5)]],
-    constant float& scale [[buffer(6)]],
+    constant float2& scales [[buffer(6)]],
     constant bool& align_corners [[buffer(7)]],
     uint thread_index [[thread_position_in_grid]]) {
   auto output_x = thread_index;
   auto real_x = area_pixel_compute_source_index(
-      scale, output_x, align_corners, /*cubic=*/false);
+      scales.x, output_x, align_corners, /*cubic=*/false);
   auto t_x = fract(real_x);
 
-  for (int n = 0; n < output_sizes.z; n++) {
+  for (int n = 0; n < output_sizes.x; n++) {
     for (int c = 0; c < output_sizes.y; c++) {
       auto i00 = upsample_get_value_bounded<T>(
-          inputData, input_sizes.x, input_strides, n, c, real_x);
+          inputData, input_sizes.z, input_strides, n, c, real_x);
       auto i01 = upsample_get_value_bounded<T>(
-          inputData, input_sizes.x, input_strides, n, c, real_x + 1);
+          inputData, input_sizes.z, input_strides, n, c, real_x + 1);
       auto res = linear_interp(i00, i01, t_x);
       outputData
-          [n * output_strides.z + c * output_strides.y +
-           output_x * output_strides.x] = static_cast<T>(res);
+          [n * output_strides.x + c * output_strides.y +
+           output_x * output_strides.z] = static_cast<T>(res);
     }
   }
 }
@@ -232,26 +232,26 @@ kernel void upsample_bilinear2d(
     constant float2& scales [[buffer(6)]],
     constant bool& align_corners [[buffer(7)]],
     uint thread_index [[thread_position_in_grid]]) {
-  auto output_x = thread_index % output_sizes.x;
-  auto output_y = thread_index / output_sizes.x;
+  auto output_x = thread_index % output_sizes.w;
+  auto output_y = thread_index / output_sizes.w;
   auto real_x = area_pixel_compute_source_index(
       scales.x, output_x, align_corners, /*cubic=*/false);
   auto t_x = fract(real_x);
 
   auto real_y = area_pixel_compute_source_index(
       scales.y, output_y, align_corners, /*cubic=*/false);
   auto t_y = fract(real_y);
-  for (int n = 0; n < output_sizes.w; n++) {
-    for (int c = 0; c < output_sizes.z; c++) {
+  for (int n = 0; n < output_sizes.x; n++) {
+    for (int c = 0; c < output_sizes.y; c++) {
       auto i00 = upsample_get_value_bounded<T>(
-          inputData, input_sizes.xy, input_strides, n, c, real_y, real_x);
+          inputData, input_sizes.wz, input_strides, n, c, real_y, real_x);
       auto i01 = upsample_get_value_bounded<T>(
-          inputData, input_sizes.xy, input_strides, n, c, real_y, real_x + 1);
+          inputData, input_sizes.wz, input_strides, n, c, real_y, real_x + 1);
       auto i10 = upsample_get_value_bounded<T>(
-          inputData, input_sizes.xy, input_strides, n, c, real_y + 1, real_x);
+          inputData, input_sizes.wz, input_strides, n, c, real_y + 1, real_x);
       auto i11 = upsample_get_value_bounded<T>(
           inputData,
-          input_sizes.xy,
+          input_sizes.wz,
           input_strides,
           n,
           c,
@@ -261,8 +261,8 @@ kernel void upsample_bilinear2d(
       auto i1_l = linear_interp(i10, i11, t_x);
       auto res = linear_interp(i0_l, i1_l, t_y);
       outputData
-          [n * output_strides.w + c * output_strides.z +
-           output_x * output_strides.x + output_y * output_strides.y] =
+          [n * output_strides.x + c * output_strides.y +
+           output_y * output_strides.z + output_x * output_strides.w] =
               static_cast<T>(res);
     }
   }
@@ -283,36 +283,36 @@ kernel void upsample_bilinear2d_aa(
     constant float2& scales [[buffer(6)]],
     constant bool& align_corners [[buffer(7)]],
     uint thread_index [[thread_position_in_grid]]) {
-  auto output_x = thread_index % output_sizes.x;
-  auto output_y = thread_index / output_sizes.x;
+  auto output_x = thread_index % output_sizes.w;
+  auto output_y = thread_index / output_sizes.w;
   (void)align_corners; // Align corners is unused for AA algorithm
   auto x_center = area_pixel_compute_source_index(
       scales.x, output_x, /*align_corners=*/false, /*cubic=*/false);
   auto y_center = area_pixel_compute_source_index(
       scales.y, output_y, /*align_corners=*/false, /*cubic=*/false);
   auto clamped_scales = max(1.0, scales);
   auto x_min = max(0L, long(floor(x_center - clamped_scales.x + 1)));
-  auto x_max = min(input_sizes.x, long(ceil(x_center + clamped_scales.x)));
+  auto x_max = min(input_sizes.w, long(ceil(x_center + clamped_scales.x)));
   auto y_min = max(0L, long(floor(y_center - clamped_scales.y + 1)));
-  auto y_max = min(input_sizes.y, long(ceil(y_center + clamped_scales.y)));
-  for (int n = 0; n < output_sizes.w; n++) {
-    for (int c = 0; c < output_sizes.z; c++) {
+  auto y_max = min(input_sizes.z, long(ceil(y_center + clamped_scales.y)));
+  for (int n = 0; n < output_sizes.x; n++) {
+    for (int c = 0; c < output_sizes.y; c++) {
       float res = 0.0;
       float ws = 0.0;
       constant auto* input =
-          inputData + n * input_strides.w + c * input_strides.z;
+          inputData + n * input_strides.x + c * input_strides.y;
       for (auto y = y_min; y < y_max; ++y) {
         auto dy = bilinear_functor((y - y_center) / clamped_scales.y);
         for (auto x = x_min; x < x_max; ++x) {
           auto dx = bilinear_functor((x - x_center) / clamped_scales.x);
-          auto val = input[x * input_strides.x + y * input_strides.y];
+          auto val = input[x * input_strides.w + y * input_strides.z];
           res += val * dx * dy;
           ws += dx * dy;
         }
       }
       outputData
-          [n * output_strides.w + c * output_strides.z +
-           output_x * output_strides.x + output_y * output_strides.y] =
+          [n * output_strides.x + c * output_strides.y +
+           output_y * output_strides.z + output_x * output_strides.w] =
               static_cast<T>(res / ws);
     }
   }
@@ -329,8 +329,8 @@ kernel void upsample_bicubic2d(
     constant float2& scales [[buffer(6)]],
     constant bool& align_corners [[buffer(7)]],
     uint thread_index [[thread_position_in_grid]]) {
-  auto output_x = thread_index % output_sizes.x;
-  auto output_y = thread_index / output_sizes.x;
+  auto output_x = thread_index % output_sizes.w;
+  auto output_y = thread_index / output_sizes.w;
   auto real_x = area_pixel_compute_source_index(
       scales.x, output_x, align_corners, /*cubic=*/true);
   int in_x = floor(real_x);
@@ -340,38 +340,38 @@ kernel void upsample_bicubic2d(
       scales.y, output_y, align_corners, /*cubic=*/true);
   int in_y = floor(real_y);
   auto t_y = real_y - in_y;
-  for (int n = 0; n < output_sizes.w; n++) {
-    for (int c = 0; c < output_sizes.z; c++) {
+  for (int n = 0; n < output_sizes.x; n++) {
+    for (int c = 0; c < output_sizes.y; c++) {
       float coefficients[4];
       for (int k = 0; k < 4; k++) {
         coefficients[k] = cubic_interp1d(
             upsample_get_value_bounded<T>(
                 inputData,
-                input_sizes.xy,
+                input_sizes.wz,
                 input_strides,
                 n,
                 c,
                 in_y - 1 + k,
                 in_x - 1),
             upsample_get_value_bounded<T>(
                 inputData,
-                input_sizes.xy,
+                input_sizes.wz,
                 input_strides,
                 n,
                 c,
                 in_y - 1 + k,
                 in_x + 0),
             upsample_get_value_bounded<T>(
                 inputData,
-                input_sizes.xy,
+                input_sizes.wz,
                 input_strides,
                 n,
                 c,
                 in_y - 1 + k,
                 in_x + 1),
             upsample_get_value_bounded<T>(
                 inputData,
-                input_sizes.xy,
+                input_sizes.wz,
                 input_strides,
                 n,
                 c,
@@ -386,8 +386,8 @@ kernel void upsample_bicubic2d(
           coefficients[3],
           t_y));
       outputData
-          [n * output_strides.w + c * output_strides.z +
-           output_x * output_strides.x + output_y * output_strides.y] = inp;
+          [n * output_strides.x + c * output_strides.y +
+           output_y * output_strides.z + output_x * output_strides.w] = inp;
     }
   }
 }
@@ -403,8 +403,8 @@ kernel void upsample_bicubic2d_backward(
     constant float2& scales [[buffer(6)]],
     constant bool& align_corners [[buffer(7)]],
     uint thread_index [[thread_position_in_grid]]) {
-  auto output_x = thread_index % output_sizes.x;
-  auto output_y = thread_index / output_sizes.x;
+  auto output_x = thread_index % output_sizes.w;
+  auto output_y = thread_index / output_sizes.w;
   auto real_x = area_pixel_compute_source_index<float>(
       scales.x, output_x, align_corners, /*cubic=*/true);
   int input_x = floor(real_x);
@@ -421,16 +421,16 @@ kernel void upsample_bicubic2d_backward(
   get_cubic_upsampling_coefficients(x_coeffs, t_x);
   get_cubic_upsampling_coefficients(y_coeffs, t_y);
 
-  for (int n = 0; n < output_sizes.w; n++) {
-    for (int c = 0; c < output_sizes.z; ++c) {
+  for (int n = 0; n < output_sizes.x; n++) {
+    for (int c = 0; c < output_sizes.y; ++c) {
       auto out_value = gradOutputData
-          [n * output_strides.w + c * output_strides.z +
-           output_x * output_strides.x + output_y * output_strides.y];
+          [n * output_strides.x + c * output_strides.y +
+           output_y * output_strides.z + output_x * output_strides.w];
       for (int i = 0; i < 4; i++) {
         for (int j = 0; j < 4; j++) {
           upsample_increment_value_bounded<T>(
               gradInputData,
-              input_sizes.xy,
+              input_sizes.wz,
               input_strides,
               n,
               c,
@@ -478,7 +478,7 @@ kernel void upsample_bicubic2d_backward(
       constant ulong3 & output_strides [[buffer(3)]],             \
       constant long3 & input_sizes [[buffer(4)]],                 \
       constant long3 & output_sizes [[buffer(5)]],                \
-      constant float& scale [[buffer(6)]],                        \
+      constant float2 & scales [[buffer(6)]],                     \
       constant bool& align_corners [[buffer(7)]],                 \
       uint thread_index [[thread_position_in_grid]])
 

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

@@ -263,62 +263,28 @@ static void upsample_kernel_out_template(const Tensor& input,
     return;
   }
   std::array<float, 2> scales = {
-      area_pixel_compute_scale<float>(input.size(3), output.size(3), align_corners, scale_w_opt),
+      area_pixel_compute_scale<float>(input.size(-1), output.size(-1), align_corners, scale_w_opt),
       area_pixel_compute_scale<float>(input.size(2), output.size(2), align_corners, scale_h_opt)};
   auto upsamplePSO = lib.getPipelineStateForFunc(fmt::format("upsample_{}_{}", name, scalarToMetalTypeString(input)));
   auto stream = getCurrentMPSStream();
   dispatch_sync_with_rethrow(stream->queue(), ^() {
     @autoreleasepool {
-      std::array<int64_t, 4> output_strides = {output.stride(3), output.stride(2), output.stride(1), output.stride(0)};
-      std::array<int64_t, 4> output_sizes = {output.size(3), output.size(2), output.size(1), output.size(0)};
-      std::array<int64_t, 4> input_sizes = {input.size(3), input.size(2), input.size(1), input.size(0)};
-      std::array<int64_t, 4> input_strides = {input.stride(3), input.stride(2), input.stride(1), input.stride(0)};
       auto computeEncoder = stream->commandEncoder();
       [computeEncoder setComputePipelineState:upsamplePSO];
       mtl_setArgs(computeEncoder,
                   input,
                   output,
-                  input_strides,
-                  output_strides,
-                  input_sizes,
-                  output_sizes,
+                  input.strides(),
+                  output.strides(),
+                  input.sizes(),
+                  output.sizes(),
                   scales,
                   align_corners);
-      mtl_dispatch1DJob(computeEncoder, upsamplePSO, output_size[0] * output_size[1]);
-    }
-  });
-}
-
-static void upsample_kernel_out_template(const Tensor& input,
-                                         IntArrayRef output_size,
-                                         bool align_corners,
-                                         std::optional<double> scale_opt,
-                                         const Tensor& output,
-                                         const std::string name) {
-  if (output.numel() == 0) {
-    return;
-  }
-  float scale = area_pixel_compute_scale<float>(input.size(2), output.size(2), align_corners, scale_opt);
-  auto upsamplePSO = lib.getPipelineStateForFunc(fmt::format("upsample_{}_{}", name, scalarToMetalTypeString(input)));
-  auto stream = getCurrentMPSStream();
-  dispatch_sync_with_rethrow(stream->queue(), ^() {
-    @autoreleasepool {
-      std::array<int64_t, 3> output_strides = {output.stride(2), output.stride(1), output.stride(0)};
-      std::array<int64_t, 3> output_sizes = {output.size(2), output.size(1), output.size(0)};
-      std::array<int64_t, 3> input_sizes = {input.size(2), input.size(1), input.size(0)};
-      std::array<int64_t, 3> input_strides = {input.stride(2), input.stride(1), input.stride(0)};
-      auto computeEncoder = stream->commandEncoder();
-      [computeEncoder setComputePipelineState:upsamplePSO];
-      mtl_setArgs(computeEncoder,
-                  input,
-                  output,
-                  input_strides,
-                  output_strides,
-                  input_sizes,
-                  output_sizes,
-                  scale,
-                  align_corners);
-      mtl_dispatch1DJob(computeEncoder, upsamplePSO, output_size[0]);
+      if (output.ndimension() == 4) {
+        mtl_dispatch1DJob(computeEncoder, upsamplePSO, output_size[0] * output_size[1]);
+      } else {
+        mtl_dispatch1DJob(computeEncoder, upsamplePSO, output_size[0]);
+      }
     }
   });
 }
@@ -343,23 +309,15 @@ static void upsample_kernel_backward_out_template(const Tensor& grad_input,
   auto stream = getCurrentMPSStream();
   dispatch_sync_with_rethrow(stream->queue(), ^() {
     @autoreleasepool {
-      std::array<int64_t, 4> output_strides = {
-          grad_output.stride(3), grad_output.stride(2), grad_output.stride(1), grad_output.stride(0)};
-      std::array<int64_t, 4> output_sizes = {
-          grad_output.size(3), grad_output.size(2), grad_output.size(1), grad_output.size(0)};
-      std::array<int64_t, 4> input_sizes = {
-          grad_input.size(3), grad_input.size(2), grad_input.size(1), grad_input.size(0)};
-      std::array<int64_t, 4> input_strides = {
-          grad_input.stride(3), grad_input.stride(2), grad_input.stride(1), grad_input.stride(0)};
       auto computeEncoder = stream->commandEncoder();
       [computeEncoder setComputePipelineState:upsamplePSO];
       mtl_setArgs(computeEncoder,
                   grad_input,
                   grad_output,
-                  input_strides,
-                  output_strides,
-                  input_sizes,
-                  output_sizes,
+                  grad_input.strides(),
+                  grad_output.strides(),
+                  grad_input.sizes(),
+                  grad_output.sizes(),
                   scales,
                   align_corners);
       mtl_dispatch1DJob(computeEncoder, upsamplePSO, output_size[0] * output_size[1]);
@@ -439,7 +397,7 @@ static void upsample_kernel_backward_out_template(const Tensor& grad_input,
 
 TORCH_IMPL_FUNC(upsample_linear1d_out_mps)
 (const Tensor& input, IntArrayRef output_size, bool align_corners, std::optional<double> scale, const Tensor& output) {
-  mps::upsample_kernel_out_template(input, output_size, align_corners, scale, output, "linear1d");
+  mps::upsample_kernel_out_template(input, output_size, align_corners, scale, scale, output, "linear1d");
 }
 
 TORCH_IMPL_FUNC(upsample_linear1d_backward_out_mps)