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 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
16 #include <executorch/runtime/platform/runtime.h>
17
18 #include <gtest/gtest.h>
19 #include <cstdint>
20 #include <limits>
21
22 using namespace ::testing;
23 using exec_aten::ArrayRef;
24 using exec_aten::Scalar;
25 using exec_aten::ScalarType;
26 using exec_aten::Tensor;
27
28 using torch::executor::testing::TensorFactory;
29
30 class OpArangeOutTest : public OperatorTest {
31 protected:
op_arange_out(const Scalar & end,Tensor & out)32 Tensor& op_arange_out(const Scalar& end, Tensor& out) {
33 return torch::executor::aten::arange_outf(context_, end, out);
34 }
35
36 template <class CTYPE, exec_aten::ScalarType DTYPE>
test_arange_dtype()37 void test_arange_dtype() {
38 TensorFactory<DTYPE> tf;
39
40 Scalar end = Scalar(static_cast<CTYPE>(10));
41
42 Tensor out = tf.zeros({10});
43
44 Tensor ret = op_arange_out(end, out);
45
46 // Should always return the provided out Tensor.
47 EXPECT_TENSOR_EQ(ret, out);
48
49 // Expected tensor, filled with 0, 1, ..., 9
50 Tensor expected = tf.make({10}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
51
52 EXPECT_TENSOR_EQ(out, expected);
53 }
54 };
55
56 class OpArangeStartOutTest : public OperatorTest {
57 protected:
op_arange_start_out(const Scalar & start,const Scalar & end,const Scalar & step,Tensor & out)58 Tensor& op_arange_start_out(
59 const Scalar& start,
60 const Scalar& end,
61 const Scalar& step,
62 Tensor& out) {
63 return torch::executor::aten::arange_outf(context_, start, end, step, out);
64 }
65
66 template <class CTYPE, exec_aten::ScalarType DTYPE>
test_arange_start_dtype()67 void test_arange_start_dtype() {
68 TensorFactory<DTYPE> tf;
69
70 Scalar start = Scalar(static_cast<CTYPE>(0));
71 Scalar end = Scalar(static_cast<CTYPE>(10));
72 Scalar step = Scalar(static_cast<CTYPE>(1));
73
74 Tensor out = tf.zeros({10});
75
76 Tensor ret = op_arange_start_out(start, end, step, out);
77
78 // Should always return the provided out Tensor.
79 EXPECT_TENSOR_EQ(ret, out);
80
81 // Expected tensor, filled with 0, 1, ..., 9
82 Tensor expected = tf.make({10}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
83
84 EXPECT_TENSOR_EQ(out, expected);
85 }
86 };
87
88 /// A generic smoke test that works for any dtype that supports zeros().
TEST_F(OpArangeOutTest,AllRealDtypesSupported)89 TEST_F(OpArangeOutTest, AllRealDtypesSupported) {
90 #define TEST_ENTRY(ctype, dtype) test_arange_dtype<ctype, ScalarType::dtype>();
91 ET_FORALL_REAL_TYPES(TEST_ENTRY);
92 #undef TEST_ENTRY
93 }
94
TEST_F(OpArangeOutTest,FloatNumberNotEqualIntSupport)95 TEST_F(OpArangeOutTest, FloatNumberNotEqualIntSupport) {
96 TensorFactory<ScalarType::Float> tf;
97
98 // end = any floating point number between [a, a+1) where a is an arbitrary
99 // integer should have same result as end = a. So here arage(end = 5.5) ==
100 // arange(5)
101 Scalar end = Scalar(5.5);
102
103 Tensor out = tf.zeros({6});
104
105 Tensor ret = op_arange_out(end, out);
106
107 // Should always return the provided out Tensor.
108 EXPECT_TENSOR_EQ(ret, out);
109
110 // Expected tensor, equal
111 Tensor expected = tf.make({6}, {0.0, 1.0, 2.0, 3.0, 4.0, 5.0});
112
113 EXPECT_TENSOR_EQ(out, expected);
114 }
115
TEST_F(OpArangeOutTest,OutDimUnsupportedDie)116 TEST_F(OpArangeOutTest, OutDimUnsupportedDie) {
117 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
118 GTEST_SKIP() << "ATen kernel can handle mismatched out dim";
119 }
120 TensorFactory<ScalarType::Float> tf;
121
122 Scalar end = Scalar(5);
123
124 Tensor out = tf.zeros({5, 1});
125
126 // out.dim() should be 1, not 2
127 ET_EXPECT_KERNEL_FAILURE(context_, op_arange_out(end, out));
128 }
129
TEST_F(OpArangeOutTest,DynamicShapeUpperBoundSameAsExpected)130 TEST_F(OpArangeOutTest, DynamicShapeUpperBoundSameAsExpected) {
131 TensorFactory<ScalarType::Float> tf;
132
133 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
134
135 Tensor out =
136 tf.zeros({5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
137 Tensor ret = op_arange_out(Scalar(5), out);
138 EXPECT_TENSOR_CLOSE(out, expected_result);
139 }
140
TEST_F(OpArangeOutTest,DynamicShapeUpperBoundLargerThanExpected)141 TEST_F(OpArangeOutTest, DynamicShapeUpperBoundLargerThanExpected) {
142 TensorFactory<ScalarType::Float> tf;
143
144 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
145
146 Tensor out =
147 tf.zeros({10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
148 Tensor ret = op_arange_out(Scalar(5), out);
149 EXPECT_TENSOR_CLOSE(out, expected_result);
150 }
151
TEST_F(OpArangeOutTest,DynamicShapeUnbound)152 TEST_F(OpArangeOutTest, DynamicShapeUnbound) {
153 if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
154 GTEST_SKIP() << "Dynamic Unbound not supported";
155 }
156 TensorFactory<ScalarType::Float> tf;
157
158 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
159
160 Tensor out =
161 tf.zeros({1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
162 Tensor ret = op_arange_out(Scalar(5), out);
163 EXPECT_TENSOR_CLOSE(out, expected_result);
164 }
165
166 /// A generic smoke test that works for any dtype that supports zeros().
TEST_F(OpArangeStartOutTest,AllRealDtypesSupported)167 TEST_F(OpArangeStartOutTest, AllRealDtypesSupported) {
168 #define TEST_ENTRY(ctype, dtype) \
169 test_arange_start_dtype<ctype, ScalarType::dtype>();
170 ET_FORALL_REAL_TYPES(TEST_ENTRY);
171 #undef TEST_ENTRY
172 }
173
TEST_F(OpArangeStartOutTest,FloatNumberNotEqualIntSupport)174 TEST_F(OpArangeStartOutTest, FloatNumberNotEqualIntSupport) {
175 TensorFactory<ScalarType::Float> tf;
176
177 // Tested in bento:
178 // import torch
179 // torch.arange(5.5)
180 // >> tensor([0., 1., 2., 3., 4., 5.])
181 Scalar start = Scalar(0);
182 Scalar end = Scalar(5.5);
183 Scalar step = Scalar(1);
184
185 Tensor out = tf.zeros({6});
186
187 Tensor ret = op_arange_start_out(start, end, step, out);
188
189 // Should always return the provided out Tensor.
190 EXPECT_TENSOR_EQ(ret, out);
191
192 // Expected tensor, equal
193 Tensor expected = tf.make({6}, {0.0, 1.0, 2.0, 3.0, 4.0, 5.0});
194
195 EXPECT_TENSOR_EQ(out, expected);
196 }
197
TEST_F(OpArangeStartOutTest,OutDimUnsupportedDie)198 TEST_F(OpArangeStartOutTest, OutDimUnsupportedDie) {
199 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
200 GTEST_SKIP() << "ATen kernel can handle mismatched out dim";
201 }
202 TensorFactory<ScalarType::Float> tf;
203
204 Scalar start = Scalar(0);
205 Scalar end = Scalar(5);
206 Scalar step = Scalar(1);
207
208 Tensor out = tf.zeros({5, 1});
209
210 // out.dim() should be 1, not 2
211 ET_EXPECT_KERNEL_FAILURE(
212 context_, op_arange_start_out(start, end, step, out));
213 }
214
TEST_F(OpArangeStartOutTest,DynamicShapeUpperBoundSameAsExpected)215 TEST_F(OpArangeStartOutTest, DynamicShapeUpperBoundSameAsExpected) {
216 TensorFactory<ScalarType::Float> tf;
217
218 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
219
220 Tensor out =
221 tf.zeros({5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
222 Tensor ret = op_arange_start_out(Scalar(0), Scalar(5), Scalar(1), out);
223 EXPECT_TENSOR_CLOSE(out, expected_result);
224 }
225
TEST_F(OpArangeStartOutTest,DynamicShapeUpperBoundLargerThanExpected)226 TEST_F(OpArangeStartOutTest, DynamicShapeUpperBoundLargerThanExpected) {
227 TensorFactory<ScalarType::Float> tf;
228
229 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
230
231 Tensor out =
232 tf.zeros({10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
233 Tensor ret = op_arange_start_out(Scalar(0), Scalar(5), Scalar(1), out);
234 EXPECT_TENSOR_CLOSE(out, expected_result);
235 }
236
TEST_F(OpArangeStartOutTest,DynamicShapeUnbound)237 TEST_F(OpArangeStartOutTest, DynamicShapeUnbound) {
238 if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
239 GTEST_SKIP() << "Dynamic Unbound not supported";
240 }
241 TensorFactory<ScalarType::Float> tf;
242
243 Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
244
245 Tensor out =
246 tf.zeros({1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
247 Tensor ret = op_arange_start_out(Scalar(0), Scalar(5), Scalar(1), out);
248 EXPECT_TENSOR_CLOSE(out, expected_result);
249 }
250
TEST_F(OpArangeStartOutTest,StartOut)251 TEST_F(OpArangeStartOutTest, StartOut) {
252 TensorFactory<ScalarType::Float> tf;
253
254 Scalar start = Scalar(1.1);
255 Scalar end = Scalar(5.5);
256 Scalar step = Scalar(1.1);
257
258 Tensor out = tf.zeros({4});
259
260 Tensor ret = op_arange_start_out(start, end, step, out);
261
262 // Should always return the provided out Tensor.
263 EXPECT_TENSOR_EQ(ret, out);
264
265 // Expected tensor, equal
266 Tensor expected = tf.make({4}, {1.1, 2.2, 3.3, 4.4});
267
268 EXPECT_TENSOR_EQ(out, expected);
269
270 end = Scalar(5.51);
271 out = tf.zeros({5});
272
273 ret = op_arange_start_out(start, end, step, out);
274
275 // Should always return the provided out Tensor.
276 EXPECT_TENSOR_EQ(ret, out);
277
278 // Expected tensor, equal
279 expected = tf.make({5}, {1.1, 2.2, 3.3, 4.4, 5.5});
280
281 EXPECT_TENSOR_EQ(out, expected);
282 }
283
TEST_F(OpArangeStartOutTest,StartOutNegativeStep)284 TEST_F(OpArangeStartOutTest, StartOutNegativeStep) {
285 TensorFactory<ScalarType::Float> tf;
286
287 Scalar start = Scalar(5.5);
288 Scalar end = Scalar(1.1);
289 Scalar step = Scalar(-1.1);
290
291 Tensor out = tf.zeros({4});
292
293 Tensor ret = op_arange_start_out(start, end, step, out);
294
295 // Should always return the provided out Tensor.
296 EXPECT_TENSOR_EQ(ret, out);
297
298 // Expected tensor, equal
299 Tensor expected = tf.make({4}, {5.5, 4.4, 3.3, 2.2});
300
301 EXPECT_TENSOR_EQ(out, expected);
302
303 end = Scalar(1.09);
304 out = tf.zeros({5});
305
306 ret = op_arange_start_out(start, end, step, out);
307
308 // Should always return the provided out Tensor.
309 EXPECT_TENSOR_EQ(ret, out);
310
311 // Expected tensor, equal
312 expected = tf.make({5}, {5.5, 4.4, 3.3, 2.2, 1.1});
313
314 EXPECT_TENSOR_EQ(out, expected);
315 }
316