Enable fix for UpsampleNearest functions in Pytorch

[mikaylagawarecki]

🚀 Feature

In response to pytorch/pytorch#62396, we opened pytorch/pytorch#97869 and pytorch/pytorch#97868 in the PyTorch repo

Motivation

Upsample with mode nearest in Pytorch has a bug where the output size is rounded incorrectly using floor() rather than round() which makes it inconsistent with other frameworks like openCV, tensorflow etc. We are fixing this by adding a kwarg per the approach in pytorch/pytorch#64983 (comment). This is breaking xla CI, we need to fix this on XLA side as well in order to land the PR

Pitch

The specific code is here:

xla/torch_xla/csrc/aten_xla_type.cpp

Lines 177 to 188 in 7d9d43e

	std::vector<int64_t> GetOutputSizeWithScale(
	absl::Span<const int64_t> input_size, const c10::optional<double> scales_h,
	const c10::optional<double> scales_w,
	const std::vector<int64_t>& output_size) {
	XLA_CHECK(scales_h);
	XLA_CHECK(scales_w);
	// Calculate the output size from input_shape and scale_factors
	XLA_CHECK_EQ(input_size.size(), 4);
	int64_t output_h = input_size[2] * (*scales_h);
	int64_t output_w = input_size[3] * (*scales_w);
	return {output_h, output_w};
	}

which should compute output size with either round or implicit floor
int64_t output_h = 0.5 + input_size[2] * (*scales_h);
or
int64_t output_h = input_size[2] * (*scales_h);
in the UpsampleNearest variants depending on new flag round_with_scale_factor

However, this flag can't be exposed in, for example, upsample_bilinear2d function:

xla/torch_xla/csrc/aten_xla_type.cpp

Line 3112 in 7d9d43e

at::Tensor XLANativeFunctions::upsample_bilinear2d(

as output size was already correctly computed before dispatching to XLA.

Alternatives

Additional context

cc @vfdev-5 @JackCaoG

[JackCaoG]

Thanks @mikaylagawarecki for raising the issue, @wonjoolee95 do you have cycle to pick this us?

[wonjoo-wj]

Sure thing, I'll get a fix out by next week. I'll let you know once that's ready for merge, @mikaylagawarecki.

[wonjoo-wj]

Ok, spent some time looking onto this and it's a bit tricky. The newly added flag round_with_scale_factor is added to the upsample_bilinear2d.vecop, but this op never seems to be dispatched to XLA due to fuctionalization.

For example, running the test TORCH_SHOW_DISPATCH_TRACE=1 python test/test_nn.py TestNNDeviceTypeXLA.test_upsampling_output_size_round_with_scale_factor_True_mode_bilinear_xla, the test fails due to this new flag (failure shown at upstream PR CI: https://github.com/pytorch/pytorch/actions/runs/4806459855/jobs/8554437961?pr=97869). In the dispatch trace, we can see that the op upsample_bilinear2d.vec does not get dispatched to XLA:

 [call] op=[aten::upsample_bilinear2d.vec], key=[Functionalize]
  [call] op=[aten::upsample_bilinear2d], key=[AutogradXLA]
   [redispatch] op=[aten::upsample_bilinear2d], key=[Functionalize]
    [callBoxed] op=[aten::upsample_bilinear2d], key=[XLA]

With functionalization, this op seems to be decomposed into the aten::upsample_bilinear2d op. And this op doesn't have this new round_with_scale_factor flag. @bdhirsh, do you have any pointers on how we can get by this issue?

And @mikaylagawarecki, is this urgent? If so, we can disable this specific test on XLA for now until we have a more permanent solution for now, so we can unblock your PR. Let me know, thanks!

cc @JackCaoG, just fyi

[vfdev-5]

@wonjoolee95, I agree with your findings, yes,aten::upsample_bilinear2d.vec is composite, here is the code:

https://github.com/pytorch/pytorch/blob/52a36a98d9425479f62b6e2d1a59e434b85f7f7e/aten/src/ATen/native/UpSampleNearest2d.cpp#L150-L158

This function signature has output_size arg as optional:

Tensor upsample_nearest2d(
    const Tensor& input,
    at::OptionalIntArrayRef output_size,
    c10::optional<ArrayRef<double>> scale_factors)

so, it calls inside another method to compute actual output size if scale_factors are provided:

auto osize = compute_output_size(input.sizes(), output_size, scale_factors);

then it calls another publicly available "upsample_nearest2d" where output size is not optional anymore:

TORCH_META_FUNC(upsample_nearest2d) (
    const Tensor& input, IntArrayRef output_size, c10::optional<double> scales_h, c10::optional<double> scales_w
) {
  auto full_output_size = native::upsample_2d_common_check(input.sizes(), output_size);

  // Allow for empty batch size but not other dimensions
  TORCH_CHECK(
      input.numel() != 0 || c10::multiply_integers(input.sizes().begin() + 1, input.sizes().end()),
      "Non-empty 4D data tensor expected but got a tensor with sizes ",
      input.sizes());

  set_output_raw_strided(0, full_output_size, {}, input.options().memory_format(input.suggest_memory_format()));
}

Latter function is then dispatched to XLANativeFunctions::upsample_nearest2d.

In XLA upsample_nearest2d there is an additional check that is not present in CPU or CUDA implementations:

xla/torch_xla/csrc/aten_xla_type.cpp

Lines 3179 to 3183 in 7d9d43e

	if (!output_size.empty()) {
	XLA_CHECK(scaled_output_size.at(0) == output_size.at(0) &&
	scaled_output_size.at(1) == output_size.at(1))
	<< "Inferred output size and output_size from upstream are different";
	}

As far as I can see from the code, XLA upsample_nearest2d/bilinear2d finally does the upsampling/downsampling using input size and output size params only vs CPU/CUDA implementations where they keep input size, output size and also user scales to compute output index -> input index mapping:

https://github.com/pytorch/pytorch/blob/52a36a98d9425479f62b6e2d1a59e434b85f7f7e/aten/src/ATen/native/UpSample.h#L263-L279

https://github.com/pytorch/pytorch/blob/52a36a98d9425479f62b6e2d1a59e434b85f7f7e/aten/src/ATen/native/UpSample.h#L250-L260

(this may potentially lead to a certain output inconsistency between CPU vs XLA for floating-point scales like s=1.2345 or s=2.3456)

As for a "fix" in XLA, we can add round_with_scale_factor arg into GetOutputSizeWithScale and loose a bit the assert such that we check for both possible outputs:

    scaled_output_size = GetOutputSizeWithScale(input_dims, scales_h, scales_w,
                                                scaled_output_size, /*round_with_scale_factor=*/false);
    if (!output_size.empty()) {
      if (! (scaled_output_size.at(0) == output_size.at(0) && scaled_output_size.at(1) == output_size.at(1)) ) {
          scaled_output_size = GetOutputSizeWithScale(input_dims, scales_h, scales_w,
                                                scaled_output_size, /*round_with_scale_factor=*/true);
          XLA_CHECK(scaled_output_size.at(0) == output_size.at(0) && scaled_output_size.at(1) == output_size.at(1))
                << "Inferred output size and output_size from upstream are different";
      }
    }

What do you think ?

And @mikaylagawarecki, is this urgent? If so, we can disable this specific test on XLA for now until we have a more permanent solution for now, so we can unblock your PR. Let me know, thanks!

I think we can also go this way and disable this specific XLA test until a better solution could be found.

[bdhirsh]

@wonjoolee95 for what it's worth - functionalization isn't actually the culprit in that decomposition getting run 😛. upsample_nearest2d.vec has a CompositeImplicitAutograd decomposition, and CompositeImplicitAutograd always run before hitting XLA (unless XLA overrides the AutogradXLA dispatch key, which it does for ~3-4 ops).

[wonjoo-wj]

Thanks for the inputs, @vfdev-5 and @bdhirsh! That makes a lot of sense.

@vfdev-5, the proposed fix makes sense. Thanks for catching the issue and the detailed write-up. Let me get a working PR out and run the CI to verify the change.

[wonjoo-wj]

@mikaylagawarecki and @vfdev-5, the fix PR on our side is ready -- #4977. Could you update the XLA CI pin on your PR to 7eb437d76db298d3bfc28f5ff61a3bb748b721f7 and retrigger the CI? Once both CIs are green, we can coordinate a merge.

[mikaylagawarecki]

Thank you so much for the fix @wonjoolee95! So sorry I'm still resolving some internal errors with the initial PR, will ping here when we're ready to coordinate the merge!

[JackCaoG]

@wonjoolee95 can we close this one?

[wonjoo-wj]

The upstream PyTorch PRs got closed, so closing ours as well.