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::Scalar;
20 using exec_aten::ScalarType;
21 using exec_aten::Tensor;
22 using torch::executor::testing::SupportedFeatures;
23 using torch::executor::testing::TensorFactory;
24
25 class OpPixelShuffleOutTest : public OperatorTest {
26 protected:
op_pixel_shuffle_out(const Tensor & self,int64_t upscale_factor,Tensor & out)27 Tensor& op_pixel_shuffle_out(
28 const Tensor& self,
29 int64_t upscale_factor,
30 Tensor& out) {
31 return torch::executor::aten::pixel_shuffle_outf(
32 context_, self, upscale_factor, out);
33 }
34
35 template <ScalarType DTYPE_IN>
test_pixel_shuffle()36 void test_pixel_shuffle() {
37 TensorFactory<DTYPE_IN> tf_in;
38
39 const std::vector<int32_t> sizes = {1, 4, 2, 2};
40 const std::vector<int32_t> out_sizes = {1, 1, 4, 4};
41
42 // Destination for the pixel_shuffle.
43 Tensor out = tf_in.zeros(out_sizes);
44
45 op_pixel_shuffle_out(
46 tf_in.make(
47 sizes, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}),
48 2,
49 out);
50 EXPECT_TENSOR_EQ(
51 out,
52 // Pixel shuffle distributes channels amongst the spatial dimensions.
53 tf_in.make(
54 out_sizes, {0, 4, 1, 5, 8, 12, 9, 13, 2, 6, 3, 7, 10, 14, 11, 15}));
55 }
56 };
57
58 //
59 // Correctness Tests
60 //
61
62 /**
63 * Uses the function templates above to test all input dtypes.
64 */
TEST_F(OpPixelShuffleOutTest,AllRealDtypesSupported)65 TEST_F(OpPixelShuffleOutTest, AllRealDtypesSupported) {
66 #define ENUMERATE_TEST_ENTRY(ctype, dtype) \
67 test_pixel_shuffle<ScalarType::dtype>();
68
69 ET_FORALL_REAL_TYPES(ENUMERATE_TEST_ENTRY)
70
71 #undef ENUMERATE_TEST_ENTRY
72 }
73
TEST_F(OpPixelShuffleOutTest,LargerInputRank)74 TEST_F(OpPixelShuffleOutTest, LargerInputRank) {
75 TensorFactory<ScalarType::Int> tf;
76
77 // Pixel shuffle allows a 3D (or higher) input tensor, make sure the extra
78 // dimensions don't cause issues.
79 Tensor a = tf.ones(/*sizes=*/{1, 4, 1, 4, 2, 2});
80
81 const std::vector<int32_t> out_sizes = {1, 4, 1, 1, 4, 4};
82 Tensor out = tf.zeros(out_sizes);
83
84 op_pixel_shuffle_out(a, 2, out);
85 EXPECT_TENSOR_EQ(out, tf.ones(out_sizes));
86 }
87
88 // Mismatched shape tests.
TEST_F(OpPixelShuffleOutTest,InvalidInputChannelsDies)89 TEST_F(OpPixelShuffleOutTest, InvalidInputChannelsDies) {
90 TensorFactory<ScalarType::Int> tf;
91
92 // Input tensors with invalid shapes. 7 is not divisible by upsample_factor
93 // ** 2.
94 Tensor a = tf.ones(/*sizes=*/{1, 7, 4, 4});
95
96 Tensor out = tf.zeros(/*sizes=*/{1, 1, 8, 8});
97
98 // Using the wrong input shape should exit with an error code.
99 ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_shuffle_out(a, 2, out));
100 }
101
TEST_F(OpPixelShuffleOutTest,WrongInputRankDies)102 TEST_F(OpPixelShuffleOutTest, WrongInputRankDies) {
103 TensorFactory<ScalarType::Int> tf;
104
105 // Pixel shuffle requires a 3D or higher input tensor.
106 Tensor a = tf.ones(/*sizes=*/{1, 2});
107 Tensor out = tf.zeros(/*sizes=*/{1, 2});
108
109 // Using the wrong input shape should exit with an error code.
110 ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_shuffle_out(a, 2, out));
111 }
112
TEST_F(OpPixelShuffleOutTest,DifferentDtypeDies)113 TEST_F(OpPixelShuffleOutTest, DifferentDtypeDies) {
114 TensorFactory<ScalarType::Int> tf;
115 TensorFactory<ScalarType::Float> tf_float;
116
117 Tensor a = tf.ones(/*sizes=*/{1, 18, 4, 4});
118
119 // Pixel shuffle requires two tensors with the same dtype.
120 Tensor out = tf_float.zeros(/*sizes=*/{1, 2, 12, 12});
121
122 // Using the wrong output shape should exit with an error code.
123 ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_shuffle_out(a, 3, out));
124 }
125
TEST_F(OpPixelShuffleOutTest,NegativeUpscaleFactorDies)126 TEST_F(OpPixelShuffleOutTest, NegativeUpscaleFactorDies) {
127 TensorFactory<ScalarType::Int> tf;
128 Tensor a = tf.ones(/*sizes=*/{1, 18, 4, 4});
129 Tensor out = tf.zeros(/*sizes=*/{1, 2, 12, 12});
130 // Using a negative upscale factor should exit with an error code.
131 ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_shuffle_out(a, -3, out));
132 }
133