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::TensorFactory;
23
24 class OpFillTest : public OperatorTest {
25 protected:
26 Tensor&
op_fill_scalar_out(const Tensor & self,const Scalar & other,Tensor & out)27 op_fill_scalar_out(const Tensor& self, const Scalar& other, Tensor& out) {
28 return torch::executor::aten::fill_outf(context_, self, other, out);
29 }
30
31 Tensor&
op_fill_tensor_out(const Tensor & self,const Tensor & other,Tensor & out)32 op_fill_tensor_out(const Tensor& self, const Tensor& other, Tensor& out) {
33 return torch::executor::aten::fill_outf(context_, self, other, out);
34 }
35
36 template <ScalarType DTYPE>
test_fill_scalar_out(std::vector<int32_t> && sizes)37 void test_fill_scalar_out(std::vector<int32_t>&& sizes) {
38 TensorFactory<DTYPE> tf;
39
40 // Before: `out` consists of 0s.
41 Tensor self = tf.zeros(sizes);
42 Tensor out = tf.zeros(sizes);
43
44 // After: `out` consists of 1s.
45 Scalar other = 1;
46 if (DTYPE == ScalarType::Bool) {
47 other = false;
48 }
49 op_fill_scalar_out(self, other, out);
50
51 Tensor exp_out = tf.full(sizes, 1);
52 if (DTYPE == ScalarType::Bool) {
53 exp_out = tf.full(sizes, false);
54 }
55
56 // Check `out` matches expected output.
57 EXPECT_TENSOR_EQ(out, exp_out);
58 }
59
60 template <ScalarType DTYPE>
test_fill_tensor_out(std::vector<int32_t> && sizes)61 void test_fill_tensor_out(std::vector<int32_t>&& sizes) {
62 TensorFactory<DTYPE> tf;
63
64 // Before: `out` consists of 0s.
65 Tensor self = tf.zeros(sizes);
66 Tensor out = tf.zeros(sizes);
67
68 // After: `out` consists of 1s.
69 Tensor other = tf.ones({});
70 op_fill_tensor_out(self, other, out);
71
72 Tensor exp_out = tf.full(sizes, 1);
73
74 // Check `out` matches expected output.
75 EXPECT_TENSOR_EQ(out, exp_out);
76 }
77 };
78
79 // A macro for defining tests for both scalar and tensor variants of
80 // `fill_out`. Here the `self` and `out` tensors will be created according
81 // to the sizes provided, while the scalar/tensor will be a singleton.
82 #define TEST_FILL_OUT(FN, DTYPE) \
83 FN<ScalarType::DTYPE>({}); \
84 FN<ScalarType::DTYPE>({1}); \
85 FN<ScalarType::DTYPE>({1, 1, 1}); \
86 FN<ScalarType::DTYPE>({2, 0, 4}); \
87 FN<ScalarType::DTYPE>({2, 3, 4});
88
89 // Create input support tests for scalar variant.
90 #define GENERATE_SCALAR_INPUT_SUPPORT_TEST(_, DTYPE) \
91 TEST_F(OpFillTest, DTYPE##ScalarInputSupport) { \
92 TEST_FILL_OUT(test_fill_scalar_out, DTYPE); \
93 }
94
ET_FORALL_REAL_TYPES_AND(Bool,GENERATE_SCALAR_INPUT_SUPPORT_TEST)95 ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_SCALAR_INPUT_SUPPORT_TEST)
96
97 // Create input support tests for tensor variant.
98 #define GENERATE_TENSOR_INPUT_SUPPORT_TEST(_, DTYPE) \
99 TEST_F(OpFillTest, DTYPE##TensorInputSupport) { \
100 TEST_FILL_OUT(test_fill_tensor_out, DTYPE); \
101 }
102
103 ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_TENSOR_INPUT_SUPPORT_TEST)
104
105 TEST_F(OpFillTest, MismatchedOtherPropertiesDies) {
106 TensorFactory<ScalarType::Int> tf;
107
108 // `self` and `out` have different shapes but same dtype.
109 Tensor self = tf.zeros({1});
110 Tensor out = tf.zeros({1});
111
112 // Create `other` tensors with incompatible shapes (`dim()` >=1) and/or
113 // elements (`numel()` > 1).
114
115 Tensor other1 = tf.zeros({1});
116 EXPECT_EQ(other1.dim(), 1);
117 EXPECT_EQ(other1.numel(), 1);
118
119 Tensor other2 = tf.zeros({2});
120 EXPECT_EQ(other2.dim(), 1);
121 EXPECT_EQ(other2.numel(), 2);
122
123 Tensor other3 = tf.zeros({3, 3});
124 EXPECT_EQ(other3.dim(), 2);
125 EXPECT_EQ(other3.numel(), 9);
126
127 // Assert `other` tensors with incompatible properties fails.
128 ET_EXPECT_KERNEL_FAILURE(context_, op_fill_tensor_out(self, other1, out));
129 ET_EXPECT_KERNEL_FAILURE(context_, op_fill_tensor_out(self, other2, out));
130 ET_EXPECT_KERNEL_FAILURE(context_, op_fill_tensor_out(self, other3, out));
131 }
132
TEST_F(OpFillTest,MismatchedOutputShapesDies)133 TEST_F(OpFillTest, MismatchedOutputShapesDies) {
134 // Skip ATen test since it supports `self` and `out` having different shapes.
135 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
136 GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
137 }
138
139 TensorFactory<ScalarType::Int> tf;
140
141 // `self` and `out` have different shapes but same dtype.
142 Tensor self = tf.zeros({1});
143 Tensor out = tf.zeros({2, 2});
144
145 // Assert `out` can't be filled due to incompatible shapes.
146 ET_EXPECT_KERNEL_FAILURE(context_, op_fill_scalar_out(self, 0, out));
147 }
148
TEST_F(OpFillTest,MismatchedOutputDtypeDies)149 TEST_F(OpFillTest, MismatchedOutputDtypeDies) {
150 TensorFactory<ScalarType::Byte> tf_byte;
151 TensorFactory<ScalarType::Float> tf_float;
152
153 // `self` and `out` have different dtypes but same shape.
154 Tensor self = tf_byte.zeros({2, 2});
155 Tensor out = tf_float.ones({2, 2});
156
157 // Assert `out` can't be filled due to incompatible dtype.
158 ET_EXPECT_KERNEL_FAILURE(context_, op_fill_scalar_out(self, 0.0, out));
159 }
160