xref: /aosp_15_r20/external/executorch/kernels/test/op_permute_copy_test.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 <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
10 #include <executorch/kernels/test/TestUtil.h>
11 #include <executorch/kernels/test/supported_features.h>
12 #include <executorch/runtime/core/exec_aten/exec_aten.h>
13 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
14 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
15 
16 #include <gtest/gtest.h>
17 
18 using namespace ::testing;
19 using exec_aten::ArrayRef;
20 using exec_aten::IntArrayRef;
21 using exec_aten::ScalarType;
22 using exec_aten::Tensor;
23 using torch::executor::testing::TensorFactory;
24 
25 class OpPermuteCopyTest : public OperatorTest {
26  protected:
27   Tensor&
op_permute_copy_out(const Tensor & self,IntArrayRef dims,Tensor & out)28   op_permute_copy_out(const Tensor& self, IntArrayRef dims, Tensor& out) {
29     return torch::executor::aten::permute_copy_outf(context_, self, dims, out);
30   }
31 };
32 
TEST_F(OpPermuteCopyTest,OneDPermute)33 TEST_F(OpPermuteCopyTest, OneDPermute) {
34   TensorFactory<ScalarType::Int> tf;
35 
36   const std::vector<int64_t> new_dim = {0};
37 
38   const std::vector<int32_t> sizes = {2};
39   Tensor t_int = tf.make(sizes, {1, 2});
40 
41   Tensor out = tf.zeros(sizes);
42 
43   op_permute_copy_out(
44       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
45   EXPECT_TENSOR_EQ(out, tf.make(sizes, {1, 2}));
46 }
47 
TEST_F(OpPermuteCopyTest,PermuteWithNoDataReorder)48 TEST_F(OpPermuteCopyTest, PermuteWithNoDataReorder) {
49   TensorFactory<ScalarType::Int> tf;
50 
51   const std::vector<int64_t> new_dim = {1, 0, 2};
52 
53   // clang-format off
54   Tensor t_int = tf.make({1,4,5}, {
55     0,  1,  2,  3,  4,
56     5,  6,  7,  8,  9,
57     10, 11, 12, 13, 14,
58     15, 16, 17,18, 19});
59   // clang-format on
60 
61   const std::vector<int32_t> new_sizes = {4, 1, 5};
62   Tensor out = tf.zeros(new_sizes);
63 
64   op_permute_copy_out(
65       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
66   // clang-format off
67   EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
68     0,  1,  2,  3,  4,
69     5,  6,  7,  8,  9,
70     10, 11, 12, 13, 14,
71     15, 16, 17, 18, 19}));
72   // clang-format on
73 }
74 
TEST_F(OpPermuteCopyTest,TwoDPermute)75 TEST_F(OpPermuteCopyTest, TwoDPermute) {
76   TensorFactory<ScalarType::Int> tf;
77 
78   const std::vector<int64_t> new_dim = {1, 0};
79 
80   // clang-format off
81   Tensor t_int = tf.make({2, 3}, {
82     // 2x3 data block
83     0, 1, 2,
84     3, 4, 5
85   });
86   // clang-format on
87 
88   const std::vector<int32_t> new_sizes = {3, 2};
89   Tensor out = tf.zeros(new_sizes);
90 
91   op_permute_copy_out(
92       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
93   // clang-format off
94   EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
95     // 3x2 data block
96     0, 3,
97     1, 4,
98     2, 5
99   }));
100   // clang-format on
101 }
102 
TEST_F(OpPermuteCopyTest,ThreeDPermute)103 TEST_F(OpPermuteCopyTest, ThreeDPermute) {
104   TensorFactory<ScalarType::Int> tf;
105 
106   const std::vector<int64_t> new_dim = {2, 0, 1};
107 
108   // clang-format off
109   Tensor t_int = tf.make({2, 1, 3}, {
110     // 2 1x3 data blocks
111     0, 1, 2,
112 
113     3, 4, 5
114   });
115   // clang-format on
116 
117   const std::vector<int32_t> new_sizes = {3, 2, 1};
118   Tensor out = tf.zeros(new_sizes);
119 
120   op_permute_copy_out(
121       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
122   // clang-format off
123   EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
124   // 3 2x1 data blocks
125     0,
126     3,
127 
128     1,
129     4,
130 
131     2,
132     5
133   }));
134   // clang-format on
135 }
136 
TEST_F(OpPermuteCopyTest,FourDPermute)137 TEST_F(OpPermuteCopyTest, FourDPermute) {
138   TensorFactory<ScalarType::Int> tf;
139 
140   const std::vector<int64_t> new_dim = {0, 3, 2, 1};
141 
142   // clang-format off
143   Tensor t_int = tf.make(
144       {2, 3, 3, 4},
145       // 2 groupings of 3 3x4 data blocks
146       {0,  1,  2,  3,
147        4,  5,  6,  7,
148        8,  9,  10, 11,
149 
150        12, 13, 14, 15,
151        16, 17, 18, 19,
152        20, 21, 22, 23,
153 
154        24, 25, 26, 27,
155        28, 29, 30, 31,
156        32, 33, 34, 35,
157 
158 
159        36, 37, 38, 39,
160        40, 41, 42, 43,
161        44, 45, 46, 47,
162 
163        48, 49, 50, 51,
164        52, 53, 54, 55,
165        56, 57, 58, 59,
166 
167        60, 61, 62, 63,
168        64, 65, 66, 67,
169        68, 69, 70, 71});
170   // clang-format on
171 
172   const std::vector<int32_t> new_sizes = {2, 4, 3, 3};
173   Tensor out = tf.zeros(new_sizes);
174 
175   // Long results like this are gotten by running torch.permute in a notebook
176   // and copy pasting the result. Ex:
177   //   import torch
178   //   x = torch.arange(0, 72, 1).view(2, 3, 3, 4).contiguous()
179   //   print(x.flatten().contiguous())
180   //   z = torch.permute(x, (0, 3, 2, 1))
181   //   print(z.flatten().contiguous())
182   op_permute_copy_out(
183       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
184   EXPECT_TENSOR_EQ(
185       out,
186       // clang-format off
187       tf.make(new_sizes, {
188         // 2 groupings of 4 3x3 data blocks
189         0, 12, 24,
190         4, 16, 28,
191         8, 20, 32,
192 
193         1, 13, 25,
194         5, 17, 29,
195         9, 21, 33,
196 
197         2, 14, 26,
198         6, 18, 30,
199         10, 22, 34,
200 
201         3, 15, 27,
202         7, 19, 31,
203         11, 23, 35,
204 
205 
206         36, 48, 60,
207         40, 52, 64,
208         44, 56, 68,
209 
210         37, 49, 61,
211         41, 53, 65,
212         45, 57, 69,
213 
214         38, 50, 62,
215         42, 54, 66,
216         46, 58, 70,
217 
218         39, 51, 63,
219         43, 55, 67,
220         47, 59, 71}));
221   // clang-format on
222 }
223 
TEST_F(OpPermuteCopyTest,FiveDPermute)224 TEST_F(OpPermuteCopyTest, FiveDPermute) {
225   TensorFactory<ScalarType::Int> tf;
226 
227   const std::vector<int64_t> new_dim = {4, 3, 2, 1, 0};
228 
229   const std::vector<int32_t> sizes = {2, 2, 2, 2, 2};
230   // clang-format off
231   Tensor t_int = tf.make(
232       sizes, {
233         0,  1,
234         2,  3,
235 
236         4,  5,
237         6,  7,
238 
239 
240         8,  9,
241         10, 11,
242 
243         12, 13,
244         14, 15,
245 
246 
247         16, 17,
248         18, 19,
249 
250         20, 21,
251         22, 23,
252 
253 
254         24, 25,
255         26, 27,
256 
257         28, 29,
258         30, 31});
259   // clang-format on
260 
261   Tensor out = tf.zeros(sizes);
262 
263   op_permute_copy_out(
264       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
265   // Long results like this are gotten by running torch.permute in a notebook
266   // and copy pasting the result. Ex:
267   //   import torch
268   //   x = torch.arange(0, 32, 1).view(2, 2, 2, 2, 2).contiguous()
269   //   print(x.flatten().contiguous())
270   //   z = torch.permute(x, (4, 3, 2, 1, 0))
271   //   print(z.flatten().contiguous())
272   // clang-format off
273   EXPECT_TENSOR_EQ(
274       out, tf.make(sizes, {
275         0,  16,
276         8,  24,
277 
278         4,  20,
279         12, 28,
280 
281 
282         2,  18,
283         10, 26,
284 
285         6,  22,
286         14, 30,
287 
288 
289         1,  17,
290         9,  25,
291 
292         5,  21,
293         13, 29,
294 
295 
296         3,  19,
297         11, 27,
298 
299         7,  23,
300         15, 31}));
301   // clang-format on
302 }
303 
TEST_F(OpPermuteCopyTest,AllDimensionsSizeOne)304 TEST_F(OpPermuteCopyTest, AllDimensionsSizeOne) {
305   TensorFactory<ScalarType::Int> tf;
306 
307   const std::vector<int64_t> new_dim = {4, 3, 2, 1, 0};
308 
309   const std::vector<int32_t> sizes = {1, 1, 1, 1, 1};
310   Tensor t_int = tf.make(sizes, {1});
311 
312   Tensor out = tf.zeros(sizes);
313 
314   op_permute_copy_out(
315       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out);
316   EXPECT_TENSOR_EQ(out, tf.make(sizes, {1}));
317 }
318 
TEST_F(OpPermuteCopyTest,DupeDimensionPos)319 TEST_F(OpPermuteCopyTest, DupeDimensionPos) {
320   TensorFactory<ScalarType::Int> tf;
321 
322   const std::vector<int64_t> new_dim = {0, 1, 1};
323 
324   const std::vector<int32_t> sizes = {1, 1, 1};
325   Tensor t_int = tf.make(sizes, {1});
326 
327   Tensor out = tf.zeros(sizes);
328 
329   ET_EXPECT_KERNEL_FAILURE(
330       context_,
331       op_permute_copy_out(
332           t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
333 }
334 
TEST_F(OpPermuteCopyTest,DupeDimensionPos2)335 TEST_F(OpPermuteCopyTest, DupeDimensionPos2) {
336   TensorFactory<ScalarType::Int> tf;
337 
338   const std::vector<int64_t> new_dim = {1, 1, 1};
339 
340   const std::vector<int32_t> sizes = {1, 1, 1};
341   Tensor t_int = tf.make(sizes, {1});
342 
343   Tensor out = tf.zeros(sizes);
344 
345   ET_EXPECT_KERNEL_FAILURE(
346       context_,
347       op_permute_copy_out(
348           t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
349 }
350 
TEST_F(OpPermuteCopyTest,DupeDimensionNeg)351 TEST_F(OpPermuteCopyTest, DupeDimensionNeg) {
352   TensorFactory<ScalarType::Int> tf;
353 
354   const std::vector<int64_t> new_dim = {0, 1, -2};
355 
356   const std::vector<int32_t> sizes = {1, 1, 1};
357   Tensor t_int = tf.make(sizes, {1});
358 
359   Tensor out = tf.zeros(sizes);
360 
361   ET_EXPECT_KERNEL_FAILURE(
362       context_,
363       op_permute_copy_out(
364           t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
365 }
366 
TEST_F(OpPermuteCopyTest,DupeDimensionNeg2)367 TEST_F(OpPermuteCopyTest, DupeDimensionNeg2) {
368   TensorFactory<ScalarType::Int> tf;
369 
370   const std::vector<int64_t> new_dim = {0, 1, -5};
371 
372   const std::vector<int32_t> sizes = {1, 1, 1};
373   Tensor t_int = tf.make(sizes, {1});
374 
375   Tensor out = tf.zeros(sizes);
376 
377   ET_EXPECT_KERNEL_FAILURE(
378       context_,
379       op_permute_copy_out(
380           t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
381 }
382 
TEST_F(OpPermuteCopyTest,MismatchDim)383 TEST_F(OpPermuteCopyTest, MismatchDim) {
384   TensorFactory<ScalarType::Int> tf;
385 
386   const std::vector<int64_t> new_dim = {0, 1, 2};
387 
388   const std::vector<int32_t> sizes = {1, 1};
389   Tensor t_int = tf.make(sizes, {1});
390 
391   Tensor out = tf.zeros(sizes);
392 
393   ET_EXPECT_KERNEL_FAILURE(
394       context_,
395       op_permute_copy_out(
396           t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
397 }
398 
399 /* %python
400 import torch
401 torch.manual_seed(0)
402 x = torch.randint(10, (2, 3, 4))
403 res = torch.permute(x, (2, 0, 1))
404 op = "op_permute_copy_out"
405 opt_setup_params = f"""
406   {declare_array_ref([2, 0, 1], "int64_t", "perm_aref")}
407 """
408 opt_extra_params = "perm_aref,"
409 dtype = "ScalarType::Int"
410 check = "EXPECT_TENSOR_EQ" */
411 
TEST_F(OpPermuteCopyTest,DynamicShapeUpperBoundSameAsExpected)412 TEST_F(OpPermuteCopyTest, DynamicShapeUpperBoundSameAsExpected) {
413   /* %python
414   out_args = "{4, 2, 3}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
415   %rewrite(unary_op) */
416 
417   TensorFactory<ScalarType::Int> tf;
418 
419   Tensor x = tf.make({2, 3, 4}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6,
420                                  6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4});
421   Tensor expected = tf.make({4, 2, 3}, {4, 3, 7, 6, 6, 6, 9, 9, 3, 9, 8, 9,
422                                         3, 7, 1, 8, 4, 1, 0, 3, 6, 6, 3, 4});
423 
424   std::vector<int64_t> perm_arefv = {2, 0, 1};
425   ArrayRef<int64_t> perm_aref(perm_arefv.data(), perm_arefv.size());
426 
427   Tensor out =
428       tf.zeros({4, 2, 3}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
429   op_permute_copy_out(x, perm_aref, out);
430   EXPECT_TENSOR_EQ(out, expected);
431 }
432 
TEST_F(OpPermuteCopyTest,DynamicShapeUpperBoundLargerThanExpected)433 TEST_F(OpPermuteCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
434   /* %python
435   out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
436   %rewrite(unary_op) */
437 
438   TensorFactory<ScalarType::Int> tf;
439 
440   Tensor x = tf.make({2, 3, 4}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6,
441                                  6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4});
442   Tensor expected = tf.make({4, 2, 3}, {4, 3, 7, 6, 6, 6, 9, 9, 3, 9, 8, 9,
443                                         3, 7, 1, 8, 4, 1, 0, 3, 6, 6, 3, 4});
444 
445   std::vector<int64_t> perm_arefv = {2, 0, 1};
446   ArrayRef<int64_t> perm_aref(perm_arefv.data(), perm_arefv.size());
447 
448   Tensor out =
449       tf.zeros({5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
450   op_permute_copy_out(x, perm_aref, out);
451   EXPECT_TENSOR_EQ(out, expected);
452 }
453 
TEST_F(OpPermuteCopyTest,DynamicShapeUnbound)454 TEST_F(OpPermuteCopyTest, DynamicShapeUnbound) {
455   if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
456     GTEST_SKIP() << "Dynamic shape unbound not supported";
457   }
458   /* %python
459   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
460   %rewrite(unary_op) */
461 
462   TensorFactory<ScalarType::Int> tf;
463 
464   Tensor x = tf.make({2, 3, 4}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6,
465                                  6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4});
466   Tensor expected = tf.make({4, 2, 3}, {4, 3, 7, 6, 6, 6, 9, 9, 3, 9, 8, 9,
467                                         3, 7, 1, 8, 4, 1, 0, 3, 6, 6, 3, 4});
468 
469   std::vector<int64_t> perm_arefv = {2, 0, 1};
470   ArrayRef<int64_t> perm_aref(perm_arefv.data(), perm_arefv.size());
471 
472   Tensor out = tf.zeros(
473       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
474   op_permute_copy_out(x, perm_aref, out);
475   EXPECT_TENSOR_EQ(out, expected);
476 }
477