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::ScalarType;
21 using exec_aten::Tensor;
22 using torch::executor::testing::TensorFactory;
23
24 class OpSqueezeTest : public OperatorTest {
25 protected:
26 Tensor&
op_squeeze_copy_dim_out(const Tensor & self,int64_t dim,Tensor & out)27 op_squeeze_copy_dim_out(const Tensor& self, int64_t dim, Tensor& out) {
28 return torch::executor::aten::squeeze_copy_outf(context_, self, dim, out);
29 }
30 };
31
32 class OpSqueezeCopyDimsOutTest : public OperatorTest {
33 protected:
op_squeeze_copy_dims_out(const Tensor & self,exec_aten::ArrayRef<int64_t> dims,Tensor & out)34 Tensor& op_squeeze_copy_dims_out(
35 const Tensor& self,
36 exec_aten::ArrayRef<int64_t> dims,
37 Tensor& out) {
38 return torch::executor::aten::squeeze_copy_outf(context_, self, dims, out);
39 }
40 };
41
42 namespace {
43
TEST_F(OpSqueezeTest,DTypesMismatchDies)44 TEST_F(OpSqueezeTest, DTypesMismatchDies) {
45 TensorFactory<ScalarType::Int> tf_int;
46 TensorFactory<ScalarType::Double> tf_d;
47 Tensor t_in = tf_int.ones({2});
48 Tensor t_out = tf_d.ones({2});
49 int64_t dim = 0;
50
51 ET_EXPECT_KERNEL_FAILURE(context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
52 }
53
54 TEST_F(OpSqueezeTest, 0DTensorSqueeze) {
55 TensorFactory<ScalarType::Int> tf;
56 Tensor t_in = tf.ones({});
57 Tensor t_out = tf.zeros({});
58 Tensor t_expected = tf.ones({});
59 int64_t dim = 0;
60
61 op_squeeze_copy_dim_out(t_in, dim, t_out);
62 EXPECT_TENSOR_EQ(t_expected, t_out);
63 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
64 }
65
66 TEST_F(OpSqueezeTest, 0DTensorSqueezeInvalidDim1Dies) {
67 TensorFactory<ScalarType::Int> tf;
68 Tensor t_in = tf.ones({});
69 Tensor t_out = tf.ones({});
70 int64_t dim = 1;
71
72 ET_EXPECT_KERNEL_FAILURE(context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
73 }
74
75 TEST_F(OpSqueezeTest, 1DTensorSqueezeTo0D) {
76 TensorFactory<ScalarType::Int> tf;
77 Tensor t_in = tf.ones({1});
78 Tensor t_out = tf.make({}, {99});
79 Tensor t_expected = tf.make({}, {1});
80 int64_t dim = 0;
81
82 op_squeeze_copy_dim_out(t_in, dim, t_out);
83 EXPECT_TENSOR_EQ(t_expected, t_out);
84 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
85 }
86
87 TEST_F(OpSqueezeTest, 2DTensorSqueezeUnchange) {
88 TensorFactory<ScalarType::Int> tf;
89 Tensor t_in = tf.ones({2, 1});
90 Tensor t_out = tf.make({2, 1}, {4, 3});
91 Tensor t_expected = t_in;
92 int64_t dim = 0;
93
94 op_squeeze_copy_dim_out(t_in, dim, t_out);
95 EXPECT_TENSOR_EQ(t_expected, t_out);
96 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
97 }
98
99 TEST_F(OpSqueezeTest, 2DTensorSqueezeTo1D) {
100 TensorFactory<ScalarType::Int> tf;
101 Tensor t_in = tf.ones({2, 1});
102 Tensor t_out = tf.make({2}, {4, 3});
103 Tensor t_expected = tf.ones({2});
104 int64_t dim = 1;
105
106 op_squeeze_copy_dim_out(t_in, dim, t_out);
107 EXPECT_TENSOR_EQ(t_expected, t_out);
108 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
109 }
110
111 #ifndef USE_ATEN_LIB
112 TEST_F(OpSqueezeTest, 2DTensorSqueezeDownwardDimResizeOut) {
113 TensorFactory<ScalarType::Int> tf;
114 Tensor t_in = tf.ones({2, 1});
115 Tensor t_out = tf.zeros(
116 {4, 1},
117 torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); // okay to dwonward
118 // resize to (2, 1)
119 Tensor t_expected = tf.ones({2, 1});
120 int64_t dim = 0;
121
122 op_squeeze_copy_dim_out(t_in, dim, t_out);
123 EXPECT_TENSOR_EQ(t_expected, t_out);
124 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
125 }
126
127 TEST_F(OpSqueezeTest, 2DTensorSqueezeUpwardDimResizeOutDie) {
128 TensorFactory<ScalarType::Int> tf;
129 Tensor t_in = tf.ones({2, 1});
130 Tensor t_out = tf.zeros(
131 {1, 1},
132 torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); // can NOT upward
133 // resize 0th dim
134 Tensor t_expected = tf.ones({2, 1});
135 int64_t dim = 0;
136
137 ET_EXPECT_KERNEL_FAILURE(context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
138 }
139
140 TEST_F(OpSqueezeTest, 2DTensorSqueezeRemoveADimResizeOutDie) {
141 TensorFactory<ScalarType::Int> tf;
142 Tensor t_in = tf.ones({2, 1});
143 Tensor t_out = tf.zeros(
144 {2, 1, 3},
145 torch::executor::TensorShapeDynamism::
146 DYNAMIC_BOUND); // can NOT remove the 2nd dim via resizing
147 Tensor t_expected = tf.ones({2, 1});
148 int64_t dim = 0;
149
150 ET_EXPECT_KERNEL_FAILURE(context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
151 }
152
153 TEST_F(OpSqueezeTest, 2DTensorSqueezeAddDimsResizeOutDie) {
154 TensorFactory<ScalarType::Int> tf;
155 Tensor t_in = tf.ones({2, 1});
156 Tensor t_out = tf.zeros(
157 {2},
158 torch::executor::TensorShapeDynamism::
159 DYNAMIC_BOUND); // can NOT add dim(s) via resizing
160 Tensor t_expected = tf.ones({2, 1});
161 int64_t dim = 0;
162
163 ET_EXPECT_KERNEL_FAILURE(context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
164 }
165 #endif
166
TEST_F(OpSqueezeTest,TensorSqueeze)167 TEST_F(OpSqueezeTest, TensorSqueeze) {
168 TensorFactory<ScalarType::Int> tf;
169 Tensor t_in = tf.make({3, 1, 2, 1}, {1, 2, 3, 4, 5, 6});
170 Tensor t_out = tf.zeros({3, 2, 1});
171 Tensor t_expected = tf.make({3, 2, 1}, {1, 2, 3, 4, 5, 6});
172 int64_t dim = 1;
173
174 op_squeeze_copy_dim_out(t_in, dim, t_out);
175 EXPECT_TENSOR_EQ(t_expected, t_out);
176 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
177 }
178
TEST_F(OpSqueezeTest,TensorSqueezeNegativeDim)179 TEST_F(OpSqueezeTest, TensorSqueezeNegativeDim) {
180 TensorFactory<ScalarType::Int> tf;
181 Tensor t_in = tf.make({3, 1, 2, 1}, {1, 2, 3, 4, 5, 6});
182 Tensor t_out = tf.zeros({3, 2, 1});
183 Tensor t_expected = tf.make({3, 2, 1}, {1, 2, 3, 4, 5, 6});
184 int64_t dim = -3;
185
186 op_squeeze_copy_dim_out(t_in, dim, t_out);
187 EXPECT_TENSOR_EQ(t_expected, t_out);
188 EXPECT_TENSOR_DATA_EQ(t_expected, t_out);
189 }
190
TEST_F(OpSqueezeTest,TensorSqueezeInvaidDim)191 TEST_F(OpSqueezeTest, TensorSqueezeInvaidDim) {
192 TensorFactory<ScalarType::Int> tf;
193 Tensor t_in = tf.make({3, 1, 2, 1}, {1, 2, 3, 4, 5, 6});
194 Tensor t_out = tf.zeros({3, 2, 1});
195 Tensor t_expected = tf.make({3, 2, 1}, {1, 2, 3, 4, 5, 6});
196 std::vector<int64_t> invalid_dims = {t_in.dim(), -t_in.dim() - 1};
197
198 for (const auto dim : invalid_dims) {
199 ET_EXPECT_KERNEL_FAILURE(
200 context_, op_squeeze_copy_dim_out(t_in, dim, t_out));
201 }
202 }
203
204 /* %python
205 import torch
206 torch.manual_seed(0)
207 x = torch.rand(2, 1, 4)
208 res = torch.squeeze(x, 1)
209 op = "op_squeeze_copy_dim_out"
210 opt_extra_params = "1,"
211 dtype = "ScalarType::Float"
212 check = "EXPECT_TENSOR_EQ" */
213
TEST_F(OpSqueezeTest,DynamicShapeUpperBoundSameAsExpected)214 TEST_F(OpSqueezeTest, DynamicShapeUpperBoundSameAsExpected) {
215 /* %python
216 out_args = "{2, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
217 %rewrite(unary_op) */
218
219 TensorFactory<ScalarType::Float> tf;
220
221 Tensor x = tf.make(
222 {2, 1, 4},
223 {0.49625658988952637,
224 0.7682217955589294,
225 0.08847743272781372,
226 0.13203048706054688,
227 0.30742281675338745,
228 0.6340786814689636,
229 0.4900934100151062,
230 0.8964447379112244});
231 Tensor expected = tf.make(
232 {2, 4},
233 {0.49625658988952637,
234 0.7682217955589294,
235 0.08847743272781372,
236 0.13203048706054688,
237 0.30742281675338745,
238 0.6340786814689636,
239 0.4900934100151062,
240 0.8964447379112244});
241
242 Tensor out =
243 tf.zeros({2, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
244 op_squeeze_copy_dim_out(x, 1, out);
245 EXPECT_TENSOR_EQ(out, expected);
246 }
247
TEST_F(OpSqueezeTest,DynamicShapeUpperBoundLargerThanExpected)248 TEST_F(OpSqueezeTest, DynamicShapeUpperBoundLargerThanExpected) {
249 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
250 GTEST_SKIP() << "Dynamic shape not supported";
251 }
252 /* %python
253 out_args = "{5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
254 %rewrite(unary_op) */
255
256 TensorFactory<ScalarType::Float> tf;
257
258 Tensor x = tf.make(
259 {2, 1, 4},
260 {0.49625658988952637,
261 0.7682217955589294,
262 0.08847743272781372,
263 0.13203048706054688,
264 0.30742281675338745,
265 0.6340786814689636,
266 0.4900934100151062,
267 0.8964447379112244});
268 Tensor expected = tf.make(
269 {2, 4},
270 {0.49625658988952637,
271 0.7682217955589294,
272 0.08847743272781372,
273 0.13203048706054688,
274 0.30742281675338745,
275 0.6340786814689636,
276 0.4900934100151062,
277 0.8964447379112244});
278
279 Tensor out =
280 tf.zeros({5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
281 op_squeeze_copy_dim_out(x, 1, out);
282 EXPECT_TENSOR_EQ(out, expected);
283 }
284
TEST_F(OpSqueezeTest,DynamicShapeUnbound)285 TEST_F(OpSqueezeTest, DynamicShapeUnbound) {
286 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
287 GTEST_SKIP() << "Dynamic shape not supported";
288 }
289 /* %python
290 out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
291 %rewrite(unary_op) */
292
293 TensorFactory<ScalarType::Float> tf;
294
295 Tensor x = tf.make(
296 {2, 1, 4},
297 {0.49625658988952637,
298 0.7682217955589294,
299 0.08847743272781372,
300 0.13203048706054688,
301 0.30742281675338745,
302 0.6340786814689636,
303 0.4900934100151062,
304 0.8964447379112244});
305 Tensor expected = tf.make(
306 {2, 4},
307 {0.49625658988952637,
308 0.7682217955589294,
309 0.08847743272781372,
310 0.13203048706054688,
311 0.30742281675338745,
312 0.6340786814689636,
313 0.4900934100151062,
314 0.8964447379112244});
315
316 Tensor out =
317 tf.zeros({1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
318 op_squeeze_copy_dim_out(x, 1, out);
319 EXPECT_TENSOR_EQ(out, expected);
320 }
321
322 } // namespace
323
TEST_F(OpSqueezeCopyDimsOutTest,SanityTest4D)324 TEST_F(OpSqueezeCopyDimsOutTest, SanityTest4D) {
325 torch::executor::testing::TensorFactory<exec_aten::ScalarType::Float> tfFloat;
326
327 exec_aten::Tensor self = tfFloat.make(
328 {1, 2, 1, 5},
329 {-26.5,
330 5.75,
331 95.75,
332 92.625,
333 -97.25,
334 65.5,
335 -92.25,
336 -67.625,
337 54.75,
338 27.125});
339 ::std::vector<int64_t> dim_vec = {0, 2};
340 exec_aten::ArrayRef<int64_t> dim =
341 exec_aten::ArrayRef<int64_t>(dim_vec.data(), dim_vec.size());
342 exec_aten::Tensor out = tfFloat.zeros({2, 5});
343 exec_aten::Tensor out_expected = tfFloat.make(
344 {2, 5},
345 {-26.5,
346 5.75,
347 95.75,
348 92.625,
349 -97.25,
350 65.5,
351 -92.25,
352 -67.625,
353 54.75,
354 27.125});
355 op_squeeze_copy_dims_out(self, dim, out);
356 EXPECT_TENSOR_CLOSE(out, out_expected);
357 }
358
TEST_F(OpSqueezeCopyDimsOutTest,SanityCheck5D)359 TEST_F(OpSqueezeCopyDimsOutTest, SanityCheck5D) {
360 torch::executor::testing::TensorFactory<exec_aten::ScalarType::Float> tfFloat;
361
362 exec_aten::Tensor self = tfFloat.make(
363 {1, 2, 1, 5, 4},
364 {-73.5, -67.625, -54.375, 51.625, -11.125, -28.625, -40.75, 45.625,
365 84.375, 65.625, 95.125, -47.125, -21.25, 32.25, -86.125, 55.875,
366 -62.25, 47.125, -71.875, 43.0, 47.875, -73.375, 97.75, 69.25,
367 64.125, -59.875, 59.75, -52.25, 59.5, 44.875, -51.25, 20.875,
368 -67.0, 32.5, -26.625, 83.75, 45.5, 85.5, -92.875, 60.0});
369 ::std::vector<int64_t> dim_vec = {0, 3, 2, 1};
370 exec_aten::ArrayRef<int64_t> dim =
371 exec_aten::ArrayRef<int64_t>(dim_vec.data(), dim_vec.size());
372 exec_aten::Tensor out = tfFloat.zeros({2, 5, 4});
373 exec_aten::Tensor out_expected = tfFloat.make(
374 {2, 5, 4},
375 {-73.5, -67.625, -54.375, 51.625, -11.125, -28.625, -40.75, 45.625,
376 84.375, 65.625, 95.125, -47.125, -21.25, 32.25, -86.125, 55.875,
377 -62.25, 47.125, -71.875, 43.0, 47.875, -73.375, 97.75, 69.25,
378 64.125, -59.875, 59.75, -52.25, 59.5, 44.875, -51.25, 20.875,
379 -67.0, 32.5, -26.625, 83.75, 45.5, 85.5, -92.875, 60.0});
380 op_squeeze_copy_dims_out(self, dim, out);
381 EXPECT_TENSOR_CLOSE(out, out_expected);
382 }
383
TEST_F(OpSqueezeCopyDimsOutTest,SanityCheck5DUnchanged)384 TEST_F(OpSqueezeCopyDimsOutTest, SanityCheck5DUnchanged) {
385 torch::executor::testing::TensorFactory<exec_aten::ScalarType::Float> tfFloat;
386
387 exec_aten::Tensor self = tfFloat.make(
388 {1, 2, 1, 5, 4},
389 {-0.375, -40.125, 5.75, 21.25, -34.875, -19.375, 15.75, -60.75,
390 -41.75, 53.125, -76.0, -64.25, -84.5, -37.25, -39.125, 22.875,
391 -69.0, 30.25, -21.25, 85.5, 8.875, 41.625, 12.375, -1.125,
392 -14.875, 78.5, 43.0, -78.625, -58.625, -58.375, 47.5, -67.375,
393 -82.375, 35.0, 83.25, 49.625, -9.875, -46.75, 17.875, -68.375});
394 ::std::vector<int64_t> dim_vec = {1, 4, 3};
395 exec_aten::ArrayRef<int64_t> dim =
396 exec_aten::ArrayRef<int64_t>(dim_vec.data(), dim_vec.size());
397 exec_aten::Tensor out = tfFloat.zeros({1, 2, 1, 5, 4});
398 exec_aten::Tensor out_expected = tfFloat.make(
399 {1, 2, 1, 5, 4},
400 {-0.375, -40.125, 5.75, 21.25, -34.875, -19.375, 15.75, -60.75,
401 -41.75, 53.125, -76.0, -64.25, -84.5, -37.25, -39.125, 22.875,
402 -69.0, 30.25, -21.25, 85.5, 8.875, 41.625, 12.375, -1.125,
403 -14.875, 78.5, 43.0, -78.625, -58.625, -58.375, 47.5, -67.375,
404 -82.375, 35.0, 83.25, 49.625, -9.875, -46.75, 17.875, -68.375});
405 op_squeeze_copy_dims_out(self, dim, out);
406 EXPECT_TENSOR_CLOSE(out, out_expected);
407 }
408