1 // Adapted from interp.cpp from Caffe util by Pauline Luc
2 // Originally developed by George Papandreou
3 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
4
5 #include <ATen/core/Tensor.h>
6 #include <ATen/TensorMeta.h>
7 #include <ATen/native/UpSample.h>
8 #include <c10/util/irange.h>
9
10 #ifndef AT_PER_OPERATOR_HEADERS
11 #include <ATen/Functions.h>
12 #include <ATen/NativeFunctions.h>
13 #else
14 #include <ATen/ops/upsample_trilinear3d.h>
15 #include <ATen/ops/upsample_trilinear3d_backward.h>
16 #include <ATen/ops/upsample_trilinear3d_backward_native.h>
17 #include <ATen/ops/upsample_trilinear3d_native.h>
18 #endif
19
20 namespace at::meta {
21
TORCH_META_FUNC(upsample_trilinear3d)22 TORCH_META_FUNC(upsample_trilinear3d) (
23 const Tensor& input,
24 IntArrayRef output_size,
25 bool align_corners,
26 std::optional<double> scales_d,
27 std::optional<double> scales_h,
28 std::optional<double> scales_w
29 ) {
30 auto full_output_size = native::upsample_3d_common_check(input.sizes(), output_size);
31
32 // Allow for empty batch size but not other dimensions
33 TORCH_CHECK(
34 input.numel() != 0 || c10::multiply_integers(input.sizes().begin() + 1, input.sizes().end()),
35 "Non-empty 5D data tensor expected but got a tensor with sizes ",
36 input.sizes());
37
38 set_output_raw_strided(0, full_output_size, {}, input.options().memory_format(input.suggest_memory_format()));
39 }
40
TORCH_META_FUNC(upsample_trilinear3d_backward)41 TORCH_META_FUNC(upsample_trilinear3d_backward) (
42 const Tensor& grad_output,
43 IntArrayRef output_size,
44 IntArrayRef input_size,
45 bool align_corners,
46 std::optional<double> scales_d,
47 std::optional<double> scales_h,
48 std::optional<double> scales_w
49 ) {
50 auto full_output_size = native::upsample_3d_common_check(input_size, output_size);
51
52 TORCH_CHECK(
53 grad_output.dim() == 5,
54 "Expected grad_output to be a tensor of dimension 5 but got: dimension ", grad_output.dim());
55
56 for (const auto i : c10::irange(5)) {
57 TORCH_CHECK(
58 grad_output.size(i) == full_output_size[i],
59 "Expected grad_output to have the same shape as output;",
60 " output.size(", i, ") = ", full_output_size[i],
61 " but got grad_output.size(", i, ") = ", grad_output.size(i));
62 }
63
64 set_output_raw_strided(0, input_size, {}, grad_output.options().memory_format(grad_output.suggest_memory_format()));
65 }
66
67 } // namespace at::meta
68 namespace at::native {
69
TORCH_IMPL_FUNC(upsample_trilinear3d_out_cpu)70 TORCH_IMPL_FUNC(upsample_trilinear3d_out_cpu) (
71 const Tensor& input,
72 IntArrayRef output_size,
73 bool align_corners,
74 std::optional<double> scales_d,
75 std::optional<double> scales_h,
76 std::optional<double> scales_w,
77 const Tensor& output
78 ) {
79 upsample_trilinear3d_kernel(kCPU, output, input, align_corners, scales_d, scales_h, scales_w);
80 }
81
TORCH_IMPL_FUNC(upsample_trilinear3d_backward_out_cpu)82 TORCH_IMPL_FUNC(upsample_trilinear3d_backward_out_cpu) (
83 const Tensor& grad_output,
84 IntArrayRef output_size,
85 IntArrayRef input_size,
86 bool align_corners,
87 std::optional<double> scales_d,
88 std::optional<double> scales_h,
89 std::optional<double> scales_w,
90 const Tensor& grad_input
91 ) {
92 grad_input.zero_();
93 upsample_trilinear3d_backward_kernel(kCPU, grad_input, grad_output, align_corners, scales_d, scales_h, scales_w);
94 }
95
96 // vec variants
97
98 using at::native::upsample::compute_output_size;
99 using at::native::upsample::get_scale_value;
100
upsample_trilinear3d(const Tensor & input,at::OptionalIntArrayRef output_size,bool align_corners,std::optional<ArrayRef<double>> scale_factors)101 Tensor upsample_trilinear3d(
102 const Tensor& input,
103 at::OptionalIntArrayRef output_size,
104 bool align_corners,
105 std::optional<ArrayRef<double>> scale_factors) {
106 auto osize = compute_output_size(input.sizes(), output_size, scale_factors);
107 auto scale_d = get_scale_value(scale_factors, 0);
108 auto scale_h = get_scale_value(scale_factors, 1);
109 auto scale_w = get_scale_value(scale_factors, 2);
110 return at::upsample_trilinear3d(input, osize, align_corners, scale_d, scale_h, scale_w);
111 }
112
113 DEFINE_DISPATCH(upsample_trilinear3d_kernel);
114 DEFINE_DISPATCH(upsample_trilinear3d_backward_kernel);
115
116 } // namespace at::native
117