Skip to content

Enable lowering for upsample_bilinear2d with scale factor #4464

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jan 18, 2023
Merged

Enable lowering for upsample_bilinear2d with scale factor #4464

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d219896
Enable lowering for upsample_bilinear2d with scale factor
lsy323 Jan 17, 2023
0b47f41
fix linter
lsy323 Jan 17, 2023
5226826
Merge branch 'pytorch:master' into merge-upsample-bilinear2d-with-scale
lsy323 Jan 17, 2023
9c2807c
add shape validation
lsy323 Jan 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions test/cpp/test_aten_xla_tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4311,6 +4311,47 @@ TEST_F(AtenXlaTensorTest, TestUpsampleBilinear2D) {
cpp_test::GetIgnoredCounters());
}

TEST_F(AtenXlaTensorTest, TestUpsampleBilinear2DWithScale) {
struct ImageInfo {
int batch_size;
int h;
int w;
int chans;
double scale_h;
double scale_w;
};

/* clang-format off */
std::vector<ImageInfo> inputs = {
{/*batch_size=*/2, /*h=*/5, /*w=*/5, /*chans=*/2, /*scale_h*/8.0/5, /*scale_w*/8.0/5},
{/*batch_size=*/2, /*h=*/1335, /*w=*/1335, /*chans=*/3, /*scale_h*/255.0/1335, /*scale_w*/255.0/1335},
{/*batch_size=*/2, /*h=*/255, /*w=*/255, /*chans=*/3, /*scale_h*/1335.0/255, /*scale_w*/1335.0/255},
{/*batch_size=*/2, /*h=*/254, /*w=*/243, /*chans=*/3, /*scale_h*/784.0/254, /*scale_w*/214.0/243}
};
/* clang-format on */

for (const auto& img_info : inputs) {
for (bool align_corners : {true, false}) {
torch::Tensor input = torch::rand(
{img_info.batch_size, img_info.chans, img_info.h, img_info.w},
torch::TensorOptions(torch::kFloat));
ForEachDevice([&](const torch::Device& device) {
torch::Tensor xla_input = CopyToDevice(input, device);
torch::Tensor result = torch::upsample_bilinear2d(
input, c10::nullopt, align_corners,
at::ArrayRef<double>{img_info.scale_h, img_info.scale_w});
torch::Tensor xla_result = torch::upsample_bilinear2d(
xla_input, c10::nullopt, align_corners,
at::ArrayRef<double>{img_info.scale_h, img_info.scale_w});
AllClose(result, xla_result, /*rtol=*/1e-4, /*atol=*/1e-4);
});
}
}
ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
ExpectCounterChanged("xla::upsample_bilinear2d",
cpp_test::GetIgnoredCounters());
}

TEST_F(AtenXlaTensorTest, TestUpsampleBilinear2DBackward) {
int batch_size = 2;
int h = 5;
Expand Down
19 changes: 12 additions & 7 deletions torch_xla/csrc/aten_xla_type.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2904,16 +2904,21 @@ at::Tensor XLANativeFunctions::upsample_bilinear2d(
c10::optional<double> scales_h, c10::optional<double> scales_w) {
TORCH_LAZY_FN_COUNTER("xla::");
XLATensorPtr self_tensor = bridge::GetXlaTensor(self);
Copy link
Collaborator

@JackCaoG JackCaoG Jan 17, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we need to make sure only one of (scales_h + scales_w ) and output_size is specified, otherwise I don't know which one we should believe. Ideally upsteram already does that check but we can make it more explict.

Copy link
Collaborator Author

@lsy323 lsy323 Jan 18, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After investigating using GDB, I found the output_size will always be filled by upstream (at: https://github.com/pytorch/pytorch/blob/88366a907549abdd7e2c402a961b60c2be910824/aten/src/ATen/native/UpSampleBilinear2d.cpp#L166)

With the current upstream implementation, we can rely on the output_size inferred from scales_h/w by upstream. However, I think we can keep the scale factor shape inference here to make it future-proof.

In the scale factor shape inference block, I think we can do a shape validation if output_size is not empty, to make sure the output shape and the inferred shape are the same.

absl::Span<const int64_t> input_dims =
self_tensor->shape().get().dimensions();
std::vector<int64_t> scaled_output_size =
torch::lazy::ToVector<int64_t>(output_size);
if ((scales_h && *scales_h != 1.0) || (scales_w && *scales_w != 1.0)) {
return at::native::call_fallback_fn<
&xla_cpu_fallback, ATEN_OP(upsample_bilinear2d)>::call(self,
output_size,
align_corners,
scales_h,
scales_w);
scaled_output_size = GetOutputSizeWithScale(input_dims, scales_h, scales_w,
scaled_output_size);
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";
}
}
return bridge::AtenFromXlaTensor(tensor_methods::upsample_bilinear2d(
self_tensor, torch::lazy::ToVector<int64_t>(output_size), align_corners));
self_tensor, scaled_output_size, align_corners));
}

at::Tensor XLANativeFunctions::upsample_bilinear2d_backward(
Expand Down