1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/xla/client/lib/quantize.h"
17
18 #include <limits>
19
20 #include "tensorflow/compiler/xla/client/xla_builder.h"
21 #include "tensorflow/compiler/xla/test.h"
22 #include "tensorflow/compiler/xla/tests/client_library_test_base.h"
23 #include "tensorflow/compiler/xla/tests/test_macros.h"
24 #include "tensorflow/compiler/xla/types.h"
25 #include "tensorflow/compiler/xla/util.h"
26
27 namespace xla {
28 namespace {
29
30 using bfloat16 = tensorflow::bfloat16;
31
32 template <typename NativeT>
GenerateInput()33 std::vector<NativeT> GenerateInput() {
34 std::vector<NativeT> input;
35
36 for (int64 i = std::numeric_limits<NativeT>::min();
37 i < std::numeric_limits<NativeT>::max(); ++i) {
38 input.push_back(static_cast<NativeT>(i));
39 }
40
41 return input;
42 }
43
44 template <typename NativeT>
GenerateLargeSizeInput(int num_columns,int num_rows)45 Array2D<NativeT> GenerateLargeSizeInput(int num_columns, int num_rows) {
46 Array2D<NativeT> input(num_columns, num_rows);
47
48 input.FillRandom(6, 128);
49
50 return input;
51 }
52
53 template <typename NativeT>
PackLargeInput(Array2D<NativeT> & input)54 Array2D<uint32> PackLargeInput(Array2D<NativeT> &input) {
55 const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT);
56 int64 width = input.width();
57
58 int64 padded_output_width = CeilOfRatio(width, size_per_pack);
59
60 Array2D<uint32> pack_input(input.height(), padded_output_width);
61
62 for (int h = 0; h < input.height(); h++) {
63 std::vector<NativeT> input_row;
64 for (int w = 0; w < width; w++) {
65 input_row.push_back(input({h, w}));
66 }
67
68 auto pack_input_vec = PackToUint32<uint8>(input_row);
69
70 for (int w = 0; w < padded_output_width; w++) {
71 pack_input(h, w) = pack_input_vec[w];
72 }
73 }
74
75 return pack_input;
76 }
77
78 template <typename NativeT>
GenerateLargeSizeMinCombinedOutput(Array2D<NativeT> & input,const QuantizedRange & range,bool transpose_output=false)79 Array2D<bfloat16> GenerateLargeSizeMinCombinedOutput(
80 Array2D<NativeT> &input, const QuantizedRange &range,
81 bool transpose_output = false) {
82 const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT);
83 int64 width = input.width();
84
85 int64 padded_output_width = CeilOfRatio(width, size_per_pack) * size_per_pack;
86
87 int64 output_height;
88 int64 output_width;
89
90 if (transpose_output) {
91 output_height = padded_output_width;
92 output_width = input.height();
93 } else {
94 output_height = input.height();
95 output_width = padded_output_width;
96 }
97
98 Array2D<bfloat16> output(output_height, output_width, bfloat16(0.0));
99
100 float half_range =
101 !std::is_signed<NativeT>::value
102 ? 0.0f
103 : (static_cast<float>(std::numeric_limits<NativeT>::max() -
104 std::numeric_limits<NativeT>::min() + 1)) /
105 2.0f;
106 const bfloat16 scale_factor =
107 (range.max - range.min) /
108 (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
109 std::numeric_limits<NativeT>::min()));
110
111 for (int h = 0; h < input.height(); h++) {
112 std::vector<NativeT> input_row;
113 for (int w = 0; w < width; w++) {
114 bfloat16 result =
115 static_cast<bfloat16>(input(h, w) + half_range) * scale_factor +
116 range.min;
117 if (transpose_output) {
118 output(w, h) = result;
119 } else {
120 output(h, w) = result;
121 }
122 }
123 }
124
125 return output;
126 }
127
128 template <typename NativeT>
GenerateMinCombinedOutput(const QuantizedRange & range)129 std::vector<bfloat16> GenerateMinCombinedOutput(const QuantizedRange &range) {
130 float half_range =
131 !std::is_signed<NativeT>::value
132 ? 0.0f
133 : (static_cast<float>(std::numeric_limits<NativeT>::max() -
134 std::numeric_limits<NativeT>::min() + 1)) /
135 2.0f;
136 const bfloat16 scale_factor =
137 (range.max - range.min) /
138 (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
139 std::numeric_limits<NativeT>::min()));
140 std::vector<bfloat16> output;
141 for (int64 i = std::numeric_limits<NativeT>::min();
142 i < std::numeric_limits<NativeT>::max(); ++i) {
143 bfloat16 result =
144 static_cast<bfloat16>(i + half_range) * scale_factor + range.min;
145 output.push_back(result);
146 }
147
148 const int64 pack_size = sizeof(uint32) / sizeof(NativeT);
149 const int64 output_size = output.size();
150
151 int64 num_tailing_zeros =
152 CeilOfRatio(output_size, pack_size) * pack_size - output_size;
153
154 output.insert(output.end(), num_tailing_zeros, bfloat16(0.0));
155 return output;
156 }
157
158 // TODO(wangtao): add a test to make sure this op is the inverse of the existing
159 // TF quantize op defined in: third_party/tensorflow/core/kernels/quantize_op.cc
160
161 using DequantizeTest = ClientLibraryTestBase;
162
TEST(PackTest,PackUint8ToUint32)163 TEST(PackTest, PackUint8ToUint32) {
164 std::vector<uint8> input = {0xAB, 0x0B, 0x00, 0xF0, 0x01};
165 auto output = PackToUint32<uint8>(input);
166 EXPECT_THAT(output, ::testing::ElementsAre(0xAB0B00F0, 0x01000000));
167 }
168
TEST(PackTest,PackInt8ToUint32)169 TEST(PackTest, PackInt8ToUint32) {
170 std::vector<int8> input = {static_cast<signed char>(0x81), 0x0B, 0x00, 0x20,
171 0x01};
172 auto output = PackToUint32<int8>(input);
173 EXPECT_THAT(output, ::testing::ElementsAre(0x810B0020, 0x01000000));
174 }
175
TEST(PackTest,PackUint8ToUint32PerfectSize)176 TEST(PackTest, PackUint8ToUint32PerfectSize) {
177 std::vector<uint8> input = {3, 2, 1, 0};
178 auto output = PackToUint32<uint8>(input);
179 EXPECT_THAT(output, ::testing::ElementsAre(0x03020100));
180 }
181
XLA_TEST_F(DequantizeTest,MinCombinedUint16R1)182 XLA_TEST_F(DequantizeTest, MinCombinedUint16R1) {
183 XlaBuilder builder(TestName());
184 auto input = GenerateInput<uint16>();
185 auto x = ConstantR1<uint32>(&builder, PackToUint32<uint16>(input));
186 QuantizedRange range(0, 255.0f);
187 xla::Dequantize<uint16>(x, range, "MIN_COMBINED");
188 auto expected = GenerateMinCombinedOutput<uint16>(range);
189 ComputeAndCompareR1<bfloat16>(&builder, expected, {});
190 }
191
XLA_TEST_F(DequantizeTest,MinCombinedUint8R1)192 XLA_TEST_F(DequantizeTest, MinCombinedUint8R1) {
193 XlaBuilder builder(TestName());
194 auto input = GenerateInput<uint8>();
195 auto x = ConstantR1<uint32>(&builder, PackToUint32<uint8>(input));
196 QuantizedRange range(0, 127.0f);
197 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
198 auto expected = GenerateMinCombinedOutput<uint8>(range);
199 ComputeAndCompareR1<bfloat16>(&builder, expected, {});
200 }
201
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2)202 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2) {
203 XlaBuilder builder(TestName());
204 std::vector<std::vector<uint8>> input = {
205 {0, 1, 2, 3},
206 {4, 5, 6, 7},
207 {8, 9, 10, 11},
208 {12, 13, 16, 15},
209 };
210 auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
211 {PackToUint32<uint8>(input[1])[0]},
212 {PackToUint32<uint8>(input[2])[0]},
213 {PackToUint32<uint8>(input[3])[0]}});
214 QuantizedRange range(0, 255.0f);
215 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
216 const Array2D<bfloat16> expected = {
217 {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0)},
218 {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0)},
219 {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0)},
220 {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0)},
221 };
222 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
223 }
224
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TransposeOutput)225 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TransposeOutput) {
226 XlaBuilder builder(TestName());
227 std::vector<std::vector<uint8>> input = {
228 {0, 1, 2, 3},
229 {4, 5, 6, 7},
230 {8, 9, 10, 11},
231 {12, 13, 16, 15},
232 };
233 auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
234 {PackToUint32<uint8>(input[1])[0]},
235 {PackToUint32<uint8>(input[2])[0]},
236 {PackToUint32<uint8>(input[3])[0]}});
237 QuantizedRange range(0, 255.0f);
238 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
239 const Array2D<bfloat16> expected = {
240 {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
241 {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
242 {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
243 {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
244 };
245 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
246 }
247
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TailingZero)248 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZero) {
249 XlaBuilder builder(TestName());
250 std::vector<std::vector<uint8>> input = {
251 {0, 1, 2, 3, 16},
252 {4, 5, 6, 7, 17},
253 {8, 9, 10, 11, 18},
254 {12, 13, 16, 15, 19},
255 };
256 auto x = ConstantR2<uint32>(
257 &builder,
258 {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
259 {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
260 {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
261 {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
262 QuantizedRange range(0, 255.0f);
263 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
264
265 const Array2D<bfloat16> expected = {
266 {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0),
267 bfloat16(16.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
268 {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0),
269 bfloat16(17.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
270 {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0),
271 bfloat16(18.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
272 {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0),
273 bfloat16(19.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
274 };
275 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
276 }
277
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TailingZeroTransposeOutput)278 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZeroTransposeOutput) {
279 XlaBuilder builder(TestName());
280 std::vector<std::vector<uint8>> input = {
281 {0, 1, 2, 3, 16},
282 {4, 5, 6, 7, 17},
283 {8, 9, 10, 11, 18},
284 {12, 13, 16, 15, 19},
285 };
286 auto x = ConstantR2<uint32>(
287 &builder,
288 {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
289 {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
290 {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
291 {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
292 QuantizedRange range(0, 255.0f);
293 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
294
295 const Array2D<bfloat16> expected = {
296 {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
297 {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
298 {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
299 {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
300 {bfloat16(16.0), bfloat16(17.0), bfloat16(18.0), bfloat16(19.0)},
301 {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
302 {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
303 {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
304 };
305 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
306 }
307
XLA_TEST_F(DequantizeTest,MinCombinedUint8LargeSizeTest)308 XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTest) {
309 XlaBuilder builder(TestName());
310 Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
311 Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
312
313 auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
314 QuantizedRange range(0, 255.0f);
315 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
316
317 const Array2D<bfloat16> expected =
318 GenerateLargeSizeMinCombinedOutput<uint8>(input, range);
319 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
320 }
321
XLA_TEST_F(DequantizeTest,MinCombinedUint8LargeSizeTestTransposeOutput)322 XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTestTransposeOutput) {
323 XlaBuilder builder(TestName());
324 Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
325 Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
326
327 auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
328 QuantizedRange range(0, 255.0f);
329 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
330
331 const Array2D<bfloat16> expected = GenerateLargeSizeMinCombinedOutput<uint8>(
332 input, range, /*transpose_output=*/true);
333 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
334 }
335
336 } // namespace
337 } // namespace xla
338