diff --git a/kernels/portable/cpu/op_upsample_bilinear2d_aa.cpp b/kernels/portable/cpu/op_upsample_bilinear2d_aa.cpp index 1553471e79a..b5404d74697 100644 --- a/kernels/portable/cpu/op_upsample_bilinear2d_aa.cpp +++ b/kernels/portable/cpu/op_upsample_bilinear2d_aa.cpp @@ -30,17 +30,19 @@ inline T bilinear_aa_filter(T x) { : static_cast(0.0); } -// Compute anti-aliasing weights exactly matching PyTorch's algorithm template -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 +AAKernelParams +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(0)); const int64_t xmax = std::min( static_cast(center + support + static_cast(0.5)), input_size); - - *num_contributors = std::min(xmax - xmin, static_cast(4)); + const T invscale = (scale >= static_cast(1.0)) + ? (static_cast(1.0) / scale) + : static_cast(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(0), std::min(static_cast(center), input_size - 1)); - weights[0] = static_cast(1.0); - // Clear unused weight slots - for (int64_t j = 1; j < 4; ++j) { - weights[j] = static_cast(0.0); - } - return; + return { + fallback_index, + fallback_index + 1, + fallback_index, + center, + invscale, + static_cast(1.0)}; } - // PyTorch's weight computation T total_weight = static_cast(0.0); - const T invscale = (scale >= static_cast(1.0)) - ? (static_cast(1.0) / scale) - : static_cast(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(j) + static_cast(xmin) - center + - static_cast(0.5)) * - invscale; - T weight = bilinear_aa_filter(arg); - indices[j] = x; - weights[j] = weight; - total_weight += weight; + for (int64_t x = xmin; x < xmax; ++x) { + total_weight += bilinear_aa_filter( + (static_cast(x) - center + static_cast(0.5)) * invscale); } - // Normalize weights to sum to 1 (PyTorch does this) - if (total_weight > static_cast(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(1.0) / static_cast(*num_contributors); - for (int64_t j = 0; j < *num_contributors; ++j) { - weights[j] = equal_weight; - } + if (total_weight <= static_cast(0.0)) { + const int64_t fallback_index = std::max( + static_cast(0), + std::min(static_cast(center), input_size - 1)); + return { + fallback_index, + fallback_index + 1, + fallback_index, + center, + invscale, + static_cast(1.0)}; } - // Clear unused weight slots - for (int64_t j = *num_contributors; j < 4; ++j) { - weights[j] = static_cast(0.0); + return {xmin, xmax, -1, center, invscale, total_weight}; +} + +template +inline T compute_normalized_aa_weight( + const AAKernelParams& params, + int64_t input_idx) { + if (params.fallback_index >= 0) { + return (input_idx == params.fallback_index) ? static_cast(1.0) + : static_cast(0.0); } + return bilinear_aa_filter( + (static_cast(input_idx) - params.center + static_cast(0.5)) * + params.invscale) / + params.total_weight; } template @@ -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( - 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( - 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(oh, scale_h, in.size(2)); + for (int64_t ow = 0; ow < out.size(3); ++ow) { + const auto w_params = + compute_aa_kernel_params(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( - oh, scale_h, in.size(2), h_indices, h_weights, &h_num_contributors); + const auto h_params = + compute_aa_kernel_params(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( - ow, - scale_w, - in.size(3), - w_indices, - w_weights, - &w_num_contributors); + const auto w_params = + compute_aa_kernel_params(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; diff --git a/kernels/portable/test/op_upsample_bilinear2d_aa_test.py b/kernels/portable/test/op_upsample_bilinear2d_aa_test.py index c6e09af3b5c..8d0c94258a9 100644 --- a/kernels/portable/test/op_upsample_bilinear2d_aa_test.py +++ b/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.""" diff --git a/kernels/test/op_upsample_bilinear2d_aa_test.cpp b/kernels/test/op_upsample_bilinear2d_aa_test.cpp index c0cc69f35d9..896dab3061d 100644 --- a/kernels/test/op_upsample_bilinear2d_aa_test.cpp +++ b/kernels/test/op_upsample_bilinear2d_aa_test.cpp @@ -626,6 +626,69 @@ TEST_F(OpUpsampleBilinear2dAAOutTest, TestPrecisionConsistency) { } } +TEST_F(OpUpsampleBilinear2dAAOutTest, Simple8x1To3x8) { + TensorFactory 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 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(); + 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 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 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(); + 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 tf;