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/quantized/NativeFunctions.h> // Declares the operator
10 #include <executorch/runtime/core/exec_aten/exec_aten.h>
11 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
12 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
13 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
14 #include <executorch/test/utils/DeathTest.h>
15
16 #include <gtest/gtest.h>
17 #include <limits>
18
19 using namespace ::testing;
20 using exec_aten::ArrayRef;
21 using exec_aten::Scalar;
22 using exec_aten::ScalarType;
23 using exec_aten::Tensor;
24 using torch::executor::native::quantize_per_channel_out;
25 using torch::executor::native::quantize_per_tensor_out;
26 using torch::executor::native::quantize_per_tensor_tensor_args_out;
27 using torch::executor::testing::TensorFactory;
28
29 /// A generic smoke test that works for any dtype that supports ones() and
30 /// zeros().
31 template <ScalarType DTYPE>
test_dtype()32 void test_dtype() {
33 TensorFactory<ScalarType::Float> tf;
34
35 Tensor input = tf.full({3, 5}, 4);
36 double scale = 0.5;
37
38 int64_t zero_point = 108;
39 int64_t quant_min = 0;
40 int64_t quant_max = 127;
41
42 TensorFactory<DTYPE> tfo;
43 Tensor out = tfo.zeros({3, 5});
44 // 4 / 0.5 + 127
45 Tensor expected = tfo.full({3, 5}, 116);
46 quantize_per_tensor_out(
47 input, scale, zero_point, quant_min, quant_max, DTYPE, out);
48
49 EXPECT_TENSOR_EQ(out, expected);
50 }
51
TEST(OpQuantizeOutTest,AllDtypesSupported)52 TEST(OpQuantizeOutTest, AllDtypesSupported) {
53 test_dtype<ScalarType::Byte>();
54 test_dtype<ScalarType::Char>();
55 test_dtype<ScalarType::Short>();
56 test_dtype<ScalarType::Bits16>();
57 test_dtype<ScalarType::UInt16>();
58 test_dtype<ScalarType::Int>();
59 }
60
TEST(OpQuantizeOutTest,TensorArgOverload)61 TEST(OpQuantizeOutTest, TensorArgOverload) {
62 TensorFactory<ScalarType::Float> tf_float;
63 TensorFactory<ScalarType::Double> tf_double;
64 TensorFactory<ScalarType::Long> tf_long;
65
66 Tensor input = tf_float.full({3, 5}, 4);
67 Tensor scale = tf_double.make({1}, {0.5});
68 Tensor zero_point = tf_long.make({1}, {127});
69 int64_t quant_min = 0;
70 int64_t quant_max = 255;
71
72 TensorFactory<ScalarType::Byte> tfo;
73 Tensor out = tfo.zeros({3, 5});
74 // 4 / 0.5 + 127
75 Tensor expected = tfo.full({3, 5}, 135);
76 auto context = torch::executor::KernelRuntimeContext();
77 quantize_per_tensor_tensor_args_out(
78 context,
79 input,
80 scale,
81 zero_point,
82 quant_min,
83 quant_max,
84 ScalarType::Byte,
85 out);
86
87 EXPECT_TENSOR_EQ(out, expected);
88 }
89
TEST(OpQuantizeOutTest,TestOutOfBounds)90 TEST(OpQuantizeOutTest, TestOutOfBounds) {
91 // Test where 1.0 / epsilon is larger than 8bit integer.
92
93 TensorFactory<ScalarType::Float> tf_float;
94 TensorFactory<ScalarType::Double> tf_double;
95 TensorFactory<ScalarType::Long> tf_long;
96
97 Tensor input = tf_float.ones({1, 3, 256, 256});
98
99 Tensor scale = tf_double.make({1}, {0.0011316323652863503});
100 Tensor zero_point = tf_long.make({1}, {0});
101 int64_t quant_min = -128;
102 int64_t quant_max = 127;
103
104 TensorFactory<ScalarType::Char> tfo;
105 Tensor out = tfo.zeros({1, 3, 256, 256});
106
107 Tensor expected = tfo.full({1, 3, 256, 256}, 127);
108
109 auto context = torch::executor::KernelRuntimeContext();
110 quantize_per_tensor_tensor_args_out(
111 context,
112 input,
113 scale,
114 zero_point,
115 quant_min,
116 quant_max,
117 ScalarType::Char,
118 out);
119
120 EXPECT_TENSOR_EQ(out, expected);
121 }
122
TEST(OpQuantizeOutTest,QuantizePerChannel)123 TEST(OpQuantizeOutTest, QuantizePerChannel) {
124 TensorFactory<ScalarType::Float> tf_float;
125 TensorFactory<ScalarType::Double> tf_double;
126 TensorFactory<ScalarType::Long> tf_long;
127
128 Tensor input = tf_float.full({3, 2}, 4);
129 Tensor scale = tf_double.make({2}, {0.5, 1});
130 Tensor zero_point = tf_long.make({2}, {127, 63});
131 int64_t quant_min = 0;
132 int64_t quant_max = 255;
133
134 TensorFactory<ScalarType::Byte> tfo;
135 Tensor out = tfo.zeros({3, 2});
136 // 4 / 0.5 + 127
137 // 4 / 1 + 63
138 Tensor expected = tfo.make({3, 2}, {135, 67, 135, 67, 135, 67});
139 quantize_per_channel_out(
140 input, scale, zero_point, 1, quant_min, quant_max, ScalarType::Byte, out);
141
142 EXPECT_TENSOR_EQ(out, expected);
143 }
144