1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <cstring>
10
11 #include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
12 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
13 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
14
15 namespace torch {
16 namespace executor {
17
18 using Tensor = exec_aten::Tensor;
19
20 namespace {
21
as_strided_copy_compute_storage_nbytes(IntArrayRef sizes,IntArrayRef strides,size_t itemsize_bytes)22 size_t as_strided_copy_compute_storage_nbytes(
23 IntArrayRef sizes,
24 IntArrayRef strides,
25 size_t itemsize_bytes) {
26 // size of the underlying storage is 1 bigger than the offset
27 // of the last element according to stride
28 size_t size = 1;
29 for (size_t i = 0; i < sizes.size(); ++i) {
30 if (sizes[i] == 0) {
31 return 0;
32 }
33 size += strides[i] * (sizes[i] - 1);
34 }
35 return size * itemsize_bytes;
36 }
37
38 } // namespace
39
check_as_strided_copy_args(const Tensor & in,ArrayRef<int64_t> size,ArrayRef<int64_t> stride,optional<int64_t> storage_offset,Tensor & out)40 bool check_as_strided_copy_args(
41 const Tensor& in,
42 ArrayRef<int64_t> size,
43 ArrayRef<int64_t> stride,
44 optional<int64_t> storage_offset,
45 Tensor& out) {
46 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
47 ET_LOG_MSG_AND_RETURN_IF_FALSE(
48 size.size() == stride.size(), "mismatch in length of strides and shape");
49 for (const auto& val : stride) {
50 ET_LOG_MSG_AND_RETURN_IF_FALSE(
51 val >= 0,
52 "as_strided: Negative strides are not supported at the moment");
53 }
54
55 int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
56 ET_LOG_MSG_AND_RETURN_IF_FALSE(offset >= 0, "Negative storage offset");
57
58 // Check that the requested storage is within bounds of input storage
59 size_t storage_size_bytes =
60 as_strided_copy_compute_storage_nbytes(size, stride, in.element_size());
61 size_t storage_offset_bytes = offset * in.element_size();
62 if (storage_size_bytes == 0) {
63 return true;
64 }
65 size_t new_storage_size_bytes = in.nbytes();
66 ET_LOG_MSG_AND_RETURN_IF_FALSE(
67 storage_size_bytes + storage_offset_bytes <= new_storage_size_bytes,
68 "Requiring a storage size of %zd are out of bounds for storage of size %zd",
69 storage_size_bytes + storage_offset_bytes,
70 new_storage_size_bytes);
71 return true;
72 }
73
check_cat_args(exec_aten::ArrayRef<Tensor> tensors,int64_t dim,Tensor & out)74 bool check_cat_args(
75 exec_aten::ArrayRef<Tensor> tensors,
76 int64_t dim,
77 Tensor& out) {
78 // Ensure the input tensors list is non-empty
79 ET_LOG_AND_RETURN_IF_FALSE(tensors.size() > 0);
80
81 // Find the first non-empty tensor in the list to use as a reference
82 size_t ref_i = 0;
83 for (size_t i = 0; i < tensors.size(); ++i) {
84 if (tensors[i].numel() > 0) {
85 ref_i = i;
86 break;
87 }
88 }
89
90 // "All tensors must either have the same shape (except in the concatenating
91 // dimension) or be empty."
92 // https://pytorch.org/docs/stable/generated/torch.cat.html
93 for (size_t i = 0; i < tensors.size(); ++i) {
94 // All input dtypes must be castable to the output dtype.
95 ET_LOG_AND_RETURN_IF_FALSE(
96 canCast(tensors[i].scalar_type(), out.scalar_type()));
97
98 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dim_order(tensors[i], out));
99
100 // Empty tensors have no shape constraints.
101 if (tensors[i].numel() == 0) {
102 continue;
103 }
104
105 // All input tensors must have the same number of dimensions.
106 ET_LOG_AND_RETURN_IF_FALSE(
107 tensor_is_rank(tensors[ref_i], tensors[i].dim()));
108
109 for (size_t d = 0; d < tensors[i].dim(); ++d) {
110 if (d != dim) {
111 ET_LOG_AND_RETURN_IF_FALSE(
112 tensors_have_same_size_at_dims(tensors[i], d, tensors[ref_i], d));
113 }
114 }
115 }
116
117 // Ensure dim is in range.
118 ET_LOG_AND_RETURN_IF_FALSE(
119 tensors[ref_i].numel() == 0 || tensors[ref_i].dim() > dim);
120 ET_LOG_AND_RETURN_IF_FALSE(dim >= 0);
121
122 return true;
123 }
124
get_cat_out_target_size(exec_aten::ArrayRef<Tensor> tensors,int64_t dim,exec_aten::SizesType * out_sizes,size_t * out_ndim)125 void get_cat_out_target_size(
126 exec_aten::ArrayRef<Tensor> tensors,
127 int64_t dim,
128 exec_aten::SizesType* out_sizes,
129 size_t* out_ndim) {
130 // Find the first non-1D-or-empty tensor in the list to use as a reference
131 // because an 1D empty tensor is a wildcard and should be ignored when we
132 // calculate out dim
133 size_t ref_i = 0;
134 size_t cat_dim_size = 0;
135 for (size_t i = 0; i < tensors.size(); ++i) {
136 if (tensors[i].numel() > 0) {
137 cat_dim_size += tensors[i].size(dim);
138 }
139 if (tensors[i].dim() != 1 || tensors[i].numel() != 0) {
140 ref_i = i;
141 }
142 }
143
144 *out_ndim = tensors[ref_i].dim();
145
146 for (size_t d = 0; d < *out_ndim; ++d) {
147 if (d != dim) {
148 out_sizes[d] = tensors[ref_i].size(d);
149 } else {
150 out_sizes[d] = cat_dim_size;
151 }
152 }
153 }
154
check_expand_copy_args(const Tensor & input,ArrayRef<int64_t> expand_sizes,bool implicit,Tensor & out)155 bool check_expand_copy_args(
156 const Tensor& input,
157 ArrayRef<int64_t> expand_sizes,
158 bool implicit,
159 Tensor& out) {
160 (void)out;
161
162 ET_LOG_MSG_AND_RETURN_IF_FALSE(
163 implicit == false,
164 "This operator is not implemented for when implicit == true.");
165
166 ET_LOG_MSG_AND_RETURN_IF_FALSE(
167 expand_sizes.size() >= input.sizes().size(),
168 "The number of sizes provided (%zu) must at least be equal to the number of dimensions in the tensor (%zu)",
169 expand_sizes.size(),
170 input.sizes().size());
171
172 ET_LOG_MSG_AND_RETURN_IF_FALSE(
173 expand_sizes.size() <= kTensorDimensionLimit,
174 "The number of expanded dims (%zu) exceeds the configured maximum (%zu). Increase this limit.",
175 expand_sizes.size(),
176 kTensorDimensionLimit);
177
178 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(input, out));
179
180 return true;
181 }
182
get_expand_copy_out_target_size(exec_aten::ArrayRef<exec_aten::SizesType> self_sizes,exec_aten::ArrayRef<int64_t> expand_sizes,exec_aten::SizesType * output_sizes,size_t * output_rank)183 bool get_expand_copy_out_target_size(
184 exec_aten::ArrayRef<exec_aten::SizesType> self_sizes,
185 exec_aten::ArrayRef<int64_t> expand_sizes,
186 exec_aten::SizesType* output_sizes,
187 size_t* output_rank) {
188 auto j{expand_sizes.size()};
189 *output_rank = 0;
190
191 for (size_t i{self_sizes.size()}; i > 0 && j > 0;) {
192 --i;
193 --j;
194
195 output_sizes[j] = expand_sizes[j];
196
197 if (expand_sizes[j] == -1) {
198 // -1 can use for replacing any corresponding dimension
199 output_sizes[j] = self_sizes[i];
200 } else if (self_sizes[i] != 1) {
201 ET_LOG_MSG_AND_RETURN_IF_FALSE(
202 expand_sizes[j] == self_sizes[i],
203 "The expanded size of the tensor (%zu) must match the existing size (%zu) at non-singleton dimension %zu.",
204 (size_t)expand_sizes[j],
205 (size_t)self_sizes[i],
206 i);
207 }
208 }
209
210 // The leading expand_sizes cannot be negative
211 while (j > 0) {
212 --j;
213 output_sizes[j] = expand_sizes[j];
214 ET_LOG_MSG_AND_RETURN_IF_FALSE(
215 expand_sizes[j] >= 0,
216 "The expanded size of the tensor (%zu) isn't allowed in a leading, non-existing dimension %zu",
217 (size_t)expand_sizes[j],
218 j);
219 }
220
221 *output_rank = expand_sizes.size();
222 return true;
223 }
224
check_permute_copy_args(const Tensor & in,IntArrayRef dims,Tensor & out)225 bool check_permute_copy_args(const Tensor& in, IntArrayRef dims, Tensor& out) {
226 ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(in, dims.size()));
227 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
228
229 // Make sure no dimensions are duplicated and all in the range [-in.dim(),
230 // in.dim() - 1].
231 bool dim_exist[kTensorDimensionLimit];
232 memset(dim_exist, false, sizeof(dim_exist));
233
234 for (int i = 0; i < dims.size(); i++) {
235 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dims[i]));
236 // Convert dimension to a non-negative number in the range
237 // [0 .. in.dim() - 1].
238 size_t dim = dims[i] >= 0 ? dims[i] : in.dim() + dims[i];
239
240 // Internal check, since we have already validated this
241 ET_LOG_AND_RETURN_IF_FALSE(dim < kTensorDimensionLimit && dim >= 0);
242
243 // Check that the dimension hasn't been seen previously.
244 ET_LOG_MSG_AND_RETURN_IF_FALSE(
245 dim_exist[dim] == false, "duplicate dims are not allowed.");
246
247 dim_exist[dim] = true;
248 }
249
250 return true;
251 }
252
check_unbind_copy_args(const Tensor & in,int64_t dim,TensorList out)253 bool check_unbind_copy_args(const Tensor& in, int64_t dim, TensorList out) {
254 ET_LOG_MSG_AND_RETURN_IF_FALSE(
255 in.dim() > 0, "in must have at least one dimension; saw %zd", in.dim());
256
257 ET_LOG_AND_RETURN_IF_FALSE(dim_is_valid(dim, in.dim()));
258
259 const ssize_t dim_size = in.size(dim);
260 ET_LOG_MSG_AND_RETURN_IF_FALSE(
261 dim_size == out.size(),
262 "out tensorlist's length %zd must equal unbind dim %" PRId64
263 " size = %zd.",
264 out.size(),
265 dim,
266 dim_size);
267
268 // Validate each output.
269 for (size_t i = 0; i < out.size(); ++i) {
270 // All output dtypes must be the same.
271 ET_LOG_MSG_AND_RETURN_IF_FALSE(
272 out[i].scalar_type() == out[0].scalar_type(),
273 "out[%zu] dtype %" PRId8 " != out[0] dtype %" PRId8,
274 i,
275 static_cast<int8_t>(out[i].scalar_type()),
276 static_cast<int8_t>(out[0].scalar_type()));
277
278 // output tensor must have # of dims = in.dim() -1
279 ET_LOG_MSG_AND_RETURN_IF_FALSE(
280 out[i].dim() == (in.dim() - 1),
281 "out[%zu] dim %zd != in dim %zd",
282 i,
283 out[i].dim(),
284 in.dim() - 1);
285
286 // Check the shape of the output.
287 for (ssize_t d = 0, out_d = 0; d < in.dim(); ++d) {
288 if (d != dim) {
289 ET_LOG_MSG_AND_RETURN_IF_FALSE(
290 out[i].size(out_d) == in.size(d),
291 "out[%zu].size(%zd) %zd != in.size(%zd) %zd",
292 i,
293 d,
294 out[i].size(out_d),
295 d,
296 in.size(d));
297 out_d++;
298 }
299 }
300 }
301
302 return true;
303 }
304
get_permute_copy_out_target_size(const Tensor & in,IntArrayRef dims,exec_aten::SizesType * out_sizes,size_t * out_ndim)305 void get_permute_copy_out_target_size(
306 const Tensor& in,
307 IntArrayRef dims,
308 exec_aten::SizesType* out_sizes,
309 size_t* out_ndim) {
310 *out_ndim = in.dim();
311
312 for (size_t i = 0; i < in.dim(); ++i) {
313 out_sizes[i] = in.size(dims[i] >= 0 ? dims[i] : dims[i] + in.dim());
314 }
315 }
316
check_pixel_shuffle_args(const Tensor & in,int64_t upscale_factor,Tensor & out)317 bool check_pixel_shuffle_args(
318 const Tensor& in,
319 int64_t upscale_factor,
320 Tensor& out) {
321 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
322 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 3));
323 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(out, 3));
324 ET_LOG_AND_RETURN_IF_FALSE(upscale_factor > 0);
325 ET_LOG_AND_RETURN_IF_FALSE(
326 in.size(in.dim() - 3) % (upscale_factor * upscale_factor) == 0);
327 return true;
328 }
329
check_pixel_unshuffle_args(const Tensor & in,int64_t downscale_factor,Tensor & out)330 bool check_pixel_unshuffle_args(
331 const Tensor& in,
332 int64_t downscale_factor,
333 Tensor& out) {
334 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
335 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 3));
336 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(out, 3));
337 ET_LOG_AND_RETURN_IF_FALSE(downscale_factor > 0);
338 ET_LOG_AND_RETURN_IF_FALSE(in.size(in.dim() - 1) % downscale_factor == 0);
339 ET_LOG_AND_RETURN_IF_FALSE(in.size(in.dim() - 2) % downscale_factor == 0);
340 return true;
341 }
342
get_pixel_shuffle_out_target_size(const Tensor & in,int64_t upscale_factor,exec_aten::SizesType * out_sizes,size_t * out_ndim)343 void get_pixel_shuffle_out_target_size(
344 const Tensor& in,
345 int64_t upscale_factor,
346 exec_aten::SizesType* out_sizes,
347 size_t* out_ndim) {
348 *out_ndim = in.dim();
349 const exec_aten::SizesType casted_upscale_factor = upscale_factor;
350
351 size_t i = 0;
352 for (; i < in.dim() - 3; ++i) {
353 // Copy all leading dimensions in.
354 out_sizes[i] = in.size(i);
355 }
356 // The last 3 dimensions are (channel, height, width). Divide by the upscale
357 // factor squared and multiply the height and width by that factor.
358 out_sizes[i] = in.size(i) / (casted_upscale_factor * casted_upscale_factor);
359 i++;
360 out_sizes[i] = in.size(i) * casted_upscale_factor;
361 i++;
362 out_sizes[i] = in.size(i) * casted_upscale_factor;
363 }
364
get_pixel_unshuffle_out_target_size(const Tensor & in,int64_t downscale_factor,exec_aten::SizesType * out_sizes,size_t * out_ndim)365 void get_pixel_unshuffle_out_target_size(
366 const Tensor& in,
367 int64_t downscale_factor,
368 exec_aten::SizesType* out_sizes,
369 size_t* out_ndim) {
370 *out_ndim = in.dim();
371 const exec_aten::SizesType casted_factor = downscale_factor;
372
373 size_t i = 0;
374 for (; i < in.dim() - 3; ++i) {
375 // Copy all leading dimensions in.
376 out_sizes[i] = in.size(i);
377 }
378 // The last 3 dimensions are (channel, height, width). Multiply channel by
379 // the downscale factor squared and divide the height and width by that
380 // factor.
381 out_sizes[i] = in.size(i) * (casted_factor * casted_factor);
382 i++;
383 out_sizes[i] = in.size(i) / casted_factor;
384 i++;
385 out_sizes[i] = in.size(i) / casted_factor;
386 }
387
check_select_copy_out_args(const Tensor & in,int64_t dim,int64_t index,Tensor & out)388 bool check_select_copy_out_args(
389 const Tensor& in,
390 int64_t dim,
391 int64_t index,
392 Tensor& out) {
393 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 1));
394 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
395 ET_LOG_AND_RETURN_IF_FALSE(tensor_dim_has_index(in, dim, index));
396 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
397 return true;
398 }
399
get_select_copy_out_target_size(const Tensor & in,int64_t dim,exec_aten::SizesType * out_sizes,size_t * out_ndim)400 void get_select_copy_out_target_size(
401 const Tensor& in,
402 int64_t dim,
403 exec_aten::SizesType* out_sizes,
404 size_t* out_ndim) {
405 *out_ndim = in.dim() - 1;
406
407 for (size_t d = 0; d < in.dim() - 1; ++d) {
408 if (d < dim) {
409 out_sizes[d] = in.size(d);
410 } else {
411 out_sizes[d] = in.size(d + 1);
412 }
413 }
414 }
415
check_split_with_sizes_copy_args(const Tensor & in,exec_aten::ArrayRef<int64_t> split_sizes,int64_t dim,TensorList out)416 bool check_split_with_sizes_copy_args(
417 const Tensor& in,
418 exec_aten::ArrayRef<int64_t> split_sizes,
419 int64_t dim,
420 TensorList out) {
421 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 1));
422 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
423
424 ET_LOG_MSG_AND_RETURN_IF_FALSE(
425 split_sizes.size() == out.size(),
426 "Number of split sizes must match the number of output tensors");
427
428 int64_t sum = 0;
429 for (int i = 0; i < split_sizes.size(); i++) {
430 ET_LOG_MSG_AND_RETURN_IF_FALSE(
431 split_sizes[i] >= 0, "All split sizes must be non negative.");
432 sum += split_sizes[i];
433 }
434
435 const ssize_t dim_size = in.size(dim);
436 ET_LOG_MSG_AND_RETURN_IF_FALSE(
437 sum == dim_size,
438 "Sum of split sizes does not match input size at given dim");
439
440 return true;
441 }
442
get_split_with_sizes_copy_out_target_size(const Tensor & in,int64_t split_size,int64_t dim,exec_aten::SizesType * out_sizes,size_t * out_ndim)443 void get_split_with_sizes_copy_out_target_size(
444 const Tensor& in,
445 int64_t split_size,
446 int64_t dim,
447 exec_aten::SizesType* out_sizes,
448 size_t* out_ndim) {
449 *out_ndim = in.dim();
450
451 for (size_t d = 0; d < in.dim(); ++d) {
452 out_sizes[d] = in.size(d);
453 }
454 out_sizes[dim] = split_size;
455 }
456
check_squeeze_copy_dim_args(const Tensor in,int64_t dim,const Tensor out)457 bool check_squeeze_copy_dim_args(
458 const Tensor in,
459 int64_t dim,
460 const Tensor out) {
461 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
462 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
463
464 return true;
465 }
466
get_squeeze_copy_dim_out_target_size(const Tensor in,int64_t dim,exec_aten::SizesType * out_sizes,size_t * out_ndim)467 void get_squeeze_copy_dim_out_target_size(
468 const Tensor in,
469 int64_t dim,
470 exec_aten::SizesType* out_sizes,
471 size_t* out_ndim) {
472 // For 0 dim tensors, the output should also be 0 dim.
473 if (in.dim() == 0) {
474 *out_ndim = 0;
475 return;
476 }
477
478 // Specified dim is only removed if the size at the given dim is 1.
479 if (in.size(dim) == 1) {
480 *out_ndim = in.dim() - 1;
481 } else {
482 *out_ndim = in.dim();
483 }
484
485 size_t out_d = 0;
486 for (size_t in_d = 0; in_d < in.dim(); ++in_d) {
487 if (in_d != dim || in.size(in_d) != 1) {
488 out_sizes[out_d] = in.size(in_d);
489 ++out_d;
490 }
491 }
492 }
493
check_squeeze_copy_dims_args(const Tensor in,const exec_aten::ArrayRef<int64_t> dims,const Tensor out)494 bool check_squeeze_copy_dims_args(
495 const Tensor in,
496 const exec_aten::ArrayRef<int64_t> dims,
497 const Tensor out) {
498 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
499
500 for (size_t i = 0; i < dims.size(); ++i) {
501 const int64_t dim = dims[i] < 0 ? dims[i] + nonzero_dim(in) : dims[i];
502 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
503
504 // Check that a dim does not appear twice in dims
505 for (size_t j = 0; j < dims.size(); ++j) {
506 if (i != j) {
507 const int64_t dim_temp =
508 dims[j] < 0 ? dims[j] + nonzero_dim(in) : dims[j];
509 ET_LOG_MSG_AND_RETURN_IF_FALSE(
510 dim != dim_temp,
511 "dim %" PRId64 " appears multiple times in dims!",
512 dim);
513 }
514 }
515 }
516
517 return true;
518 }
519
get_squeeze_copy_dims_out_target_size(const Tensor in,const exec_aten::ArrayRef<int64_t> dims,exec_aten::SizesType * out_sizes,size_t * out_ndim)520 void get_squeeze_copy_dims_out_target_size(
521 const Tensor in,
522 const exec_aten::ArrayRef<int64_t> dims,
523 exec_aten::SizesType* out_sizes,
524 size_t* out_ndim) {
525 // For 0 dim tensors, the output should also be 0 dim.
526 if (in.dim() == 0) {
527 *out_ndim = 0;
528 return;
529 }
530
531 // A dim is only removed if the size at the given dim is 1.
532 exec_aten::SizesType dims_to_remove = 0;
533 for (size_t i = 0; i < dims.size(); ++i) {
534 int64_t dim = dims[i] < 0 ? dims[i] + nonzero_dim(in) : dims[i];
535 if (in.size(dim) == 1) {
536 ++dims_to_remove;
537 }
538 }
539 *out_ndim = in.dim() - dims_to_remove;
540
541 size_t out_d = 0;
542 for (size_t in_d = 0; in_d < in.dim(); ++in_d) {
543 bool in_d_in_dims = false;
544 for (size_t i = 0; i < dims.size(); ++i) {
545 int64_t dim = dims[i] < 0 ? dims[i] + nonzero_dim(in) : dims[i];
546 if (in_d == dim) {
547 in_d_in_dims = true;
548 break;
549 }
550 }
551 if (!in_d_in_dims || in.size(in_d) != 1) {
552 out_sizes[out_d] = in.size(in_d);
553 ++out_d;
554 }
555 }
556 }
557
check_stack_args(exec_aten::ArrayRef<Tensor> tensors,int64_t dim,Tensor & out)558 bool check_stack_args(
559 exec_aten::ArrayRef<Tensor> tensors,
560 int64_t dim,
561 Tensor& out) {
562 // Ensure the input tensors list is non-empty
563 ET_LOG_AND_RETURN_IF_FALSE(tensors.size() > 0);
564
565 // All input tensors need to be of the same size
566 // https://pytorch.org/docs/stable/generated/torch.stack.html
567 for (size_t i = 0; i < tensors.size(); i++) {
568 // All input dtypes must be castable to the output dtype.
569 ET_LOG_AND_RETURN_IF_FALSE(
570 canCast(tensors[i].scalar_type(), out.scalar_type()));
571
572 ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(tensors[i], tensors[0].dim()));
573 for (size_t d = 0; d < tensors[i].dim(); d++) {
574 ET_LOG_AND_RETURN_IF_FALSE(
575 tensors_have_same_size_at_dims(tensors[i], d, tensors[0], d));
576 }
577 }
578
579 // The output tensor will have a dimension inserted, so dim should be between
580 // 0 and ndim_of_inputs + 1
581 ET_LOG_AND_RETURN_IF_FALSE(dim >= 0 && dim < tensors[0].dim() + 1);
582
583 return true;
584 }
585
get_stack_out_target_size(exec_aten::ArrayRef<Tensor> tensors,int64_t dim,exec_aten::SizesType * out_sizes,size_t * out_ndim)586 void get_stack_out_target_size(
587 exec_aten::ArrayRef<Tensor> tensors,
588 int64_t dim,
589 exec_aten::SizesType* out_sizes,
590 size_t* out_ndim) {
591 *out_ndim = tensors[0].dim() + 1;
592
593 for (size_t d = 0; d < *out_ndim; ++d) {
594 if (d < dim) {
595 out_sizes[d] = tensors[0].size(d);
596 } else if (d == dim) {
597 out_sizes[d] = tensors.size();
598 } else {
599 out_sizes[d] = tensors[0].size(d - 1);
600 }
601 }
602 }
603
check_tril_args(const Tensor & in,Tensor & out)604 bool check_tril_args(const Tensor& in, Tensor& out) {
605 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
606 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 2));
607 return true;
608 }
609
check_split_copy_args(const Tensor & input,int64_t split_size,int64_t dim,TensorList out)610 bool check_split_copy_args(
611 const Tensor& input,
612 int64_t split_size,
613 int64_t dim,
614 TensorList out) {
615 ET_LOG_MSG_AND_RETURN_IF_FALSE(
616 input.dim() > 0,
617 "input must have at least one dimension; saw %zd",
618 input.dim());
619 ET_LOG_MSG_AND_RETURN_IF_FALSE(
620 dim >= 0 && dim < input.dim(),
621 "dim %" PRId64 " out of range [0,%zd)",
622 dim,
623 input.dim());
624
625 const ssize_t dim_size = input.size(dim);
626 ET_LOG_MSG_AND_RETURN_IF_FALSE(
627 split_size >= 0,
628 "split_size %" PRId64 " must be non-negative",
629 split_size);
630 ET_LOG_MSG_AND_RETURN_IF_FALSE(
631 split_size > 0 || dim_size == 0,
632 "split_size is zero but input.size(%" PRId64 ") %zd is non-zero",
633 dim,
634 dim_size);
635
636 // Check the number of outputs.
637 //
638 // The specified dimension will be split into split_size-sized chunks, with
639 // the final chunk possibly being smaller. So, the expected output length is
640 // ceil(dim_size / split_size).
641 //
642 // E.g., splitting dim 0 of a [5,2] tensor with split_size 2 would produce
643 // three tensors with size [2,2], [2,2], [1,2].
644 int64_t remainder; // The size of the split dimension of the final out tensor.
645 if (split_size >= dim_size) {
646 // Note that this also handles the case where split_size == 0, avoiding a
647 // division by zero in the other branch. When dim_size == 0 && split_size ==
648 // 0, core PyTorch expects 1 output element.
649 ET_LOG_MSG_AND_RETURN_IF_FALSE(
650 out.size() == 1,
651 "Unexpected out.size() %zu: should be 1 because split_size %" PRId64
652 " >= input.size(%" PRId64 ") %zd",
653 out.size(),
654 split_size,
655 dim,
656 dim_size);
657 remainder = dim_size;
658 } else {
659 int64_t expected_out_len = (dim_size + split_size - 1) / split_size;
660 ET_LOG_MSG_AND_RETURN_IF_FALSE(
661 out.size() == expected_out_len,
662 "Unexpected out.size() %zu: ceil(input.size(%" PRId64
663 ")=%zd"
664 " / split_size=%" PRId64 ") is %" PRId64,
665 out.size(),
666 dim,
667 dim_size,
668 split_size,
669 expected_out_len);
670 remainder = dim_size % split_size;
671 if (remainder == 0) {
672 remainder = split_size;
673 }
674 }
675
676 // Validate each output.
677 for (size_t i = 0; i < out.size(); ++i) {
678 // All output dtypes must be the same.
679 ET_LOG_MSG_AND_RETURN_IF_FALSE(
680 out[i].scalar_type() == out[0].scalar_type(),
681 "out[%zu] dtype %" PRId8 " != out[0] dtype %" PRId8,
682 i,
683 static_cast<int8_t>(out[i].scalar_type()),
684 static_cast<int8_t>(out[0].scalar_type()));
685
686 // All outputs must have the same number of dimensions as the input.
687 ET_LOG_MSG_AND_RETURN_IF_FALSE(
688 out[i].dim() == input.dim(),
689 "out[%zu] dim %zd != input dim %zd",
690 i,
691 out[i].dim(),
692 input.dim());
693
694 // Check the shape of the output.
695 for (ssize_t d = 0; d < out[i].dim(); ++d) {
696 if (d == dim) {
697 // This is the split dimension, which may be different.
698 if (i < out.size() - 1) {
699 // All outputs except the final one: split dimension should be
700 // split_size.
701 ET_LOG_MSG_AND_RETURN_IF_FALSE(
702 out[i].size(d) == split_size,
703 "out[%zu].size(%zd) %zd != split_size %" PRId64,
704 i,
705 d,
706 out[i].size(d),
707 split_size);
708 } else {
709 // The final output: split dimension should be the remainder of
710 // split_size.
711 ET_LOG_MSG_AND_RETURN_IF_FALSE(
712 out[i].size(d) == remainder,
713 "out[%zu].size(%zd) %zd != remainder %" PRId64,
714 i,
715 d,
716 out[i].size(d),
717 remainder);
718 }
719 } else {
720 // Non-split output dimensions must be the same as the input dimension.
721 ET_LOG_AND_RETURN_IF_FALSE(
722 tensors_have_same_size_at_dims(out[i], d, input, d));
723 }
724 }
725 }
726
727 return true;
728 }
729
check_to_copy_args(const Tensor & input,bool non_blocking,exec_aten::optional<exec_aten::MemoryFormat> memory_format,Tensor & out)730 bool check_to_copy_args(
731 const Tensor& input,
732 bool non_blocking,
733 exec_aten::optional<exec_aten::MemoryFormat> memory_format,
734 Tensor& out) {
735 (void)input;
736 (void)out;
737
738 // Right now we only support blocking data transfer
739 ET_LOG_AND_RETURN_IF_FALSE(non_blocking == false);
740
741 // Right now we only focus on contiguous memory, memory_format shall be
742 // exec::aten::MemoryFormat::Contiguous or none.
743 ET_LOG_AND_RETURN_IF_FALSE(
744 !memory_format.has_value() ||
745 memory_format.value() == MemoryFormat::Contiguous);
746
747 return true;
748 }
749
check__to_dim_order_copy_args(const Tensor & input,bool non_blocking,exec_aten::OptionalArrayRef<int64_t> dim_order,Tensor & out)750 bool check__to_dim_order_copy_args(
751 const Tensor& input,
752 bool non_blocking,
753 exec_aten::OptionalArrayRef<int64_t> dim_order,
754 Tensor& out) {
755 // Right now we only support blocking data transfer
756 ET_LOG_AND_RETURN_IF_FALSE(non_blocking == false);
757
758 if (dim_order.has_value()) {
759 exec_aten::ArrayRef<int64_t> dim_order_ref = dim_order.value();
760
761 // dim order size shall equal to input dim
762 ET_LOG_AND_RETURN_IF_FALSE(dim_order_ref.size() == input.dim());
763
764 ET_LOG_AND_RETURN_IF_FALSE(
765 is_channels_last_dim_order(
766 dim_order.value().data(), dim_order.value().size()) ||
767 is_contiguous_dim_order(
768 dim_order.value().data(), dim_order.value().size()));
769
770 // Out tensor shall have same dim order as dim_order
771 auto out_dim_order = out.dim_order();
772 ET_LOG_AND_RETURN_IF_FALSE(out_dim_order.size() == dim_order_ref.size());
773 for (size_t i = 0; i < dim_order_ref.size(); i++) {
774 ET_LOG_AND_RETURN_IF_FALSE(out_dim_order[i] == dim_order_ref[i]);
775 }
776 } else { // dim_order is not set, preserve the dim order of input
777
778 // Out tensor shall have same dim order as input dim_order
779 auto out_dim_order = out.dim_order();
780 auto input_dim_order = input.dim_order();
781 ET_LOG_AND_RETURN_IF_FALSE(out_dim_order.size() == input_dim_order.size());
782 for (size_t i = 0; i < input_dim_order.size(); i++) {
783 ET_LOG_AND_RETURN_IF_FALSE(out_dim_order[i] == input_dim_order[i]);
784 }
785 }
786 return true;
787 }
788
check_unsqueeze_copy_args(const Tensor input,int64_t dim,const Tensor out)789 bool check_unsqueeze_copy_args(
790 const Tensor input,
791 int64_t dim,
792 const Tensor out) {
793 ET_LOG_AND_RETURN_IF_FALSE(dim >= 0);
794
795 // The input and out shall share same dtype
796 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(input, out));
797
798 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(out, dim));
799
800 // The shape of input and out shall obey the relationship:
801 // 1. input.dim() == out.dim()-1
802 // 2. input.size(i) == out.size(i) for all i < dim
803 // 3. input.size(i-1) == out.size(i) for all i >= dim
804 // 4. out.size(dim) == 1
805 ET_LOG_AND_RETURN_IF_FALSE(input.dim() == out.dim() - 1);
806
807 for (size_t d = 0; d < out.dim(); d++) {
808 auto dim_normalized = dim;
809 if (dim_normalized < 0) {
810 dim_normalized += out.dim();
811 }
812
813 if (d < dim_normalized) {
814 ET_LOG_MSG_AND_RETURN_IF_FALSE(
815 input.size(d) == out.size(d),
816 "input.size(%zu) %zd != out.size(%zu) %zd | dim = %" PRId64,
817 d,
818 input.size(d),
819 d,
820 out.size(d),
821 dim);
822 } else if (d > dim_normalized) {
823 ET_LOG_MSG_AND_RETURN_IF_FALSE(
824 input.size(d - 1) == out.size(d),
825 "input.size(%zu) %zd != out.size(%zu) %zd | dim = %" PRId64,
826 d - 1,
827 input.size(d),
828 d,
829 out.size(d),
830 dim);
831 } else { // d == dim
832 ET_LOG_MSG_AND_RETURN_IF_FALSE(
833 out.size(d) == 1,
834 "out.size(%zu) %zd shall equal 1 | dim = %" PRId64,
835 d,
836 out.size(d),
837 dim);
838 }
839 }
840
841 return true;
842 }
843
check_view_copy_args(const Tensor & self,exec_aten::ArrayRef<int64_t> size_int64_t,Tensor & out)844 bool check_view_copy_args(
845 const Tensor& self,
846 exec_aten::ArrayRef<int64_t> size_int64_t,
847 Tensor& out) {
848 ET_LOG_AND_RETURN_IF_FALSE(size_int64_t.size() == out.sizes().size());
849
850 // The input and out shall share same dtype and numel
851 ET_LOG_MSG_AND_RETURN_IF_FALSE(
852 self.numel() == out.numel(),
853 "self.numel() %zd != out.numel() %zd",
854 self.numel(),
855 out.numel());
856 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(self, out));
857
858 // The size of out should equal target size.
859 bool size_inferred = false;
860 for (int i = 0; i < size_int64_t.size(); i++) {
861 // If this value is -1 it implies that this dimension is inferred.
862 if (size_int64_t[i] == -1) {
863 ET_LOG_MSG_AND_RETURN_IF_FALSE(
864 !size_inferred, "Multiple dimensions cannot be inferred.");
865 size_inferred = true;
866 }
867 ET_LOG_AND_RETURN_IF_FALSE(
868 ((int64_t)out.sizes()[i] == size_int64_t[i]) ||
869 (size_int64_t[i] == -1));
870 }
871
872 return true;
873 }
874
get_view_copy_target_size(const Tensor input,exec_aten::ArrayRef<int64_t> size_int64_t,int64_t dim,exec_aten::SizesType * out_sizes)875 bool get_view_copy_target_size(
876 const Tensor input,
877 exec_aten::ArrayRef<int64_t> size_int64_t,
878 int64_t dim,
879 exec_aten::SizesType* out_sizes) {
880 size_t out_numels_without_minus_1 = 1;
881 int32_t minus_1_dim = -1;
882
883 ET_LOG_AND_RETURN_IF_FALSE(size_int64_t.size() == dim);
884
885 for (size_t i = 0; i < dim; ++i) {
886 if (size_int64_t[i] != -1) {
887 out_sizes[i] = static_cast<exec_aten::SizesType>(size_int64_t[i]);
888 out_numels_without_minus_1 = out_numels_without_minus_1 * size_int64_t[i];
889 } else {
890 // TODO(kimishpatel): Add test to hit this line
891 ET_LOG_MSG_AND_RETURN_IF_FALSE(
892 minus_1_dim == -1, "At most one view copy dim can be -1.");
893 minus_1_dim = i;
894 }
895 }
896 if (minus_1_dim >= 0) {
897 out_sizes[minus_1_dim] = input.numel() / out_numels_without_minus_1;
898 }
899
900 return true;
901 }
902
check_diagonal_copy_args(const Tensor & in,int64_t dim1,int64_t dim2,Tensor & out)903 bool check_diagonal_copy_args(
904 const Tensor& in,
905 int64_t dim1,
906 int64_t dim2,
907 Tensor& out) {
908 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
909 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 2));
910 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim1));
911 ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim2));
912 if (dim1 < 0) {
913 dim1 += nonzero_dim(in);
914 }
915 if (dim2 < 0) {
916 dim2 += nonzero_dim(in);
917 }
918 ET_LOG_AND_RETURN_IF_FALSE(dim1 != dim2);
919 return true;
920 }
921
get_diagonal_copy_out_target_size(const Tensor & in,int64_t offset,int64_t dim1,int64_t dim2,exec_aten::SizesType * out_sizes,size_t * out_ndim)922 void get_diagonal_copy_out_target_size(
923 const Tensor& in,
924 int64_t offset,
925 int64_t dim1,
926 int64_t dim2,
927 exec_aten::SizesType* out_sizes,
928 size_t* out_ndim) {
929 *out_ndim = in.dim() - 1;
930
931 if (dim1 < 0) {
932 dim1 += nonzero_dim(in);
933 }
934 if (dim2 < 0) {
935 dim2 += nonzero_dim(in);
936 }
937
938 size_t diagonal_size = 0;
939 if (offset >= 0) {
940 if (in.size(dim2) <= offset) {
941 diagonal_size = 0;
942 } else {
943 diagonal_size = std::min<size_t>(in.size(dim1), in.size(dim2) - offset);
944 }
945 } else {
946 if (in.size(dim1) <= -offset) {
947 diagonal_size = 0;
948 } else {
949 diagonal_size = std::min<size_t>(in.size(dim1) + offset, in.size(dim2));
950 }
951 }
952
953 size_t shift = 0;
954 for (size_t d = 0; d < in.dim(); ++d) {
955 if (d == dim1 || d == dim2) {
956 shift++;
957 } else {
958 out_sizes[d - shift] = in.size(d);
959 }
960 }
961 out_sizes[in.dim() - 2] = diagonal_size;
962 }
963
964 } // namespace executor
965 } // namespace torch
966