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