Fix _upsample_bilinear2d_aa portable kernel correctness (pytorch#20733)

### Summary
Fixes pytorch#13969 

Fix _upsample_bilinear2d_aa portable kernel correctness for anti-aliased
downsampling by removing the 4 tap contributor limit and matching
PyTorch’s full support window.

Added portable tests and corrected tolerance in test file

cc @larryliu0820 @manuelcandales

Author: arkadip-maitra

kernels/portable/cpu/op_upsample_bilinear2d_aa.cpp

@@ -30,17 +30,19 @@ inline T bilinear_aa_filter(T x) {
                                    : static_cast<T>(0.0);
 }
 
-// Compute anti-aliasing weights exactly matching PyTorch's algorithm
 template <typename T>
-void compute_aa_weights_for_pixel(
-    int64_t output_idx,
-    T scale,
-    int64_t input_size,
-    int64_t* indices,
-    T* weights,
-    int64_t* num_contributors) {
-  // Use the provided scale directly instead of recalculating
+struct AAKernelParams {
+  int64_t xmin;
+  int64_t xmax;
+  int64_t fallback_index;
+  T center;
+  T invscale;
+  T total_weight;
+};
 
+template <typename T>
+AAKernelParams<T>
+compute_aa_kernel_params(int64_t output_idx, T scale, int64_t input_size) {
   // PyTorch's center calculation for anti-aliasing
   // Always uses scale * (i + 0.5) for anti-aliasing, regardless of
   // align_corners
@@ -58,58 +60,58 @@ void compute_aa_weights_for_pixel(
       static_cast<int64_t>(0));
   const int64_t xmax = std::min(
       static_cast<int64_t>(center + support + static_cast<T>(0.5)), input_size);
-
-  *num_contributors = std::min(xmax - xmin, static_cast<int64_t>(4));
+  const T invscale = (scale >= static_cast<T>(1.0))
+      ? (static_cast<T>(1.0) / scale)
+      : static_cast<T>(1.0);
 
   // Ensure we have at least one contributor
-  if (*num_contributors <= 0) {
-    *num_contributors = 1;
-    indices[0] = std::max(
+  if (xmax <= xmin) {
+    const int64_t fallback_index = std::max(
         static_cast<int64_t>(0),
         std::min(static_cast<int64_t>(center), input_size - 1));
-    weights[0] = static_cast<T>(1.0);
-    // Clear unused weight slots
-    for (int64_t j = 1; j < 4; ++j) {
-      weights[j] = static_cast<T>(0.0);
-    }
-    return;
+    return {
+        fallback_index,
+        fallback_index + 1,
+        fallback_index,
+        center,
+        invscale,
+        static_cast<T>(1.0)};
   }
 
-  // PyTorch's weight computation
   T total_weight = static_cast<T>(0.0);
-  const T invscale = (scale >= static_cast<T>(1.0))
-      ? (static_cast<T>(1.0) / scale)
-      : static_cast<T>(1.0);
-
-  for (int64_t j = 0; j < *num_contributors; ++j) {
-    int64_t x = xmin + j;
-    // PyTorch's exact weight formula: (j + xmin - center + 0.5) * invscale
-    T arg = (static_cast<T>(j) + static_cast<T>(xmin) - center +
-             static_cast<T>(0.5)) *
-        invscale;
-    T weight = bilinear_aa_filter<T>(arg);
-    indices[j] = x;
-    weights[j] = weight;
-    total_weight += weight;
+  for (int64_t x = xmin; x < xmax; ++x) {
+    total_weight += bilinear_aa_filter<T>(
+        (static_cast<T>(x) - center + static_cast<T>(0.5)) * invscale);
   }
 
-  // Normalize weights to sum to 1 (PyTorch does this)
-  if (total_weight > static_cast<T>(0.0)) {
-    for (int64_t j = 0; j < *num_contributors; ++j) {
-      weights[j] /= total_weight;
-    }
-  } else {
-    // Fallback: if total weight is 0, set equal weights
-    T equal_weight = static_cast<T>(1.0) / static_cast<T>(*num_contributors);
-    for (int64_t j = 0; j < *num_contributors; ++j) {
-      weights[j] = equal_weight;
-    }
+  if (total_weight <= static_cast<T>(0.0)) {
+    const int64_t fallback_index = std::max(
+        static_cast<int64_t>(0),
+        std::min(static_cast<int64_t>(center), input_size - 1));
+    return {
+        fallback_index,
+        fallback_index + 1,
+        fallback_index,
+        center,
+        invscale,
+        static_cast<T>(1.0)};
   }
 
-  // Clear unused weight slots
-  for (int64_t j = *num_contributors; j < 4; ++j) {
-    weights[j] = static_cast<T>(0.0);
+  return {xmin, xmax, -1, center, invscale, total_weight};
+}
+
+template <typename T>
+inline T compute_normalized_aa_weight(
+    const AAKernelParams<T>& params,
+    int64_t input_idx) {
+  if (params.fallback_index >= 0) {
+    return (input_idx == params.fallback_index) ? static_cast<T>(1.0)
+                                                : static_cast<T>(0.0);
   }
+  return bilinear_aa_filter<T>(
+             (static_cast<T>(input_idx) - params.center + static_cast<T>(0.5)) *
+             params.invscale) /
+      params.total_weight;
 }
 
 template <typename CTYPE>
@@ -129,48 +131,28 @@ void upsample_bilinear2d_aa_kernel_impl(
   if (is_nchw) {
     // NCHW layout
     for (int64_t n = 0; n < out.size(0); ++n) {
-      for (int64_t c = 0; c < out.size(1); ++c) {
-        const auto in_plane =
-            in_data + (n * in.size(1) + c) * in.size(2) * in.size(3);
-        auto out_plane =
-            out_data + (n * out.size(1) + c) * out.size(2) * out.size(3);
-
-        for (int64_t oh = 0; oh < out.size(2); ++oh) {
-          // Compute height weights for this output row
-          int64_t h_indices[4];
-          float h_weights[4];
-          int64_t h_num_contributors;
-          compute_aa_weights_for_pixel<float>(
-              oh,
-              scale_h,
-              in.size(2),
-              h_indices,
-              h_weights,
-              &h_num_contributors);
-
-          for (int64_t ow = 0; ow < out.size(3); ++ow) {
-            // Compute width weights for this output column
-            int64_t w_indices[4];
-            float w_weights[4];
-            int64_t w_num_contributors;
-            compute_aa_weights_for_pixel<float>(
-                ow,
-                scale_w,
-                in.size(3),
-                w_indices,
-                w_weights,
-                &w_num_contributors);
+      for (int64_t oh = 0; oh < out.size(2); ++oh) {
+        const auto h_params =
+            compute_aa_kernel_params<float>(oh, scale_h, in.size(2));
 
+        for (int64_t ow = 0; ow < out.size(3); ++ow) {
+          const auto w_params =
+              compute_aa_kernel_params<float>(ow, scale_w, in.size(3));
+
+          for (int64_t c = 0; c < out.size(1); ++c) {
+            const auto in_plane =
+                in_data + (n * in.size(1) + c) * in.size(2) * in.size(3);
+            auto out_plane =
+                out_data + (n * out.size(1) + c) * out.size(2) * out.size(3);
             CTYPE value = 0;
 
             // Apply anti-aliased interpolation
-            for (int64_t ih_idx = 0; ih_idx < h_num_contributors; ++ih_idx) {
-              int64_t ih = h_indices[ih_idx];
-              float h_weight = h_weights[ih_idx];
+            for (int64_t ih = h_params.xmin; ih < h_params.xmax; ++ih) {
+              const float h_weight = compute_normalized_aa_weight(h_params, ih);
 
-              for (int64_t iw_idx = 0; iw_idx < w_num_contributors; ++iw_idx) {
-                int64_t iw = w_indices[iw_idx];
-                float w_weight = w_weights[iw_idx];
+              for (int64_t iw = w_params.xmin; iw < w_params.xmax; ++iw) {
+                const float w_weight =
+                    compute_normalized_aa_weight(w_params, iw);
 
                 value += in_plane[ih * in.size(3) + iw] * h_weight * w_weight;
               }
@@ -188,37 +170,23 @@ void upsample_bilinear2d_aa_kernel_impl(
       auto out_batch = out_data + n * out.size(1) * out.size(2) * out.size(3);
 
       for (int64_t oh = 0; oh < out.size(2); ++oh) {
-        // Compute height weights for this output row
-        int64_t h_indices[4];
-        float h_weights[4];
-        int64_t h_num_contributors;
-        compute_aa_weights_for_pixel<float>(
-            oh, scale_h, in.size(2), h_indices, h_weights, &h_num_contributors);
+        const auto h_params =
+            compute_aa_kernel_params<float>(oh, scale_h, in.size(2));
 
         for (int64_t ow = 0; ow < out.size(3); ++ow) {
-          // Compute width weights for this output column
-          int64_t w_indices[4];
-          float w_weights[4];
-          int64_t w_num_contributors;
-          compute_aa_weights_for_pixel<float>(
-              ow,
-              scale_w,
-              in.size(3),
-              w_indices,
-              w_weights,
-              &w_num_contributors);
+          const auto w_params =
+              compute_aa_kernel_params<float>(ow, scale_w, in.size(3));
 
           for (int64_t c = 0; c < out.size(1); ++c) {
             CTYPE value = 0;
 
             // Apply anti-aliased interpolation
-            for (int64_t ih_idx = 0; ih_idx < h_num_contributors; ++ih_idx) {
-              int64_t ih = h_indices[ih_idx];
-              float h_weight = h_weights[ih_idx];
+            for (int64_t ih = h_params.xmin; ih < h_params.xmax; ++ih) {
+              const float h_weight = compute_normalized_aa_weight(h_params, ih);
 
-              for (int64_t iw_idx = 0; iw_idx < w_num_contributors; ++iw_idx) {
-                int64_t iw = w_indices[iw_idx];
-                float w_weight = w_weights[iw_idx];
+              for (int64_t iw = w_params.xmin; iw < w_params.xmax; ++iw) {
+                const float w_weight =
+                    compute_normalized_aa_weight(w_params, iw);
 
                 value += in_batch[(ih * in.size(3) + iw) * in.size(1) + c] *
                     h_weight * w_weight;

kernels/portable/test/op_upsample_bilinear2d_aa_test.py

@@ -161,10 +161,7 @@ def test_upsample_bilinear2d_aa_aggressive_downsampling(self):
             input_tensor,
             output_size=(2, 2),
             align_corners=False,
-            # Aggressive 4x downsampling magnifies the separable-vs-direct
-            # interpolation differences between ExecuTorch and ATen; observed
-            # max abs error reaches ~0.6 for typical N(0,1) inputs.
-            atol=1.0,
+            atol=1e-3,
         )
 
     def test_upsample_bilinear2d_aa_asymmetric_downsampling(self):
@@ -174,7 +171,7 @@ def test_upsample_bilinear2d_aa_asymmetric_downsampling(self):
             input_tensor,
             output_size=(4, 4),  # 3x downsample in H, 2x in W
             align_corners=False,
-            atol=0.25,  # Relaxed tolerance due to implementation differences in separable vs direct interpolation
+            atol=1e-3,
         )
 
     def test_upsample_bilinear2d_aa_align_corners_upsampling(self):
@@ -194,7 +191,7 @@ def test_upsample_bilinear2d_aa_align_corners_downsampling(self):
             input_tensor,
             output_size=(4, 4),
             align_corners=True,
-            atol=0.25,  # Relaxed tolerance due to implementation differences in separable vs direct interpolation
+            atol=1e-3,
         )
 
     def test_upsample_bilinear2d_aa_batched(self):
@@ -245,19 +242,12 @@ def test_upsample_bilinear2d_aa_known_values_correctness(self):
         """Test against known correct output values to catch regressions."""
         # This test case is adapted from ATen's test suite
         input_tensor = torch.arange(3 * 8 * 8, dtype=torch.float).reshape(1, 3, 8, 8)
-
-        # Test with a known downsampling case
-        try:
-            self.run_upsample_aa_test(
-                input_tensor,
-                output_size=(2, 2),
-                align_corners=False,
-                atol=1e-2,  # Slightly relaxed for implementation differences
-            )
-            # The test should pass if our implementation is close to ATen
-        except AssertionError as e:
-            # Log the difference for debugging but don't fail the test during development
-            print(f"Known values test difference (expected during development): {e}")
+        self.run_upsample_aa_test(
+            input_tensor,
+            output_size=(2, 2),
+            align_corners=False,
+            atol=1e-3,
+        )
 
     def test_upsample_bilinear2d_aa_various_dtypes(self):
         """Test with various data types."""

kernels/test/op_upsample_bilinear2d_aa_test.cpp

@@ -626,6 +626,69 @@ TEST_F(OpUpsampleBilinear2dAAOutTest, TestPrecisionConsistency) {
   }
 }
 
+TEST_F(OpUpsampleBilinear2dAAOutTest, Simple8x1To3x8) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor input = tf.make(
+      {1, 1, 8, 1},
+      {-98.5, 49.875, 17.125, -46.5, 10.625, -95.875, -3.875, -4.625});
+  Tensor out = tf.zeros({1, 1, 3, 8});
+
+  int64_t output_size_data[2] = {3, 8};
+  ArrayRef<int64_t> output_size(output_size_data, 2);
+
+  op_upsample_bilinear2d_aa_out(
+      input,
+      output_size,
+      /*align_corners=*/false,
+      std::nullopt,
+      std::nullopt,
+      out);
+
+  auto expected = tf.zeros({1, 1, 3, 8});
+  auto expected_data = expected.mutable_data_ptr<float>();
+  const float expected_rows[] = {-8.440787f, -23.13393f, -24.736841f};
+  for (int row = 0; row < 3; ++row) {
+    for (int col = 0; col < 8; ++col) {
+      expected_data[row * 8 + col] = expected_rows[row];
+    }
+  }
+
+  EXPECT_TENSOR_CLOSE_WITH_TOL(out, expected, 0, 1e-4);
+}
+
+TEST_F(OpUpsampleBilinear2dAAOutTest, Simple8x1To3x8ChannelsLast) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor input = tf.make_channels_last(
+      {1, 1, 8, 1},
+      {-98.5, 49.875, 17.125, -46.5, 10.625, -95.875, -3.875, -4.625});
+  Tensor out = tf.zeros_channels_last({1, 1, 3, 8});
+
+  int64_t output_size_data[2] = {3, 8};
+  ArrayRef<int64_t> output_size(output_size_data, 2);
+
+  op_upsample_bilinear2d_aa_out(
+      input,
+      output_size,
+      /*align_corners=*/false,
+      std::nullopt,
+      std::nullopt,
+      out);
+
+  auto expected_contiguous = tf.zeros({1, 1, 3, 8});
+  auto expected_data = expected_contiguous.mutable_data_ptr<float>();
+  const float expected_rows[] = {-8.440787f, -23.13393f, -24.736841f};
+  for (int row = 0; row < 3; ++row) {
+    for (int col = 0; col < 8; ++col) {
+      expected_data[row * 8 + col] = expected_rows[row];
+    }
+  }
+  const auto expected = tf.channels_last_like(expected_contiguous);
+
+  EXPECT_TENSOR_CLOSE_WITH_TOL(out, expected, 0, 1e-4);
+}
+
 TEST_F(OpUpsampleBilinear2dAAOutTest, TestSpecificInputCase) {
   TensorFactory<ScalarType::Float> tf;