• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 <algorithm>
16 
17 #include <gtest/gtest.h>
18 
19 using namespace ::testing;
20 using executorch::aten::Scalar;
21 using executorch::aten::ScalarType;
22 using executorch::aten::Tensor;
23 using executorch::runtime::testing::TensorFactory;
24 using torch::executor::testing::SupportedFeatures;
25 namespace etrt = executorch::runtime;
26 
27 class OpMulOutTest : public OperatorTest {
28  protected:
op_mul_out(const Tensor & self,const Tensor & other,Tensor & out)29   Tensor& op_mul_out(const Tensor& self, const Tensor& other, Tensor& out) {
30     return torch::executor::aten::mul_outf(context_, self, other, out);
31   }
32 
33   // Common testing for multipling two integer Tensors
34   template <ScalarType DTYPE_A, ScalarType DTYPE_B, ScalarType DTYPE_OUT>
test_mul()35   void test_mul() {
36     TensorFactory<DTYPE_A> tf_a;
37     TensorFactory<DTYPE_B> tf_b;
38     TensorFactory<DTYPE_OUT> tf_out;
39 
40     const std::vector<int32_t> sizes = {2, 2};
41 
42     // Destination for the mul.
43     Tensor out = tf_out.zeros(sizes);
44 
45     // Multiply two tensors
46     op_mul_out(tf_a.make(sizes, /*data=*/{1, 2, 4, 8}), tf_b.ones(sizes), out);
47     EXPECT_TENSOR_EQ(out, tf_out.make(sizes, /*data=*/{1, 2, 4, 8}));
48 
49     op_mul_out(tf_a.make(sizes, /*data=*/{1, 2, 4, 8}), tf_b.zeros(sizes), out);
50     EXPECT_TENSOR_EQ(out, tf_out.make(sizes, /*data=*/{0, 0, 0, 0}));
51 
52     op_mul_out(
53         tf_a.make(sizes, /*data=*/{1, 2, 4, 8}),
54         tf_b.make(sizes, /*data=*/{1, 2, 4, 8}),
55         out);
56     EXPECT_TENSOR_EQ(out, tf_out.make(sizes, /*data=*/{1, 4, 16, 64}));
57   }
58 
59   template <ScalarType DTYPE_A, ScalarType DTYPE_B>
test_mul_enumerate_out_types()60   void test_mul_enumerate_out_types() {
61     test_mul<DTYPE_A, DTYPE_B, ScalarType::Half>();
62     test_mul<DTYPE_A, DTYPE_B, ScalarType::Float>();
63     test_mul<DTYPE_A, DTYPE_B, ScalarType::Double>();
64     // Integral out type is only allowed if both inputs are integral types
65     if (etrt::isIntegralType(DTYPE_A, false) &&
66         etrt::isIntegralType(DTYPE_B, false)) {
67       test_mul<DTYPE_A, DTYPE_B, ScalarType::Int>();
68       test_mul<DTYPE_A, DTYPE_B, ScalarType::Long>();
69     }
70   }
71 
72   template <ScalarType DTYPE_A>
test_mul_enumerate_b_types()73   void test_mul_enumerate_b_types() {
74 #define ENUMERATE_TEST_ENTRY(ctype, dtype) \
75   test_mul_enumerate_out_types<DTYPE_A, ScalarType::dtype>();
76 
77     ET_FORALL_REALHBF16_TYPES(ENUMERATE_TEST_ENTRY)
78 
79 #undef ENUMERATE_TEST_ENTRY
80   }
81 
82   // Common testing for multipling two floating point Tensors
83   template <ScalarType DTYPE>
test_floating_point_mul_out()84   void test_floating_point_mul_out() {
85     TensorFactory<DTYPE> tf;
86 
87     const std::vector<int32_t> sizes = {2, 2};
88 
89     // Destination for the mul.
90     Tensor out = tf.zeros(sizes);
91 
92     // Multiply two tensors
93     op_mul_out(
94         tf.make(sizes, /*data=*/{1.25, 2.5, 4.75, 8.875}), tf.ones(sizes), out);
95     EXPECT_TENSOR_CLOSE(out, tf.make(sizes, /*data=*/{1.25, 2.5, 4.75, 8.875}));
96 
97     op_mul_out(
98         tf.make(sizes, /*data=*/{1.1, 2.2, 4.4, 8.8}), tf.zeros(sizes), out);
99     EXPECT_TENSOR_CLOSE(out, tf.make(sizes, /*data=*/{0.0, 0.0, 0.0, 0.0}));
100 
101     op_mul_out(
102         tf.make(sizes, /*data=*/{1.25, 2.5, 4.75, 8.875}),
103         tf.make(sizes, /*data=*/{1.25, 2.5, 4.75, 8.875}),
104         out);
105     EXPECT_TENSOR_CLOSE(
106         out, tf.make(sizes, /*data=*/{1.5625, 6.25, 22.5625, 78.765625}));
107   }
108 
test_mul_enumerate_a_types()109   void test_mul_enumerate_a_types() {
110 #define ENUMERATE_TEST_ENTRY(ctype, dtype) \
111   test_mul_enumerate_b_types<ScalarType::dtype>();
112 
113     ET_FORALL_REALHBF16_TYPES(ENUMERATE_TEST_ENTRY)
114 
115 #undef ENUMERATE_TEST_ENTRY
116   }
117 
118   template <ScalarType DTYPE>
test_optimized_path_ignores_leading_1_dimensions()119   void test_optimized_path_ignores_leading_1_dimensions() {
120     TensorFactory<DTYPE> tf;
121 
122     const std::vector<int32_t> sizes1 = {1, 1, 2, 2};
123     const std::vector<int32_t> sizes2 = {1, 2, 2};
124 
125     // Destination for the mul.
126     Tensor out = tf.zeros(sizes1);
127 
128     // Multiply two tensors
129     op_mul_out(
130         tf.make(sizes1, /*data=*/{1.1, 2.2, 4.4, 8.8}), tf.ones(sizes2), out);
131     EXPECT_TENSOR_CLOSE(out, tf.make(sizes1, /*data=*/{1.1, 2.2, 4.4, 8.8}));
132   }
133 
134   template <ScalarType DTYPE>
test_broadcast_a2b()135   void test_broadcast_a2b() {
136     TensorFactory<DTYPE> tf_a;
137 
138     std::vector<std::vector<int32_t>> b_sizeses = {
139         {2},
140         {1, 2},
141     };
142     for (const auto& b_sizes : b_sizeses) {
143       // a and b of different shapes
144       Tensor a = tf_a.make({2, 2}, /*data=*/{1, 2, 3, 4});
145       Tensor b = tf_a.make(b_sizes, /*data=*/{2, 2});
146 
147       // Destination for output of mul.
148       Tensor out = tf_a.zeros({2, 2});
149 
150       // Check that it matches the expected output.
151       EXPECT_TENSOR_CLOSE(
152           op_mul_out(a, b, out), tf_a.make({2, 2}, /*data=*/{2, 4, 6, 8}));
153     }
154   }
155 
156   template <ScalarType DTYPE>
test_broadcast_3D()157   void test_broadcast_3D() {
158     TensorFactory<DTYPE> tf_a;
159 
160     Tensor a =
161         tf_a.make({2, 2, 3}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
162     Tensor b = tf_a.make({2, 1, 3}, /*data=*/{2, 3, 4, 5, 6, 7});
163 
164     // Destination for output of mul.
165     Tensor out =
166         tf_a.make({2, 2, 3}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
167     Tensor expected = tf_a.make(
168         {2, 2, 3}, /*data=*/{2, 6, 12, 8, 15, 24, 35, 48, 63, 50, 66, 84});
169 
170     // Check that it matches the expected output.
171     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
172     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
173   }
174 
175   template <ScalarType DTYPE>
test_broadcast_4D()176   void test_broadcast_4D() {
177     TensorFactory<DTYPE> tf_a;
178 
179     Tensor a = tf_a.make(
180         {2, 2, 3, 5},
181         /*data=*/{1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
182                   16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
183                   31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
184                   46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60});
185     Tensor b = tf_a.make(
186         {2, 1, 3, 5},
187         /*data=*/{1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
188                   16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30});
189 
190     // Destination for output of mul.
191     Tensor out = tf_a.zeros({2, 2, 3, 5});
192     Tensor expected = tf_a.make(
193         {2, 2, 3, 5},
194         /*data=*/{1,    4,    9,    16,   25,   36,   49,   64,   81,   100,
195                   121,  144,  169,  196,  225,  16,   34,   54,   76,   100,
196                   126,  154,  184,  216,  250,  286,  324,  364,  406,  450,
197                   496,  544,  594,  646,  700,  756,  814,  874,  936,  1000,
198                   1066, 1134, 1204, 1276, 1350, 736,  799,  864,  931,  1000,
199                   1071, 1144, 1219, 1296, 1375, 1456, 1539, 1624, 1711, 1800});
200 
201     // Check that it matches the expected output.
202     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
203     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
204 
205     b = tf_a.make(
206         {2, 2, 1, 5}, /*data=*/{1,  2,  3,  4,  5,  6,  7,  8,  9,  10,
207                                 11, 12, 13, 14, 15, 16, 17, 18, 19, 20});
208     out = tf_a.zeros({2, 2, 3, 5});
209     expected = tf_a.make(
210         {2, 2, 3, 5},
211         /*data=*/{1,   4,   9,   16,   25,   6,   14,  24,   36,   50,
212                   11,  24,  39,  56,   75,   96,  119, 144,  171,  200,
213                   126, 154, 184, 216,  250,  156, 189, 224,  261,  300,
214                   341, 384, 429, 476,  525,  396, 444, 494,  546,  600,
215                   451, 504, 559, 616,  675,  736, 799, 864,  931,  1000,
216                   816, 884, 954, 1026, 1100, 896, 969, 1044, 1121, 1200});
217 
218     // Check that it matches the expected output.
219     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
220     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
221   }
222 
223   template <ScalarType DTYPE>
test_broadcast_last_dim()224   void test_broadcast_last_dim() {
225     TensorFactory<DTYPE> tf_a;
226 
227     Tensor a =
228         tf_a.make({4, 3}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
229     Tensor b = tf_a.make({4, 1}, /*data=*/{2, 3, 4, 5});
230 
231     // Destination for output of mul.
232     Tensor out = tf_a.zeros({4, 3});
233     Tensor expected = tf_a.make(
234         {4, 3}, /*data=*/{2, 4, 6, 12, 15, 18, 28, 32, 36, 50, 55, 60});
235 
236     // Check that it matches the expected output.
237     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
238     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
239 
240     a = tf_a.make({2, 2, 3}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
241     b = tf_a.make({2, 2, 1}, /*data=*/{2, 3, 4, 5});
242 
243     // Destination for output of mul.
244     out = tf_a.zeros({2, 2, 3});
245     expected = tf_a.make(
246         {2, 2, 3}, /*data=*/{2, 4, 6, 12, 15, 18, 28, 32, 36, 50, 55, 60});
247 
248     // Check that it matches the expected output.
249     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
250     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
251 
252     a = tf_a.make(
253         {2, 2, 3, 5},
254         /*data=*/{1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
255                   16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
256                   31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
257                   46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60});
258     b = tf_a.make(
259         {2, 2, 3, 1},
260         /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
261 
262     // Destination for output of mul.
263     out = tf_a.zeros({2, 2, 3, 5});
264     expected = tf_a.make(
265         {2, 2, 3, 5},
266         /*data=*/{1,   2,   3,   4,   5,   12,  14,  16,  18,  20,  33,  36,
267                   39,  42,  45,  64,  68,  72,  76,  80,  105, 110, 115, 120,
268                   125, 156, 162, 168, 174, 180, 217, 224, 231, 238, 245, 288,
269                   296, 304, 312, 320, 369, 378, 387, 396, 405, 460, 470, 480,
270                   490, 500, 561, 572, 583, 594, 605, 672, 684, 696, 708, 720});
271 
272     // Check that it matches the expected output.
273     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
274     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
275   }
276 
277   template <ScalarType DTYPE>
test_broadcast_b2a()278   void test_broadcast_b2a() {
279     TensorFactory<DTYPE> tf_a;
280     // a and b of different shapes
281     Tensor a = tf_a.make({2}, /*data=*/{2, 2});
282     Tensor b = tf_a.make({2, 2}, /*data=*/{1, 2, 3, 4});
283 
284     // Destination for output of mul.
285     Tensor out = tf_a.zeros({2, 2});
286 
287     // Check that it matches the expected output.
288     EXPECT_TENSOR_CLOSE(
289         op_mul_out(a, b, out), tf_a.make({2, 2}, /*data=*/{2, 4, 6, 8}));
290   }
291 
292   template <ScalarType DTYPE>
test_scalar_input_broadcast()293   void test_scalar_input_broadcast() {
294     TensorFactory<DTYPE> tf_a;
295 
296     // a is a 1d tensor and b is a scalar
297     Tensor a = tf_a.make({2}, /*data=*/{2, 2});
298     Tensor b = tf_a.make({}, /*data=*/{2});
299 
300     // Destination for output of mul.
301     Tensor out = tf_a.make({2}, /*data=*/{2, 2});
302     Tensor expected = tf_a.make({2}, /*data=*/{4, 4});
303 
304     // Check that it matches the expected output.
305     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
306     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
307   }
308 
309   template <ScalarType DTYPE>
test_both_scalar_input_broadcast()310   void test_both_scalar_input_broadcast() {
311     TensorFactory<DTYPE> tf_a;
312 
313     // a is a rank-1 scalar and b is a rank-0 scalar
314     Tensor a = tf_a.make({1}, /*data=*/{2});
315     Tensor b = tf_a.make({}, /*data=*/{2});
316 
317     // Destination for output of mul.
318     Tensor out = tf_a.make({1}, /*data=*/{2});
319     Tensor expected = tf_a.make({1}, /*data=*/{4});
320 
321     // Check that it matches the expected output.
322     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
323     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
324   }
325 };
326 
327 class OpMulScalarOutTest : public OperatorTest {
328  protected:
329   Tensor&
op_mul_scalar_out(const Tensor & self,const Scalar & other,Tensor & out)330   op_mul_scalar_out(const Tensor& self, const Scalar& other, Tensor& out) {
331     return torch::executor::aten::mul_outf(context_, self, other, out);
332   }
333 };
334 
335 //
336 // Correctness Tests
337 //
338 
339 /**
340  * Uses the function templates above to test all valid combinations of
341  * inputs*and output dtypes*/
TEST_F(OpMulOutTest,AllRealDtypesSupported)342 TEST_F(OpMulOutTest, AllRealDtypesSupported) {
343   test_mul_enumerate_a_types();
344 }
345 
TEST_F(OpMulOutTest,FloatTensors)346 TEST_F(OpMulOutTest, FloatTensors) {
347   test_floating_point_mul_out<ScalarType::Float>();
348 }
349 
TEST_F(OpMulOutTest,DoubleTensors)350 TEST_F(OpMulOutTest, DoubleTensors) {
351   test_floating_point_mul_out<ScalarType::Double>();
352 }
353 
TEST_F(OpMulOutTest,HalfTensors)354 TEST_F(OpMulOutTest, HalfTensors) {
355   test_floating_point_mul_out<ScalarType::Half>();
356 }
357 
TEST_F(OpMulOutTest,BFloat16Tensors)358 TEST_F(OpMulOutTest, BFloat16Tensors) {
359   test_floating_point_mul_out<ScalarType::BFloat16>();
360 }
361 
TEST_F(OpMulOutTest,BoolTensors)362 TEST_F(OpMulOutTest, BoolTensors) {
363   TensorFactory<ScalarType::Bool> tf;
364 
365   const std::vector<int32_t> sizes = {2, 2};
366 
367   // Destination for the mul.
368   Tensor out = tf.zeros(sizes);
369 
370   // Multiply two tensors
371   op_mul_out(
372       tf.make(sizes, /*data=*/{true, false, true, true}), tf.ones(sizes), out);
373   EXPECT_TENSOR_EQ(out, tf.make(sizes, /*data=*/{true, false, true, true}));
374 
375   op_mul_out(
376       tf.make(sizes, /*data=*/{true, false, true, true}), tf.zeros(sizes), out);
377   EXPECT_TENSOR_EQ(out, tf.make(sizes, /*data=*/{false, false, false, false}));
378 
379   op_mul_out(
380       tf.make(sizes, /*data=*/{true, false, true, true}),
381       tf.make(sizes, /*data=*/{false, false, true, false}),
382       out);
383   EXPECT_TENSOR_EQ(out, tf.make(sizes, /*data=*/{false, false, true, false}));
384 }
385 
TEST_F(OpMulOutTest,OptimizedPathIgnoresLeading1Dimensions)386 TEST_F(OpMulOutTest, OptimizedPathIgnoresLeading1Dimensions) {
387 #define ENUMERATE_TEST_ENTRY(ctype, dtype) \
388   test_optimized_path_ignores_leading_1_dimensions<ScalarType::dtype>();
389 
390   ET_FORALL_FLOATHBF16_TYPES(ENUMERATE_TEST_ENTRY);
391 
392 #undef ENUMERATE_TEST_ENTRY
393 }
394 
395 // Mismatched shape tests.
TEST_F(OpMulOutTest,MismatchedNonBroadcastableInputShapesDies)396 TEST_F(OpMulOutTest, MismatchedNonBroadcastableInputShapesDies) {
397   if (SupportedFeatures::get()->is_aten) {
398     GTEST_SKIP() << "ATen currently supports mismatched shapes";
399   }
400 
401   TensorFactory<ScalarType::Int> tf;
402 
403   // Input tensors with different shapes.
404   Tensor a = tf.ones(/*sizes=*/{4, 2});
405   Tensor b = tf.ones(/*sizes=*/{2, 2});
406 
407   // Output tensor; matches the shape of one of the inputs.
408   Tensor out = tf.zeros(/*sizes=*/{8});
409 
410   // Multiplying the two mismatched tensors should cause an assertion and kill
411   // the test process.
412   ET_EXPECT_KERNEL_FAILURE(context_, op_mul_out(a, b, out));
413 }
414 
415 // Broadcast tensor b's size to tensor a's size
TEST_F(OpMulOutTest,BroadcastA2BTest)416 TEST_F(OpMulOutTest, BroadcastA2BTest) {
417   test_broadcast_a2b<ScalarType::Int>();
418   test_broadcast_a2b<ScalarType::Half>();
419   test_broadcast_a2b<ScalarType::BFloat16>();
420 
421   // Test 3D tensors
422   test_broadcast_3D<ScalarType::Float>();
423   test_broadcast_3D<ScalarType::Half>();
424   test_broadcast_3D<ScalarType::BFloat16>();
425 
426   // Test 4D tensors
427   test_broadcast_4D<ScalarType::Float>();
428   test_broadcast_4D<ScalarType::Half>();
429   test_broadcast_4D<ScalarType::BFloat16>();
430 }
431 
432 // Broadcast tensor a's size to tensor b's size
TEST_F(OpMulOutTest,BroadcastB2ATest)433 TEST_F(OpMulOutTest, BroadcastB2ATest) {
434   test_broadcast_b2a<ScalarType::Int>();
435   test_broadcast_b2a<ScalarType::Half>();
436   test_broadcast_b2a<ScalarType::BFloat16>();
437 }
438 
TEST_F(OpMulOutTest,BroadcastNDTest)439 TEST_F(OpMulOutTest, BroadcastNDTest) {
440   // Test 3D tensors
441   test_broadcast_3D<ScalarType::Float>();
442   test_broadcast_3D<ScalarType::Half>();
443   test_broadcast_3D<ScalarType::BFloat16>();
444 
445   // Test 4D tensors
446   test_broadcast_4D<ScalarType::Float>();
447   test_broadcast_4D<ScalarType::Half>();
448   test_broadcast_4D<ScalarType::BFloat16>();
449 
450   // Test broadcasting on the last dimension
451   test_broadcast_last_dim<ScalarType::Float>();
452   test_broadcast_last_dim<ScalarType::Half>();
453   test_broadcast_last_dim<ScalarType::BFloat16>();
454 }
455 
TEST_F(OpMulOutTest,BroadcastLastDimTest)456 TEST_F(OpMulOutTest, BroadcastLastDimTest) {
457   // Test broadcasting on the last dimension
458   test_broadcast_last_dim<ScalarType::Float>();
459   test_broadcast_last_dim<ScalarType::Half>();
460   test_broadcast_last_dim<ScalarType::BFloat16>();
461 }
462 
463 // Broadcast tensor a and b's size to a new size c.
TEST_F(OpMulOutTest,BroadcastAB2CTest)464 TEST_F(OpMulOutTest, BroadcastAB2CTest) {
465   TensorFactory<ScalarType::Int> tf_a;
466 
467   // a and b of different shapes
468   Tensor a = tf_a.make({2, 1}, /*data=*/{1, 2});
469   Tensor b = tf_a.make({2, 1, 2}, /*data=*/{1, 2, 3, 4});
470 
471   // Destination for output of mul.
472   Tensor out = tf_a.zeros({2, 2, 2});
473 
474   // Check that it matches the expected output.
475   EXPECT_TENSOR_CLOSE(
476       op_mul_out(a, b, out),
477       tf_a.make({2, 2, 2}, /*data=*/{1, 2, 2, 4, 3, 4, 6, 8}));
478 }
479 
TEST_F(OpMulOutTest,ScalarInputBroadcastTest)480 TEST_F(OpMulOutTest, ScalarInputBroadcastTest) {
481   test_scalar_input_broadcast<ScalarType::Int>();
482   test_scalar_input_broadcast<ScalarType::Half>();
483   test_scalar_input_broadcast<ScalarType::BFloat16>();
484 }
485 
TEST_F(OpMulOutTest,BothScalarInputBroadcastTest)486 TEST_F(OpMulOutTest, BothScalarInputBroadcastTest) {
487   test_both_scalar_input_broadcast<ScalarType::Int>();
488   test_both_scalar_input_broadcast<ScalarType::Half>();
489   test_both_scalar_input_broadcast<ScalarType::BFloat16>();
490 }
491 
TEST_F(OpMulOutTest,MismatchedOutputShapesDies)492 TEST_F(OpMulOutTest, MismatchedOutputShapesDies) {
493   if (SupportedFeatures::get()->is_aten) {
494     GTEST_SKIP() << "ATen currently supports mismatched shapes";
495   }
496 
497   TensorFactory<ScalarType::Int> tf;
498 
499   const std::vector<int32_t> sizes = {2, 2};
500 
501   // Input tensors with the same shapes.
502   Tensor a = tf.ones(sizes);
503   Tensor b = tf.ones(sizes);
504 
505   // Output tensor with a different shape.
506   Tensor out = tf.zeros(/*sizes=*/{4});
507 
508   // Multiplying the tensors into a mismatched output should cause an assertion
509   // and kill the test process.
510   ET_EXPECT_KERNEL_FAILURE(context_, op_mul_out(a, b, out));
511 }
512 
TEST_F(OpMulOutTest,BroadcastDimSizeIsOneAB)513 TEST_F(OpMulOutTest, BroadcastDimSizeIsOneAB) {
514   TensorFactory<ScalarType::Float> tf;
515 
516   Tensor x = tf.make(
517       {3, 2},
518       {0.3200607895851135,
519        0.029979348182678223,
520        0.27112698554992676,
521        0.15423381328582764,
522        0.6920414566993713,
523        0.005174398422241211});
524   Tensor y = tf.make({1, 2}, {0.9711773991584778, 0.8632034063339233});
525   Tensor expected_result = tf.make(
526       {3, 2},
527       {0.3108358085155487,
528        0.02587827481329441,
529        0.2633123993873596,
530        0.13313515484333038,
531        0.672095000743866,
532        0.004466558340936899});
533 
534   Tensor out = tf.zeros({3, 2});
535   Tensor ret = op_mul_out(x, y, out);
536   EXPECT_TENSOR_CLOSE(out, expected_result);
537 }
538 
TEST_F(OpMulOutTest,BroadcastDimSizeMissingAB)539 TEST_F(OpMulOutTest, BroadcastDimSizeMissingAB) {
540   TensorFactory<ScalarType::Float> tf;
541 
542   Tensor x = tf.make(
543       {3, 2},
544       {0.3200607895851135,
545        0.029979348182678223,
546        0.27112698554992676,
547        0.15423381328582764,
548        0.6920414566993713,
549        0.005174398422241211});
550   Tensor y = tf.make({2}, {0.9711773991584778, 0.8632034063339233});
551   Tensor expected_result = tf.make(
552       {3, 2},
553       {0.3108358085155487,
554        0.02587827481329441,
555        0.2633123993873596,
556        0.13313515484333038,
557        0.672095000743866,
558        0.004466558340936899});
559 
560   Tensor out = tf.zeros({3, 2});
561   Tensor ret = op_mul_out(x, y, out);
562   EXPECT_TENSOR_CLOSE(out, expected_result);
563 }
564 
TEST_F(OpMulOutTest,BroadcastDimSizeIsOneBA)565 TEST_F(OpMulOutTest, BroadcastDimSizeIsOneBA) {
566   TensorFactory<ScalarType::Float> tf;
567 
568   Tensor x = tf.make({1, 2}, {0.9711773991584778, 0.8632034063339233});
569   Tensor y = tf.make(
570       {3, 2},
571       {0.3200607895851135,
572        0.029979348182678223,
573        0.27112698554992676,
574        0.15423381328582764,
575        0.6920414566993713,
576        0.005174398422241211});
577   Tensor expected_result = tf.make(
578       {3, 2},
579       {0.3108358085155487,
580        0.02587827481329441,
581        0.2633123993873596,
582        0.13313515484333038,
583        0.672095000743866,
584        0.004466558340936899});
585 
586   Tensor out = tf.zeros({3, 2});
587   Tensor ret = op_mul_out(x, y, out);
588   EXPECT_TENSOR_CLOSE(out, expected_result);
589 }
590 
TEST_F(OpMulOutTest,BroadcastDimSizeMissingBA)591 TEST_F(OpMulOutTest, BroadcastDimSizeMissingBA) {
592   TensorFactory<ScalarType::Float> tf;
593 
594   Tensor x = tf.make({1, 2}, {0.9711773991584778, 0.8632034063339233});
595   Tensor y = tf.make(
596       {3, 2},
597       {0.3200607895851135,
598        0.029979348182678223,
599        0.27112698554992676,
600        0.15423381328582764,
601        0.6920414566993713,
602        0.005174398422241211});
603   Tensor expected_result = tf.make(
604       {3, 2},
605       {0.3108358085155487,
606        0.02587827481329441,
607        0.2633123993873596,
608        0.13313515484333038,
609        0.672095000743866,
610        0.004466558340936899});
611 
612   Tensor out = tf.zeros({3, 2});
613   Tensor ret = op_mul_out(x, y, out);
614   EXPECT_TENSOR_CLOSE(out, expected_result);
615 }
616 
TEST_F(OpMulOutTest,DynamicShapeUpperBoundSameAsExpected)617 TEST_F(OpMulOutTest, DynamicShapeUpperBoundSameAsExpected) {
618   TensorFactory<ScalarType::Float> tf;
619 
620   Tensor x = tf.make(
621       {3, 2},
622       {0.6910695433616638,
623        0.6540696620941162,
624        0.8072559237480164,
625        0.8218746185302734,
626        0.9193597435951233,
627        0.4525110721588135});
628   Tensor y = tf.make(
629       {3, 2},
630       {0.9212601184844971,
631        0.2030404806137085,
632        0.34644562005996704,
633        0.4489826560020447,
634        0.5666958689689636,
635        0.5006863474845886});
636   Tensor expected_result = tf.make(
637       {3, 2},
638       {0.636654794216156,
639        0.13280262053012848,
640        0.27967026829719543,
641        0.3690074384212494,
642        0.5209973454475403,
643        0.2265661209821701});
644 
645   Tensor out =
646       tf.zeros({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
647   Tensor ret = op_mul_out(x, y, out);
648   EXPECT_TENSOR_CLOSE(out, expected_result);
649 }
650 
TEST_F(OpMulOutTest,DynamicShapeUpperBoundLargerThanExpected)651 TEST_F(OpMulOutTest, DynamicShapeUpperBoundLargerThanExpected) {
652   TensorFactory<ScalarType::Float> tf;
653 
654   Tensor x = tf.make(
655       {3, 2},
656       {0.6910695433616638,
657        0.6540696620941162,
658        0.8072559237480164,
659        0.8218746185302734,
660        0.9193597435951233,
661        0.4525110721588135});
662   Tensor y = tf.make(
663       {3, 2},
664       {0.9212601184844971,
665        0.2030404806137085,
666        0.34644562005996704,
667        0.4489826560020447,
668        0.5666958689689636,
669        0.5006863474845886});
670   Tensor expected_result = tf.make(
671       {3, 2},
672       {0.636654794216156,
673        0.13280262053012848,
674        0.27967026829719543,
675        0.3690074384212494,
676        0.5209973454475403,
677        0.2265661209821701});
678 
679   Tensor out =
680       tf.zeros({10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
681   Tensor ret = op_mul_out(x, y, out);
682   EXPECT_TENSOR_CLOSE(out, expected_result);
683 }
684 
TEST_F(OpMulOutTest,DynamicShapeUnbound)685 TEST_F(OpMulOutTest, DynamicShapeUnbound) {
686   GTEST_SKIP() << "Dynamic shape not supported";
687   TensorFactory<ScalarType::Float> tf;
688 
689   Tensor x = tf.make(
690       {3, 2},
691       {0.6910695433616638,
692        0.6540696620941162,
693        0.8072559237480164,
694        0.8218746185302734,
695        0.9193597435951233,
696        0.4525110721588135});
697   Tensor y = tf.make(
698       {3, 2},
699       {0.9212601184844971,
700        0.2030404806137085,
701        0.34644562005996704,
702        0.4489826560020447,
703        0.5666958689689636,
704        0.5006863474845886});
705   Tensor expected_result = tf.make(
706       {3, 2},
707       {0.636654794216156,
708        0.13280262053012848,
709        0.27967026829719543,
710        0.3690074384212494,
711        0.5209973454475403,
712        0.2265661209821701});
713 
714   Tensor out =
715       tf.zeros({1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
716   Tensor ret = op_mul_out(x, y, out);
717   EXPECT_TENSOR_CLOSE(out, expected_result);
718 }
719 
TEST_F(OpMulScalarOutTest,SanityCheck)720 TEST_F(OpMulScalarOutTest, SanityCheck) {
721   TensorFactory<ScalarType::Bool> tf_a;
722   TensorFactory<ScalarType::Float> tf_out;
723 
724   const std::vector<int32_t> sizes = {2, 2};
725 
726   Tensor out = tf_out.zeros(sizes);
727 
728   op_mul_scalar_out(tf_a.make(sizes, {true, false, true, false}), 2.3, out);
729 
730   // Check that it matches the expected output.
731   EXPECT_TENSOR_EQ(out, tf_out.make(sizes, {2.3, 0.0, 2.3, 0.0}));
732 }
733 
TEST_F(OpMulScalarOutTest,OptimizedSanityCheck)734 TEST_F(OpMulScalarOutTest, OptimizedSanityCheck) {
735   TensorFactory<ScalarType::Float> tf;
736 
737   const std::vector<int32_t> sizes = {2, 2};
738 
739   Tensor out = tf.zeros(sizes);
740 
741   op_mul_scalar_out(tf.make(sizes, {1.3, 2.1, 4.6, 8.2}), 2.0, out);
742 
743   // Check that it matches the expected output.
744   EXPECT_TENSOR_CLOSE(out, tf.make(sizes, {2.6, 4.2, 9.2, 16.4}));
745 }
746 
TEST_F(OpMulScalarOutTest,HalfSanityCheck)747 TEST_F(OpMulScalarOutTest, HalfSanityCheck) {
748   TensorFactory<ScalarType::Half> tf;
749 
750   const std::vector<int32_t> sizes = {2, 2};
751 
752   Tensor out = tf.zeros(sizes);
753 
754   op_mul_scalar_out(tf.make(sizes, {1.3, 2.1, 4.6, 8.2}), 2.0, out);
755 
756   // Check that it matches the expected output.
757   EXPECT_TENSOR_CLOSE(out, tf.make(sizes, {2.6, 4.2, 9.2, 16.4}));
758 }
759 
TEST_F(OpMulScalarOutTest,BFloat16SanityCheck)760 TEST_F(OpMulScalarOutTest, BFloat16SanityCheck) {
761   TensorFactory<ScalarType::BFloat16> tf;
762 
763   const std::vector<int32_t> sizes = {2, 2};
764 
765   Tensor out = tf.zeros(sizes);
766 
767   op_mul_scalar_out(tf.make(sizes, {1.3, 2.1, 4.6, 8.2}), 2.0, out);
768 
769   // Check that it matches the expected output.
770   EXPECT_TENSOR_CLOSE(out, tf.make(sizes, {2.6, 4.2, 9.2, 16.4}));
771 }
772