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 <cmath>
10 #include <ostream>
11
12 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
13 #include <executorch/kernels/test/TestUtil.h>
14 #include <executorch/kernels/test/supported_features.h>
15 #include <executorch/runtime/core/exec_aten/exec_aten.h>
16 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
17 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
18 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
19 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
20
21 #include <gtest/gtest.h>
22
23 using namespace ::testing;
24 using exec_aten::ArrayRef;
25 using exec_aten::IntArrayRef;
26 using exec_aten::nullopt;
27 using exec_aten::optional;
28 using exec_aten::Scalar;
29 using exec_aten::ScalarType;
30 using exec_aten::Tensor;
31 using torch::executor::testing::TensorFactory;
32
33 using OptScalar = exec_aten::optional<Scalar>;
34
35 class OpNativeLayerNormTest : public OperatorTest {
36 protected:
op_native_layer_norm_out(const Tensor & input,IntArrayRef normalized_shape,const optional<Tensor> & weight,const optional<Tensor> & bias,double eps,Tensor & out0,Tensor & out1,Tensor & out2)37 ::std::tuple<Tensor&, Tensor&, Tensor&> op_native_layer_norm_out(
38 const Tensor& input,
39 IntArrayRef normalized_shape,
40 const optional<Tensor>& weight,
41 const optional<Tensor>& bias,
42 double eps,
43 Tensor& out0,
44 Tensor& out1,
45 Tensor& out2) {
46 return torch::executor::aten::native_layer_norm_outf(
47 context_, input, normalized_shape, weight, bias, eps, out0, out1, out2);
48 }
49
50 template <ScalarType DTYPE>
51 struct NativeLayerNormTestCase {
52 using ctype = typename TensorFactory<DTYPE>::ctype;
53
54 // Human-readable, unique title for the test case. Printed if the test
55 // fails.
56 const std::string title;
57 // Size vector for the input/output
58 const std::vector<int32_t> sizes;
59 // Data for the input tensor; must agree with `sizes`.
60 const std::vector<ctype> input_data;
61 // The normalized shape. Only the last dim is accepted.
62 const std::vector<int32_t> normalized_shape;
63 // Affine transform weight.
64 const std::vector<ctype> weight_data;
65 // Affine transform bias.
66 const std::vector<ctype> bias_data;
67 // a value added to the denominator for numerical stability
68 const ctype eps;
69 // The expected output data.
70 const std::vector<ctype> expected_data;
71 };
72
73 /// Runs the provided test cases.
74 template <ScalarType DTYPE>
run_test_cases(std::vector<NativeLayerNormTestCase<DTYPE>> test_cases)75 void run_test_cases(std::vector<NativeLayerNormTestCase<DTYPE>> test_cases) {
76 TensorFactory<DTYPE> tf;
77 for (const auto& test_case : test_cases) {
78 SCOPED_TRACE(test_case.title); // Printed if the test fails
79
80 Tensor in = tf.make(test_case.sizes, test_case.input_data);
81 Tensor weight =
82 tf.make(test_case.normalized_shape, test_case.weight_data);
83 Tensor bias = tf.make(test_case.normalized_shape, test_case.bias_data);
84 Tensor out0 = tf.zeros(test_case.sizes);
85 Tensor out1 = tf.zeros(
86 test_case.sizes, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
87 Tensor out2 = tf.zeros(
88 test_case.sizes, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
89 auto normalized_shape_vec = std::vector<int64_t>(
90 test_case.normalized_shape.begin(), test_case.normalized_shape.end());
91 auto normalized_shape = exec_aten::ArrayRef<int64_t>(
92 normalized_shape_vec.data(), normalized_shape_vec.size());
93 auto result = op_native_layer_norm_out(
94 in, normalized_shape, weight, bias, test_case.eps, out0, out1, out2);
95 EXPECT_TENSOR_CLOSE(out0, std::get<0>(result));
96
97 Tensor expected = tf.make(test_case.sizes, test_case.expected_data);
98 EXPECT_TENSOR_CLOSE(out0, expected);
99 }
100 }
101
102 // Test cases that are compatible with float and double.
103 template <ScalarType DTYPE>
run_floating_point_test_cases()104 void run_floating_point_test_cases() {
105 constexpr auto kInfinity =
106 std::numeric_limits<typename TensorFactory<DTYPE>::ctype>::infinity();
107 // Reference colab note:
108 // https://colab.research.google.com/drive/1KZT6sEY-h7lwZlwBanbLl77M5OuzzsZI#scrollTo=18WtUPCXYCPx
109 std::vector<NativeLayerNormTestCase<DTYPE>> test_cases = {
110 {
111 std::string(__func__) + ": Simple negative/positive layer norm",
112 {2, 3}, // sizes
113 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
114 {3}, // normalized shape
115 {1.0, 1.0, 1.0}, // weights
116 {0.0, 0.0, 0.0}, // bias
117 1.0e-5, // eps
118 {1.22474,
119 0.0000,
120 -1.22474,
121 -0.925819,
122 1.38873,
123 -0.46291}, // expected_data
124 },
125 {
126 std::string(__func__) + ": non-default eps",
127 {2, 3}, // sizes
128 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
129 {3}, // normalized shape
130 {1.0, 1.0, 1.0}, // weights
131 {0.0, 0.0, 0.0}, // bias
132 1.0e-3, // eps
133 {1.22383,
134 0,
135 -1.22383,
136 -0.925721,
137 1.38858,
138 -0.46286}, // expected_data
139 },
140 {
141 std::string(__func__) + ": non-default weights",
142 {2, 3}, // sizes
143 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
144 {3}, // normalized shape
145 {2.0, 2.0, 2.0}, // weights
146 {0.0, 0.0, 0.0}, // bias
147 1.0e-5, // eps
148 {2.44947,
149 0,
150 -2.44947,
151 -1.85164,
152 2.77746,
153 -0.925819}, // expected_data
154 },
155 {
156 std::string(__func__) + ": non-default bias",
157 {2, 3}, // sizes
158 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
159 {3}, // normalized shape
160 {1.0, 1.0, 1.0}, // weights
161 {1.0, 1.0, 1.0}, // bias
162 1.0e-5, // eps
163 {2.22474,
164 1,
165 -0.224736,
166 0.0741809,
167 2.38873,
168 0.53709}, // expected_data
169 },
170 {
171 std::string(__func__) + ": infinite input brings NAN results",
172 {2, 3}, // sizes
173 {kInfinity, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
174 {3}, // normalized shape
175 {1.0, 1.0, 1.0}, // weights
176 {1.0, 1.0, 1.0}, // bias
177 1.0e-5, // eps
178 {-NAN, -NAN, -NAN, 0.0741809, 2.38873, 0.53709}, // expected_data
179 },
180 {
181 std::string(__func__) + ": NAN input brings NAN results",
182 {2, 3}, // sizes
183 {NAN, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
184 {3}, // normalized shape
185 {1.0, 1.0, 1.0}, // weights
186 {1.0, 1.0, 1.0}, // bias
187 1.0e-5, // eps
188 {-NAN, -NAN, -NAN, 0.0741809, 2.38873, 0.53709}, // expected_data
189 },
190 {
191 std::string(__func__) + ": NAN weight brings NAN results",
192 {2, 3}, // sizes
193 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
194 {3}, // normalized shape
195 {NAN, 1.0, 1.0}, // weights
196 {1.0, 1.0, 1.0}, // bias
197 1.0e-5, // eps
198 {NAN, 1, -0.224736, NAN, 2.38873, 0.53709}, // expected_data
199 },
200 {
201 std::string(__func__) + ": inf weight brings inf results",
202 {2, 3}, // sizes
203 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
204 {3}, // normalized shape
205 {kInfinity, 1.0, 1.0}, // weights
206 {0.0, 0.0, 0.0}, // bias
207 1.0e-5, // eps
208 {kInfinity,
209 0,
210 -1.22474,
211 -kInfinity,
212 1.38873,
213 -0.46291}, // expected_data
214 },
215 {
216 std::string(__func__) + ": inf bias brings inf results",
217 {2, 3}, // sizes
218 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
219 {3}, // normalized shape
220 {kInfinity, 1.0, 1.0}, // weights
221 {0.0, 0.0, 0.0}, // bias
222 1.0e-5, // eps
223 {kInfinity,
224 0,
225 -1.22474,
226 -kInfinity,
227 1.38873,
228 -0.46291}, // expected_data
229 },
230 };
231
232 run_test_cases(test_cases);
233 }
234
235 // Runs death test cases.
236 template <ScalarType DTYPE>
run_death_test_cases(std::vector<NativeLayerNormTestCase<DTYPE>> test_cases)237 void run_death_test_cases(
238 std::vector<NativeLayerNormTestCase<DTYPE>> test_cases) {
239 TensorFactory<DTYPE> tf;
240 for (const auto& test_case : test_cases) {
241 SCOPED_TRACE(test_case.title); // Printed if the test fails
242
243 Tensor in = tf.make(test_case.sizes, test_case.input_data);
244 exec_aten::optional<Tensor> weight, bias;
245 if (!test_case.weight_data.empty()) {
246 weight = tf.make(test_case.normalized_shape, test_case.weight_data);
247 }
248 if (!test_case.bias_data.empty()) {
249 bias = tf.make(test_case.normalized_shape, test_case.bias_data);
250 }
251 Tensor out0 = tf.zeros(test_case.sizes);
252 Tensor out1 = tf.zeros(test_case.sizes);
253 Tensor out2 = tf.zeros(test_case.sizes);
254 auto normalized_shape_vec = std::vector<int64_t>(
255 test_case.normalized_shape.begin(), test_case.normalized_shape.end());
256 auto normalized_shape = exec_aten::ArrayRef<int64_t>(
257 normalized_shape_vec.data(), normalized_shape_vec.size());
258 ET_EXPECT_KERNEL_FAILURE(
259 context_,
260 op_native_layer_norm_out(
261 in,
262 normalized_shape,
263 weight,
264 bias,
265 test_case.eps,
266 out0,
267 out1,
268 out2));
269 }
270 }
271
272 // Test cases with imcompatible types.
273 template <ScalarType DTYPE>
run_int_test_cases()274 void run_int_test_cases() {
275 std::vector<NativeLayerNormTestCase<DTYPE>> test_cases = {
276 {
277 std::string(__func__) + ": Simple negative/positive layer norm",
278 // Cannot be represented by a type other than float.
279 {2, 3}, // sizes
280 {1, 0, -1, -1, 4, 0}, // input_data
281 {3}, // normalized shape
282 {1, 1, 1}, // weights
283 {0, 0, 0}, // bias
284 1, // eps
285 {0, 0, 0, 0, 0, 0}, // expected_data
286 },
287 };
288 run_death_test_cases(test_cases);
289 }
290
291 // Test cases with wrong normalized shape.
292 template <ScalarType DTYPE>
run_wrong_shape_test_cases()293 void run_wrong_shape_test_cases() {
294 std::vector<NativeLayerNormTestCase<DTYPE>> test_cases = {
295 {
296 std::string(__func__) + ": Test with wrong normalized shape",
297 {2, 3}, // sizes
298 {1.0, 0.0, -1.0, -1.0, 4.0, 0.0}, // input_data
299 {1}, // wrong normalized shape
300 {1.0}, // weights
301 {0.0}, // bias
302 1.0e-5, // eps
303 {1.22474,
304 0.0000,
305 -1.22474,
306 -0.925819,
307 1.38873,
308 -0.46291}, // expected_data
309 },
310 };
311 run_death_test_cases(test_cases);
312 }
313
314 /* %python
315 import torch
316 torch.manual_seed(0)
317
318 input = torch.rand(2, 3)
319 normalized_shape = [3]
320 weight = torch.tensor([1.0, 1.0, 1.0])
321 bias = torch.tensor([0.0, 0.0, 0.0])
322 eps = 1e-05
323 expected = torch.nn.functional.layer_norm(
324 input, normalized_shape, weight=weight, bias=bias, eps=eps)
325
326 native_layer_norm_template = f"""
327 {declare_tensor_factory("ScalarType::Float", "tf")}
328
329 {declare_tensor_make_t("input", "tf")}
330 {declare_optional_tensor_make_t("weight", "tf")}
331 {declare_optional_tensor_make_t("bias", "tf")}
332 {declare_tensor_make_t("expected", "tf")}
333 {declare_tensor_zeros("out_shape, dynamism", "tf", "out0")}
334 {declare_tensor_zeros("out_shape, dynamism", "tf", "out1")}
335 {declare_tensor_zeros("out_shape, dynamism", "tf", "out2")}
336
337 int64_t normalized_shape[] = $normalized_shape$;
338
339 op_native_layer_norm_out(
340 input, normalized_shape, weight, bias, $eps$, out0, out1, out2);
341 EXPECT_TENSOR_CLOSE(out0, expected);""" */
test_dynamic_shape(const std::vector<int32_t> & out_shape,enum torch::executor::TensorShapeDynamism dynamism)342 void test_dynamic_shape(
343 const std::vector<int32_t>& out_shape,
344 enum torch::executor::TensorShapeDynamism dynamism) {
345 /* %python
346 %rewrite(native_layer_norm_template) */
347
348 TensorFactory<ScalarType::Float> tf;
349
350 Tensor input = tf.make(
351 {2, 3},
352 {0.49625658988952637,
353 0.7682217955589294,
354 0.08847743272781372,
355 0.13203048706054688,
356 0.30742281675338745,
357 0.6340786814689636});
358 optional<Tensor> weight(tf.make({3}, {1.0, 1.0, 1.0}));
359 optional<Tensor> bias(tf.make({3}, {0.0, 0.0, 0.0}));
360 Tensor expected = tf.make(
361 {2, 3},
362 {0.16205203533172607,
363 1.1355723142623901,
364 -1.2976245880126953,
365 -1.0853172540664673,
366 -0.24233698844909668,
367 1.3276543617248535});
368 Tensor out0 = tf.zeros(out_shape, dynamism);
369 Tensor out1 = tf.zeros(out_shape, dynamism);
370 Tensor out2 = tf.zeros(out_shape, dynamism);
371
372 int64_t normalized_shape[] = {3};
373
374 op_native_layer_norm_out(
375 input, normalized_shape, weight, bias, 1e-05, out0, out1, out2);
376 EXPECT_TENSOR_CLOSE(out0, expected);
377 }
378 };
379
380 namespace {
vector_32_to_64(std::vector<int32_t> vector_32)381 std::vector<int64_t> vector_32_to_64(std::vector<int32_t> vector_32) {
382 std::vector<int64_t> vector_64(vector_32.size());
383 std::transform(
384 vector_32.begin(), vector_32.end(), vector_64.begin(), [](int32_t x) {
385 return static_cast<int64_t>(x);
386 });
387 return vector_64;
388 }
389
390 } // namespace
391
392 /// Describes a test case, using tensors of the specified DTYPE.
TEST_F(OpNativeLayerNormTest,FloatTensors)393 TEST_F(OpNativeLayerNormTest, FloatTensors) {
394 run_floating_point_test_cases<ScalarType::Float>();
395 run_floating_point_test_cases<ScalarType::Double>();
396 }
397
TEST_F(OpNativeLayerNormTest,IntTensorsDies)398 TEST_F(OpNativeLayerNormTest, IntTensorsDies) {
399 // Cannot be represented by a type other than float.
400 run_int_test_cases<ScalarType::Int>();
401 }
402
TEST_F(OpNativeLayerNormTest,WrongNomalizedShape)403 TEST_F(OpNativeLayerNormTest, WrongNomalizedShape) {
404 // Normalized shape does not match last dim of input.
405 run_wrong_shape_test_cases<ScalarType::Float>();
406 }
407
TEST_F(OpNativeLayerNormTest,DynamicShapeUpperBoundSameAsExpected)408 TEST_F(OpNativeLayerNormTest, DynamicShapeUpperBoundSameAsExpected) {
409 test_dynamic_shape(
410 {2, 3}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
411 }
412
TEST_F(OpNativeLayerNormTest,DynamicShapeUpperBoundLargerThanExpected)413 TEST_F(OpNativeLayerNormTest, DynamicShapeUpperBoundLargerThanExpected) {
414 test_dynamic_shape(
415 {10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
416 }
417
TEST_F(OpNativeLayerNormTest,DynamicShapeUnbound)418 TEST_F(OpNativeLayerNormTest, DynamicShapeUnbound) {
419 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
420 GTEST_SKIP() << "Dynamic shape unbound not supported";
421 }
422 test_dynamic_shape(
423 {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
424 }
425