xref: /aosp_15_r20/external/executorch/kernels/portable/cpu/util/copy_ops_util.cpp (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
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