1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/PadNd.h>
3 #include <ATen/core/Tensor.h>
4
5 #include <c10/util/irange.h>
6
7 #ifndef AT_PER_OPERATOR_HEADERS
8 #include <ATen/Functions.h>
9 #include <ATen/NativeFunctions.h>
10 #else
11 #include <ATen/ops/_empty_affine_quantized.h>
12 #include <ATen/ops/_pad_circular.h>
13 #include <ATen/ops/_pad_circular_native.h>
14 #include <ATen/ops/_pad_enum_native.h>
15 #include <ATen/ops/constant_pad_nd.h>
16 #include <ATen/ops/constant_pad_nd_native.h>
17 #include <ATen/ops/empty.h>
18 #include <ATen/ops/pad_native.h>
19 #include <ATen/ops/reflection_pad1d.h>
20 #include <ATen/ops/reflection_pad2d.h>
21 #include <ATen/ops/reflection_pad3d.h>
22 #include <ATen/ops/replication_pad1d.h>
23 #include <ATen/ops/replication_pad2d.h>
24 #include <ATen/ops/replication_pad3d.h>
25 #endif
26
27 namespace at::native {
28
constant_pad_nd(const Tensor & self,IntArrayRef pad,const Scalar & value)29 Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value) {
30 TORCH_CHECK(pad.size() % 2 == 0, "Length of pad must be even but instead it equals ",
31 pad.size());
32
33 auto input_sizes = self.sizes();
34 auto l_inp = self.dim();
35
36 auto l_pad = pad.size() / 2;
37 auto l_diff = l_inp - l_pad;
38 TORCH_CHECK(l_inp >= (int64_t)l_pad, "Length of pad should be no more than twice the number of "
39 "dimensions of the input. Pad length is ", pad.size(), "while the input has ",
40 l_inp, "dimensions.");
41
42 std::vector<int64_t> new_shape;
43
44 bool all_pads_non_positive = true;
45
46 auto c_input = self;
47 for (const auto i : c10::irange(l_diff, l_inp)) {
48 auto pad_idx = 2 * (l_inp - i - 1);
49 if (pad[pad_idx] < 0) {
50 c_input = c_input.narrow(i, -pad[pad_idx], c_input.size(i) + pad[pad_idx]);
51 } else if (pad[pad_idx] != 0) {
52 all_pads_non_positive = false;
53 }
54 if (pad[pad_idx + 1] < 0) {
55 c_input = c_input.narrow(i, 0, c_input.size(i) + pad[pad_idx + 1]);
56 } else if (pad[pad_idx + 1] != 0) {
57 all_pads_non_positive = false;
58 }
59 }
60
61 // if none of the pads are positive we can optimize and just return the result
62 // of calling .narrow() on the input
63 if (all_pads_non_positive) {
64 return c_input.clone();
65 }
66
67
68 for (size_t i = 0; i < (size_t)l_diff; i ++) {
69 new_shape.emplace_back(input_sizes[i]);
70 }
71
72 for (const auto i : c10::irange((size_t)l_pad)) {
73 auto pad_idx = pad.size() - ((i + 1) * 2);
74 auto new_dim = input_sizes[l_diff + i] + pad[pad_idx] + pad[pad_idx + 1];
75 TORCH_CHECK(new_dim > 0, "The input size ", input_sizes[l_diff + i], ", plus negative padding ",
76 pad[pad_idx], " and ", pad[pad_idx + 1], " resulted in a negative output size, "
77 "which is invalid. Check dimension ", l_diff + i, " of your input.");
78 new_shape.emplace_back(new_dim);
79 }
80
81 at::Tensor output;
82 const auto memory_format = self.suggest_memory_format();
83 if (self.is_quantized()) {
84 const auto qscheme = self.qscheme();
85 TORCH_CHECK(qscheme == kPerTensorAffine || qscheme == kPerTensorSymmetric,
86 "Only per-tensor padding is supported.");
87 output = at::_empty_affine_quantized(
88 new_shape, self.options().memory_format(memory_format),
89 self.q_scale(), self.q_zero_point(), std::nullopt);
90 } else {
91 output = at::empty(new_shape, self.options().memory_format(memory_format));
92 }
93 output.fill_(value);
94
95 auto c_output = output;
96 for (const auto i : c10::irange(l_diff, l_inp)) {
97 auto pad_idx = 2 * (l_inp - i - 1);
98 if (pad[pad_idx] > 0) {
99 c_output = c_output.narrow(i, pad[pad_idx], c_output.size(i) - pad[pad_idx]);
100 }
101 if (pad[pad_idx + 1] > 0) {
102 c_output = c_output.narrow(i, 0, c_output.size(i) - pad[pad_idx + 1]);
103 }
104 }
105 c_output.copy_(c_input);
106 return output;
107 }
108
_pad_circular_symint(const Tensor & self,c10::SymIntArrayRef padding)109 Tensor _pad_circular_symint(const Tensor &self, c10::SymIntArrayRef padding) {
110 const auto in_shape = self.sym_sizes();
111 const auto self_ndim = static_cast<int64_t>(in_shape.size());
112
113 // number of dimensions that are padded
114 const auto ndim_padded = padding.size() / 2;
115 // number of preceding non_padded dimensions (1 for no_batch_dim case or 2)
116 const auto ndim_nonpadded = self_ndim - ndim_padded;
117
118 TORCH_CHECK(ndim_nonpadded == 1 || ndim_nonpadded == 2,
119 "Invalid padding size, expected 1 or 2 non-padded dimensions, ",
120 "which would be equivalent to padding of length ",
121 (self_ndim - 1) * 2,
122 " or ",
123 (self_ndim - 2) * 2,
124 " respectively but got ",
125 padding.size());
126
127 c10::SymDimVector out_shape(in_shape.size());
128 for (const auto i: c10::irange(ndim_nonpadded)) {
129 out_shape[i] = in_shape[i];
130 }
131
132 // Get shape of padded tensor
133 for (const auto i : c10::irange(ndim_padded)) {
134 const auto& pad_l = padding[2 * (ndim_padded - i - 1) + 0];
135 const auto& pad_r = padding[2 * (ndim_padded - i - 1) + 1];
136 const auto& size = in_shape[ndim_nonpadded + i];
137 out_shape[ndim_nonpadded + i] = size + pad_l + pad_r;
138
139 TORCH_CHECK(
140 pad_l <= size && pad_r <= size,
141 "Padding value causes wrapping around more than once.");
142 TORCH_CHECK(
143 out_shape[ndim_nonpadded + i] >= 0,
144 "Negative padding value is resulting in an empty dimension");
145 }
146
147 auto out = self.new_empty_symint(out_shape, self.options());
148
149 // Put original array into the padded array
150 Tensor out_slice = out;
151 Tensor in_slice = self;
152 const SymInt zero = 0;
153 for (const auto i : c10::irange(ndim_padded)) {
154 const auto dim = ndim_padded - i + ndim_nonpadded - 1;
155 const auto& pad_l = padding[2*i + 0];
156 const auto& pad_r = padding[2*i + 1];
157 out_slice = out_slice.slice_symint(dim, std::max(pad_l, zero), out_shape[dim] - std::max(pad_r, zero));
158 in_slice = in_slice.slice_symint(dim, std::max(-pad_l, zero), in_shape[dim] - std::max(-pad_r, zero));
159 }
160 out_slice.copy_(in_slice);
161
162 // The following steps first pad the beginning of the tensor (left side),
163 // and then pad the end of the tensor (right side).
164 // Note: Corners will be written more than once when ndim_padded > 1.
165 //
166 // Only in cases where padding values are > 0 are when additional copying
167 // is required.
168 for (const auto i : c10::irange(ndim_padded)) {
169 const auto dim = ndim_padded - i + ndim_nonpadded - 1;
170 const auto& pad_l = padding[2*i + 0];
171 const auto& pad_r = padding[2*i + 1];
172
173 if (pad_l > 0) {
174 out_slice = out.slice_symint(dim, 0, pad_l);
175 in_slice = out.slice_symint(dim,
176 out_shape[dim] - pad_l - std::max(pad_r, zero),
177 out_shape[dim] - std::max(pad_r, zero));
178 out_slice.copy_(in_slice);
179 }
180
181 if (pad_r > 0) {
182 out_slice = out.slice_symint(dim, out_shape[dim] - pad_r, out_shape[dim]);
183 in_slice = out.slice_symint(dim, std::max(pad_l, zero), std::max(pad_l, zero) + pad_r);
184 out_slice.copy_(in_slice);
185 }
186 }
187
188 return out;
189 }
190
_pad_enum_symint(const Tensor & self,c10::SymIntArrayRef pad,int64_t mode_int,std::optional<double> value)191 Tensor _pad_enum_symint(const Tensor &self, c10::SymIntArrayRef pad, int64_t mode_int, std::optional<double> value) {
192 const auto input_dim = self.dim();
193 TORCH_CHECK(pad.size() % 2 == 0, "Padding length must be divisible by 2");
194 TORCH_CHECK(static_cast<int64_t>(pad.size()) <= input_dim * 2,
195 "Padding length should be less than or equal to two times the input dimension but got padding length ", pad.size(), " and input of dimension ", input_dim);
196 auto mode = static_cast<at::padding_mode>(mode_int);
197
198 if (mode == at::padding_mode::constant) {
199 return at::constant_pad_nd_symint(self, pad, value.value_or(0.0));
200 }
201 TORCH_CHECK(!value.has_value() || *value == 0,
202 "Padding mode \"", padding_mode_string(mode),
203 "\" doesn't take in value argument");
204
205 if (pad.size() == 2 && (input_dim == 2 || input_dim == 3)) {
206 switch (mode) {
207 case at::padding_mode::reflect: return at::reflection_pad1d_symint(self, pad);
208 case at::padding_mode::replicate: return at::replication_pad1d_symint(self, pad);
209 case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
210 default: {}
211 }
212 } else if(pad.size() == 4 && (input_dim == 3 || input_dim == 4)) {
213 switch (mode) {
214 case at::padding_mode::reflect: return at::reflection_pad2d_symint(self, pad);
215 case at::padding_mode::replicate: return at::replication_pad2d_symint(self, pad);
216 case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
217 default: {}
218 }
219 } else if (pad.size() == 6 && (input_dim == 4 || input_dim == 5)) {
220 switch (mode) {
221 case at::padding_mode::reflect: return at::reflection_pad3d_symint(self, pad);
222 case at::padding_mode::replicate: return at::replication_pad3d_symint(self, pad);
223 case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
224 default: {}
225 }
226 }
227 C10_THROW_ERROR(NotImplementedError,
228 "Only 2D, 3D, 4D, 5D padding with non-constant padding are supported for now");
229 }
230
pad_symint(const Tensor & self,c10::SymIntArrayRef pad,c10::string_view mode,std::optional<double> value)231 Tensor pad_symint(const Tensor &self, c10::SymIntArrayRef pad, c10::string_view mode, std::optional<double> value) {
232 const auto mode_enum = [&] {
233 if (mode == "reflect") {
234 return at::padding_mode::reflect;
235 } else if (mode == "constant") {
236 return at::padding_mode::constant;
237 } else if (mode == "replicate") {
238 return at::padding_mode::replicate;
239 } else if (mode == "circular") {
240 return at::padding_mode::circular;
241 }
242 C10_THROW_ERROR(NotImplementedError,
243 c10::str("Unrecognised padding mode ", mode));
244 }();
245 return at::native::_pad_enum_symint(self, pad, static_cast<int64_t>(mode_enum), value);
246 }
247
248 } // namespace at::native
249