xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/UpSampleNearest1d.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/core/Tensor.h>
3 #include <ATen/TensorMeta.h>
4 #include <ATen/TensorUtils.h>
5 #include <ATen/native/UpSample.h>
6 
7 #ifndef AT_PER_OPERATOR_HEADERS
8 #include <ATen/Functions.h>
9 #include <ATen/NativeFunctions.h>
10 #else
11 #include <ATen/ops/_upsample_nearest_exact1d.h>
12 #include <ATen/ops/_upsample_nearest_exact1d_backward.h>
13 #include <ATen/ops/_upsample_nearest_exact1d_backward_native.h>
14 #include <ATen/ops/_upsample_nearest_exact1d_native.h>
15 #include <ATen/ops/upsample_nearest1d.h>
16 #include <ATen/ops/upsample_nearest1d_backward.h>
17 #include <ATen/ops/upsample_nearest1d_backward_native.h>
18 #include <ATen/ops/upsample_nearest1d_native.h>
19 #endif
20 
21 namespace at::meta {
22 
TORCH_META_FUNC(upsample_nearest1d)23 TORCH_META_FUNC(upsample_nearest1d) (
24     const Tensor& input, IntArrayRef output_size, std::optional<double> scales
25 ) {
26   auto full_output_size = native::upsample_1d_common_check(input.sizes(), output_size);
27 
28   // Allow for empty batch size but not other dimensions
29   TORCH_CHECK(
30       (input.size(1) != 0 && input.size(2) != 0) && input.dim() == 3,
31       "Non-empty 3D data tensor expected but got a tensor with sizes ",
32       input.sizes());
33 
34   set_output_raw_strided(0, full_output_size, {}, input.options());
35 }
36 
TORCH_META_FUNC(_upsample_nearest_exact1d)37 TORCH_META_FUNC(_upsample_nearest_exact1d) (
38   const Tensor& input, IntArrayRef output_size, std::optional<double> scales
39 ) {
40   auto full_output_size = native::upsample_1d_common_check(input.sizes(), output_size);
41 
42   // Allow for empty batch size but not other dimensions
43   TORCH_CHECK(
44       (input.size(1) != 0 && input.size(2) != 0) && input.dim() == 3,
45       "Non-empty 3D data tensor expected but got a tensor with sizes ",
46       input.sizes());
47 
48   set_output_raw_strided(0, full_output_size, {}, input.options());
49 }
50 
TORCH_META_FUNC(upsample_nearest1d_backward)51 TORCH_META_FUNC(upsample_nearest1d_backward) (
52     const Tensor& grad_output, IntArrayRef output_size, IntArrayRef input_size, std::optional<double> scales
53 ) {
54   auto full_output_size = native::upsample_1d_common_check(input_size, output_size);
55 
56   check_dim_size(grad_output, 3, 0, full_output_size[0]);
57   check_dim_size(grad_output, 3, 1, full_output_size[1]);
58   check_dim_size(grad_output, 3, 2, full_output_size[2]);
59 
60   set_output_raw_strided(0, input_size, {}, grad_output.options());
61 }
62 
TORCH_META_FUNC(_upsample_nearest_exact1d_backward)63 TORCH_META_FUNC(_upsample_nearest_exact1d_backward) (
64   const Tensor& grad_output, IntArrayRef output_size, IntArrayRef input_size, std::optional<double> scales
65 ) {
66   auto full_output_size = native::upsample_1d_common_check(input_size, output_size);
67 
68   check_dim_size(grad_output, 3, 0, full_output_size[0]);
69   check_dim_size(grad_output, 3, 1, full_output_size[1]);
70   check_dim_size(grad_output, 3, 2, full_output_size[2]);
71 
72   set_output_raw_strided(0, input_size, {}, grad_output.options());
73 }
74 
75 } // namespace at::meta
76 
77 
78 namespace at::native {
79 
TORCH_IMPL_FUNC(upsample_nearest1d_out_cpu)80 TORCH_IMPL_FUNC(upsample_nearest1d_out_cpu) (
81     const Tensor& input,
82     IntArrayRef output_size,
83     std::optional<double> scales,
84     const Tensor& output
85 ) {
86   upsample_nearest1d_kernel(kCPU, output, input, scales);
87 }
88 
TORCH_IMPL_FUNC(_upsample_nearest_exact1d_out_cpu)89 TORCH_IMPL_FUNC(_upsample_nearest_exact1d_out_cpu) (
90     const Tensor& input,
91     IntArrayRef output_size,
92     std::optional<double> scales,
93     const Tensor& output
94 ) {
95   _upsample_nearest_exact1d_kernel(kCPU, output, input, scales);
96 }
97 
TORCH_IMPL_FUNC(upsample_nearest1d_backward_out_cpu)98 TORCH_IMPL_FUNC(upsample_nearest1d_backward_out_cpu) (
99     const Tensor& grad_output,
100     IntArrayRef output_size,
101     IntArrayRef input_size,
102     std::optional<double> scales,
103     const Tensor& grad_input
104 ) {
105   grad_input.zero_();
106   upsample_nearest1d_backward_kernel(kCPU, grad_input, grad_output, scales);
107 }
108 
TORCH_IMPL_FUNC(_upsample_nearest_exact1d_backward_out_cpu)109 TORCH_IMPL_FUNC(_upsample_nearest_exact1d_backward_out_cpu) (
110     const Tensor& grad_output,
111     IntArrayRef output_size,
112     IntArrayRef input_size,
113     std::optional<double> scales,
114     const Tensor& grad_input
115 ) {
116   grad_input.zero_();
117   _upsample_nearest_exact1d_backward_kernel(kCPU, grad_input, grad_output, scales);
118 }
119 
120 using at::native::upsample::compute_output_size;
121 using at::native::upsample::get_scale_value;
122 
123 // vec variants
124 
upsample_nearest1d(const Tensor & input,at::OptionalIntArrayRef output_size,std::optional<ArrayRef<double>> scale_factors)125 Tensor upsample_nearest1d(
126     const Tensor& input,
127     at::OptionalIntArrayRef output_size,
128     std::optional<ArrayRef<double>> scale_factors) {
129   auto osize = compute_output_size(input.sizes(), output_size, scale_factors);
130   auto scale_w = get_scale_value(scale_factors, 0);
131   return at::upsample_nearest1d(input, osize, scale_w);
132 }
133 
_upsample_nearest_exact1d(const Tensor & input,at::OptionalIntArrayRef output_size,std::optional<ArrayRef<double>> scale_factors)134 Tensor _upsample_nearest_exact1d(
135     const Tensor& input,
136     at::OptionalIntArrayRef output_size,
137     std::optional<ArrayRef<double>> scale_factors) {
138   auto osize = compute_output_size(input.sizes(), output_size, scale_factors);
139   auto scale_w = get_scale_value(scale_factors, 0);
140   return at::_upsample_nearest_exact1d(input, osize, scale_w);
141 }
142 
143 DEFINE_DISPATCH(upsample_nearest1d_kernel);
144 DEFINE_DISPATCH(_upsample_nearest_exact1d_kernel);
145 DEFINE_DISPATCH(upsample_nearest1d_backward_kernel);
146 DEFINE_DISPATCH(_upsample_nearest_exact1d_backward_kernel);
147 
148 } // namespace at::native
149