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_t 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_t size_per_pack = sizeof(uint32) / sizeof(NativeT);
56 int64_t width = input.width();
57
58 int64_t 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_t size_per_pack = sizeof(uint32) / sizeof(NativeT);
83 int64_t width = input.width();
84
85 int64_t padded_output_width =
86 CeilOfRatio(width, size_per_pack) * size_per_pack;
87
88 int64_t output_height;
89 int64_t output_width;
90
91 if (transpose_output) {
92 output_height = padded_output_width;
93 output_width = input.height();
94 } else {
95 output_height = input.height();
96 output_width = padded_output_width;
97 }
98
99 Array2D<bfloat16> output(output_height, output_width, bfloat16(0.0));
100
101 float half_range =
102 !std::is_signed<NativeT>::value
103 ? 0.0f
104 : (static_cast<float>(std::numeric_limits<NativeT>::max() -
105 std::numeric_limits<NativeT>::min() + 1)) /
106 2.0f;
107 const bfloat16 scale_factor =
108 (range.max - range.min) /
109 (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
110 std::numeric_limits<NativeT>::min()));
111
112 for (int h = 0; h < input.height(); h++) {
113 std::vector<NativeT> input_row;
114 for (int w = 0; w < width; w++) {
115 bfloat16 result =
116 static_cast<bfloat16>(input(h, w) + half_range) * scale_factor +
117 range.min;
118 if (transpose_output) {
119 output(w, h) = result;
120 } else {
121 output(h, w) = result;
122 }
123 }
124 }
125
126 return output;
127 }
128
129 template <typename NativeT>
GenerateMinCombinedOutput(const QuantizedRange & range)130 std::vector<bfloat16> GenerateMinCombinedOutput(const QuantizedRange &range) {
131 float half_range =
132 !std::is_signed<NativeT>::value
133 ? 0.0f
134 : (static_cast<float>(std::numeric_limits<NativeT>::max() -
135 std::numeric_limits<NativeT>::min() + 1)) /
136 2.0f;
137 const bfloat16 scale_factor =
138 (range.max - range.min) /
139 (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
140 std::numeric_limits<NativeT>::min()));
141 std::vector<bfloat16> output;
142 for (int64_t i = std::numeric_limits<NativeT>::min();
143 i < std::numeric_limits<NativeT>::max(); ++i) {
144 bfloat16 result =
145 static_cast<bfloat16>(i + half_range) * scale_factor + range.min;
146 output.push_back(result);
147 }
148
149 const int64_t pack_size = sizeof(uint32) / sizeof(NativeT);
150 const int64_t output_size = output.size();
151
152 int64_t num_tailing_zeros =
153 CeilOfRatio(output_size, pack_size) * pack_size - output_size;
154
155 output.insert(output.end(), num_tailing_zeros, bfloat16(0.0));
156 return output;
157 }
158
159 // TODO(wangtao): add a test to make sure this op is the inverse of the existing
160 // TF quantize op defined in: third_party/tensorflow/core/kernels/quantize_op.cc
161
162 using DequantizeTest = ClientLibraryTestBase;
163
TEST(PackTest,PackUint8ToUint32)164 TEST(PackTest, PackUint8ToUint32) {
165 std::vector<uint8> input = {0xAB, 0x0B, 0x00, 0xF0, 0x01};
166 auto output = PackToUint32<uint8>(input);
167 EXPECT_THAT(output, ::testing::ElementsAre(0xAB0B00F0, 0x01000000));
168 }
169
TEST(PackTest,PackInt8ToUint32)170 TEST(PackTest, PackInt8ToUint32) {
171 std::vector<int8> input = {static_cast<signed char>(0x81), 0x0B, 0x00, 0x20,
172 0x01};
173 auto output = PackToUint32<int8>(input);
174 EXPECT_THAT(output, ::testing::ElementsAre(0x810B0020, 0x01000000));
175 }
176
TEST(PackTest,PackUint8ToUint32PerfectSize)177 TEST(PackTest, PackUint8ToUint32PerfectSize) {
178 std::vector<uint8> input = {3, 2, 1, 0};
179 auto output = PackToUint32<uint8>(input);
180 EXPECT_THAT(output, ::testing::ElementsAre(0x03020100));
181 }
182
XLA_TEST_F(DequantizeTest,MinCombinedUint16R1)183 XLA_TEST_F(DequantizeTest, MinCombinedUint16R1) {
184 XlaBuilder builder(TestName());
185 auto input = GenerateInput<uint16>();
186 auto x = ConstantR1<uint32>(&builder, PackToUint32<uint16>(input));
187 QuantizedRange range(0, 255.0f);
188 xla::Dequantize<uint16>(x, range, "MIN_COMBINED");
189 auto expected = GenerateMinCombinedOutput<uint16>(range);
190 ComputeAndCompareR1<bfloat16>(&builder, expected, {});
191 }
192
XLA_TEST_F(DequantizeTest,MinCombinedUint8R1)193 XLA_TEST_F(DequantizeTest, MinCombinedUint8R1) {
194 XlaBuilder builder(TestName());
195 auto input = GenerateInput<uint8>();
196 auto x = ConstantR1<uint32>(&builder, PackToUint32<uint8>(input));
197 QuantizedRange range(0, 127.0f);
198 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
199 auto expected = GenerateMinCombinedOutput<uint8>(range);
200 ComputeAndCompareR1<bfloat16>(&builder, expected, {});
201 }
202
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2)203 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2) {
204 XlaBuilder builder(TestName());
205 std::vector<std::vector<uint8>> input = {
206 {0, 1, 2, 3},
207 {4, 5, 6, 7},
208 {8, 9, 10, 11},
209 {12, 13, 16, 15},
210 };
211 auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
212 {PackToUint32<uint8>(input[1])[0]},
213 {PackToUint32<uint8>(input[2])[0]},
214 {PackToUint32<uint8>(input[3])[0]}});
215 QuantizedRange range(0, 255.0f);
216 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
217 const Array2D<bfloat16> expected = {
218 {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0)},
219 {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0)},
220 {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0)},
221 {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0)},
222 };
223 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
224 }
225
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TransposeOutput)226 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TransposeOutput) {
227 XlaBuilder builder(TestName());
228 std::vector<std::vector<uint8>> input = {
229 {0, 1, 2, 3},
230 {4, 5, 6, 7},
231 {8, 9, 10, 11},
232 {12, 13, 16, 15},
233 };
234 auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
235 {PackToUint32<uint8>(input[1])[0]},
236 {PackToUint32<uint8>(input[2])[0]},
237 {PackToUint32<uint8>(input[3])[0]}});
238 QuantizedRange range(0, 255.0f);
239 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
240 const Array2D<bfloat16> expected = {
241 {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
242 {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
243 {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
244 {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
245 };
246 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
247 }
248
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TailingZero)249 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZero) {
250 XlaBuilder builder(TestName());
251 std::vector<std::vector<uint8>> input = {
252 {0, 1, 2, 3, 16},
253 {4, 5, 6, 7, 17},
254 {8, 9, 10, 11, 18},
255 {12, 13, 16, 15, 19},
256 };
257 auto x = ConstantR2<uint32>(
258 &builder,
259 {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
260 {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
261 {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
262 {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
263 QuantizedRange range(0, 255.0f);
264 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
265
266 const Array2D<bfloat16> expected = {
267 {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0),
268 bfloat16(16.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
269 {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0),
270 bfloat16(17.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
271 {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0),
272 bfloat16(18.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
273 {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0),
274 bfloat16(19.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
275 };
276 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
277 }
278
XLA_TEST_F(DequantizeTest,MinCombinedUint8R2TailingZeroTransposeOutput)279 XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZeroTransposeOutput) {
280 XlaBuilder builder(TestName());
281 std::vector<std::vector<uint8>> input = {
282 {0, 1, 2, 3, 16},
283 {4, 5, 6, 7, 17},
284 {8, 9, 10, 11, 18},
285 {12, 13, 16, 15, 19},
286 };
287 auto x = ConstantR2<uint32>(
288 &builder,
289 {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
290 {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
291 {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
292 {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
293 QuantizedRange range(0, 255.0f);
294 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
295
296 const Array2D<bfloat16> expected = {
297 {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
298 {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
299 {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
300 {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
301 {bfloat16(16.0), bfloat16(17.0), bfloat16(18.0), bfloat16(19.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 {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
305 };
306 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
307 }
308
XLA_TEST_F(DequantizeTest,MinCombinedUint8LargeSizeTest)309 XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTest) {
310 XlaBuilder builder(TestName());
311 Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
312 Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
313
314 auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
315 QuantizedRange range(0, 255.0f);
316 xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
317
318 const Array2D<bfloat16> expected =
319 GenerateLargeSizeMinCombinedOutput<uint8>(input, range);
320 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
321 }
322
XLA_TEST_F(DequantizeTest,MinCombinedUint8LargeSizeTestTransposeOutput)323 XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTestTransposeOutput) {
324 XlaBuilder builder(TestName());
325 Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
326 Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
327
328 auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
329 QuantizedRange range(0, 255.0f);
330 xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
331
332 const Array2D<bfloat16> expected = GenerateLargeSizeMinCombinedOutput<uint8>(
333 input, range, /*transpose_output=*/true);
334 ComputeAndCompareR2<bfloat16>(&builder, expected, {});
335 }
336
337 } // namespace
338 } // namespace xla
339