1 /* Copyright 2017 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 <memory>
17 #include <vector>
18
19 #include "tensorflow/compiler/xla/array2d.h"
20 #include "tensorflow/compiler/xla/array4d.h"
21 #include "tensorflow/compiler/xla/client/lib/arithmetic.h"
22 #include "tensorflow/compiler/xla/client/local_client.h"
23 #include "tensorflow/compiler/xla/client/xla_builder.h"
24 #include "tensorflow/compiler/xla/reference_util.h"
25 #include "tensorflow/compiler/xla/tests/client_library_test_base.h"
26 #include "tensorflow/compiler/xla/tests/literal_test_util.h"
27 #include "tensorflow/compiler/xla/tests/test_macros.h"
28 #include "tensorflow/compiler/xla/xla_data.pb.h"
29 #include "tensorflow/core/platform/test.h"
30
31 namespace xla {
32 namespace {
33
34 #ifdef XLA_BACKEND_SUPPORTS_BFLOAT16
35 // Tests both F32 and BF16.
36 static std::array<bool, 2> use_bfloat16_params{false, true};
37 #else
38 // Only tests F32.
39 static std::array<bool, 1> use_bfloat16_params{false};
40 #endif
41
42 class PadTest : public ClientLibraryTestBase {
43 protected:
PadTest()44 PadTest() {
45 // Initializes the padding configuration used for R4 tests.
46 // Pad only on the dimension 0 {low: 1, high: 0, interior: 2} and
47 // dimension 1 {low: 0, high: 2, interior: 1}.
48 auto dimension0 = r4_padding_on_dim0_dim1_.add_dimensions();
49 dimension0->set_edge_padding_low(1);
50 dimension0->set_edge_padding_high(0);
51 dimension0->set_interior_padding(2);
52 auto dimension1 = r4_padding_on_dim0_dim1_.add_dimensions();
53 dimension1->set_edge_padding_low(0);
54 dimension1->set_edge_padding_high(2);
55 dimension1->set_interior_padding(1);
56 auto dimension2 = r4_padding_on_dim0_dim1_.add_dimensions();
57 dimension2->set_edge_padding_low(0);
58 dimension2->set_edge_padding_high(0);
59 dimension2->set_interior_padding(0);
60 auto dimension3 = r4_padding_on_dim0_dim1_.add_dimensions();
61 dimension3->set_edge_padding_low(0);
62 dimension3->set_edge_padding_high(0);
63 dimension3->set_interior_padding(0);
64 }
65
66 // Padding configuration for R4 that only pads dimension 0 and 1.
67 PaddingConfig r4_padding_on_dim0_dim1_;
68 };
69
70 class PadTestFloat : public PadTest,
71 public ::testing::WithParamInterface<bool> {
72 protected:
PadTestFloat()73 PadTestFloat() { set_use_bfloat16(GetParam()); }
74
DefaultErrorSpec() const75 ErrorSpec DefaultErrorSpec() const {
76 if (use_bfloat16()) {
77 return ErrorSpec(1e-3, 1e-3);
78 } else {
79 return ErrorSpec(1e-5, 1e-5);
80 }
81 }
82 };
83
84 // Tests a Pad() with a zero-element input and output.
XLA_TEST_P(PadTestFloat,Pad1DS0ToS0Array)85 XLA_TEST_P(PadTestFloat, Pad1DS0ToS0Array) {
86 XlaBuilder b(TestName());
87 // Set up the padding configuration {low: 0, high: 0, interior: 0}.
88 PaddingConfig padding_config;
89 auto dimension = padding_config.add_dimensions();
90 dimension->set_edge_padding_low(0);
91 dimension->set_edge_padding_high(0);
92 dimension->set_interior_padding(0);
93
94 Pad(AddParam(LiteralUtil::CreateR1<float>({}), &b),
95 AddParam(LiteralUtil::CreateR0<float>(0.1), &b), padding_config);
96 ComputeAndCompareR1<float>(&b, {}, {}, DefaultErrorSpec());
97 }
98
99 // Tests a Pad() with a zero-element input but a non-zero-element output.
XLA_TEST_P(PadTestFloat,Pad1DS0ToS5Array)100 XLA_TEST_P(PadTestFloat, Pad1DS0ToS5Array) {
101 XlaBuilder b(TestName());
102 // Set up the padding configuration {low: 3, high: 0, interior: 1}.
103 PaddingConfig padding_config;
104 auto dimension = padding_config.add_dimensions();
105 dimension->set_edge_padding_low(1);
106 dimension->set_edge_padding_high(4);
107 dimension->set_interior_padding(7);
108
109 Pad(AddParam(LiteralUtil::CreateR1<float>({}), &b),
110 AddParam(LiteralUtil::CreateR0<float>(0.1), &b), padding_config);
111 ComputeAndCompareR1<float>(&b, std::vector<float>(5, 0.1), {},
112 DefaultErrorSpec());
113 }
114
XLA_TEST_P(PadTestFloat,Pad1DS3Array)115 XLA_TEST_P(PadTestFloat, Pad1DS3Array) {
116 XlaBuilder b(TestName());
117 // Set up the padding configuration {low: 3, high: 0, interior: 1}.
118 PaddingConfig padding_config;
119 auto dimension = padding_config.add_dimensions();
120 dimension->set_edge_padding_low(3);
121 dimension->set_edge_padding_high(0);
122 dimension->set_interior_padding(1);
123
124 Pad(AddParam(LiteralUtil::CreateR1<float>({1, 2, 3}), &b),
125 AddParam(LiteralUtil::CreateR0<float>(0.1), &b), padding_config);
126 std::vector<float> expected({0.1, 0.1, 0.1, 1, 0.1, 2, 0.1, 3});
127 ComputeAndCompareR1<float>(&b, expected, {}, DefaultErrorSpec());
128 }
129
XLA_TEST_P(PadTestFloat,Pad4D_2x0x3x2_FloatArray)130 XLA_TEST_P(PadTestFloat, Pad4D_2x0x3x2_FloatArray) {
131 XlaBuilder b(TestName());
132 Pad(AddParam(Array4D<float>(2, 0, 3, 2), &b),
133 AddParam(LiteralUtil::CreateR0<float>(1.5), &b),
134 r4_padding_on_dim0_dim1_);
135 ComputeAndCompareR4<float>(&b, Array4D<float>(5, 2, 3, 2, 1.5f), {},
136 DefaultErrorSpec());
137 }
138
TEST_P(PadTestFloat,Pad4DFloat_1x1x3x2_Array)139 TEST_P(PadTestFloat, Pad4DFloat_1x1x3x2_Array) {
140 XlaBuilder b(TestName());
141 auto input = std::make_unique<Array4D<float>>(1, 1, 3, 2);
142 Array2D<float> input_xy({
143 {1.0f, 2.0f}, // row 0
144 {3.0f, 4.0f}, // row 1
145 {5.0f, 6.0f}, // row 2
146 });
147 input->FillWithYX(input_xy);
148
149 Pad(AddParam(*input, &b), AddParam(LiteralUtil::CreateR0<float>(1.5), &b),
150 r4_padding_on_dim0_dim1_);
151
152 auto expected = std::make_unique<Array4D<float>>(2, 3, 3, 2);
153 expected->Fill(1.5);
154 (*expected)(1, 0, 0, 0) = 1.0f;
155 (*expected)(1, 0, 0, 1) = 2.0f;
156 (*expected)(1, 0, 1, 0) = 3.0f;
157 (*expected)(1, 0, 1, 1) = 4.0f;
158 (*expected)(1, 0, 2, 0) = 5.0f;
159 (*expected)(1, 0, 2, 1) = 6.0f;
160 ComputeAndCompareR4<float>(&b, *expected, {}, DefaultErrorSpec());
161 }
162
TEST_P(PadTestFloat,Pad4DFloatArrayWithInteriorPadding)163 TEST_P(PadTestFloat, Pad4DFloatArrayWithInteriorPadding) {
164 XlaBuilder b(TestName());
165
166 const float pad_value = 1.5f;
167 Array4D<float> input(3, 2, 1, 1, {1, 2, 3, 4, 5, 6});
168 Pad(AddParam(input, &b),
169 AddParam(LiteralUtil::CreateR0<float>(pad_value), &b),
170 r4_padding_on_dim0_dim1_);
171
172 auto expected = std::make_unique<Array4D<float>>(8, 5, 1, 1);
173 expected->Fill(pad_value);
174 (*expected)(1, 0, 0, 0) = 1.0f;
175 (*expected)(1, 2, 0, 0) = 2.0f;
176 (*expected)(4, 0, 0, 0) = 3.0f;
177 (*expected)(4, 2, 0, 0) = 4.0f;
178 (*expected)(7, 0, 0, 0) = 5.0f;
179 (*expected)(7, 2, 0, 0) = 6.0f;
180 ComputeAndCompareR4<float>(&b, *expected, {}, ErrorSpec(0.0001));
181 }
182
TEST_P(PadTestFloat,Pad4DFloatArrayMinorFirstSmall)183 TEST_P(PadTestFloat, Pad4DFloatArrayMinorFirstSmall) {
184 XlaBuilder b(TestName());
185
186 PaddingConfig padding_config;
187 auto dimension0 = padding_config.add_dimensions();
188 dimension0->set_edge_padding_low(0);
189 dimension0->set_edge_padding_high(0);
190 dimension0->set_interior_padding(0);
191 auto dimension1 = padding_config.add_dimensions();
192 dimension1->set_edge_padding_low(0);
193 dimension1->set_edge_padding_high(0);
194 dimension1->set_interior_padding(0);
195 auto dimension2 = padding_config.add_dimensions();
196 dimension2->set_edge_padding_low(2);
197 dimension2->set_edge_padding_high(1);
198 dimension2->set_interior_padding(0);
199 auto dimension3 = padding_config.add_dimensions();
200 dimension3->set_edge_padding_low(2);
201 dimension3->set_edge_padding_high(3);
202 dimension3->set_interior_padding(0);
203
204 const Layout layout = LayoutUtil::MakeLayout({0, 1, 2, 3});
205
206 const float pad_value = -5.123f;
207 Array4D<float> input_array(1, 1, 2, 3, {1, 2, 3, 4, 5, 6});
208 auto input = LiteralUtil::CreateR4FromArray4D<float>(input_array);
209 input = input.Relayout(layout);
210
211 Pad(AddParam(input, &b),
212 AddParam(LiteralUtil::CreateR0<float>(pad_value), &b), padding_config);
213
214 Array4D<float> expected_array(1, 1, 5, 8);
215 expected_array.Fill(pad_value);
216 expected_array(0, 0, 2, 2) = 1.0f;
217 expected_array(0, 0, 2, 3) = 2.0f;
218 expected_array(0, 0, 2, 4) = 3.0f;
219 expected_array(0, 0, 3, 2) = 4.0f;
220 expected_array(0, 0, 3, 3) = 5.0f;
221 expected_array(0, 0, 3, 4) = 6.0f;
222 ComputeAndCompareR4<float>(&b, expected_array, {}, ErrorSpec(0.0001));
223 }
224
XLA_TEST_P(PadTestFloat,Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions)225 XLA_TEST_P(PadTestFloat, Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions) {
226 XlaBuilder b(TestName());
227
228 PaddingConfig padding_config;
229 auto dimension0 = padding_config.add_dimensions();
230 dimension0->set_edge_padding_low(0);
231 dimension0->set_edge_padding_high(0);
232 dimension0->set_interior_padding(0);
233 auto dimension1 = padding_config.add_dimensions();
234 dimension1->set_edge_padding_low(0);
235 dimension1->set_edge_padding_high(0);
236 dimension1->set_interior_padding(0);
237 auto dimension2 = padding_config.add_dimensions();
238 dimension2->set_edge_padding_low(2);
239 dimension2->set_edge_padding_high(2);
240 dimension2->set_interior_padding(1);
241 auto dimension3 = padding_config.add_dimensions();
242 dimension3->set_edge_padding_low(2);
243 dimension3->set_edge_padding_high(2);
244 dimension3->set_interior_padding(0);
245
246 const Layout layout = LayoutUtil::MakeLayout({0, 1, 2, 3});
247
248 const float pad_value = -5.123f;
249 Array4D<float> input_array(1, 25, 7, 7);
250 input_array.Fill(pad_value);
251 input_array(0, 0, 0, 0) = 1.0f;
252 input_array(0, 24, 6, 6) = 2.0f;
253 input_array(0, 17, 2, 5) = 3.0f;
254 auto input = LiteralUtil::CreateR4FromArray4D<float>(input_array);
255 input = input.Relayout(layout);
256
257 Pad(AddParam(input, &b),
258 AddParam(LiteralUtil::CreateR0<float>(pad_value), &b), padding_config);
259
260 Array4D<float> expected_array(1, 25, 17, 11);
261 expected_array.Fill(pad_value);
262 expected_array(0, 0, 2, 2) = 1.0f;
263 expected_array(0, 24, 14, 8) = 2.0f;
264 expected_array(0, 17, 6, 7) = 3.0f;
265 ComputeAndCompareR4<float>(&b, expected_array, {}, ErrorSpec(0.0001));
266 }
267
XLA_TEST_F(PadTest,Pad4DU8Array)268 XLA_TEST_F(PadTest, Pad4DU8Array) {
269 XlaBuilder b(TestName());
270 auto input = std::make_unique<Array4D<uint8_t>>(1, 1, 3, 2);
271 Array2D<uint8_t> input_xy({
272 {1, 2}, // row 0
273 {3, 4}, // row 1
274 {5, 6}, // row 2
275 });
276 input->FillWithYX(input_xy);
277
278 Pad(AddParam(*input, &b), ConstantR0<uint8_t>(&b, 35),
279 r4_padding_on_dim0_dim1_);
280
281 auto expected = std::make_unique<Array4D<uint8_t>>(2, 3, 3, 2);
282 expected->Fill(35);
283 (*expected)(1, 0, 0, 0) = 1;
284 (*expected)(1, 0, 0, 1) = 2;
285 (*expected)(1, 0, 1, 0) = 3;
286 (*expected)(1, 0, 1, 1) = 4;
287 (*expected)(1, 0, 2, 0) = 5;
288 (*expected)(1, 0, 2, 1) = 6;
289 ComputeAndCompareR4<uint8_t>(&b, *expected, {});
290 }
291
XLA_TEST_F(PadTest,Pad4DPredArray)292 XLA_TEST_F(PadTest, Pad4DPredArray) {
293 XlaBuilder b(TestName());
294
295 // Since bool is currently not well supported, use Broadcast operation to
296 // create the operand for Pad.
297 auto input = Broadcast(ConstantR0<bool>(&b, true), {1, 1, 3, 2});
298 auto padded =
299 Pad(input, ConstantR0<bool>(&b, false), r4_padding_on_dim0_dim1_);
300
301 // For the same reason, use Select to convert boolean values to int32_t.
302 auto zeros = std::make_unique<Array4D<int32_t>>(2, 3, 3, 2);
303 auto ones = std::make_unique<Array4D<int32_t>>(2, 3, 3, 2);
304 zeros->Fill(0);
305 ones->Fill(1);
306 Select(padded, AddParam(*ones, &b), AddParam(*zeros, &b));
307
308 auto expected = std::make_unique<Array4D<int32_t>>(2, 3, 3, 2);
309 expected->Fill(0);
310 (*expected)(1, 0, 0, 0) = 1;
311 (*expected)(1, 0, 0, 1) = 1;
312 (*expected)(1, 0, 1, 0) = 1;
313 (*expected)(1, 0, 1, 1) = 1;
314 (*expected)(1, 0, 2, 0) = 1;
315 (*expected)(1, 0, 2, 1) = 1;
316 ComputeAndCompareR4<int32_t>(&b, *expected, {});
317 }
318
XLA_TEST_P(PadTestFloat,Large2DPad)319 XLA_TEST_P(PadTestFloat, Large2DPad) {
320 XlaBuilder b(TestName());
321
322 auto ones = std::make_unique<Array2D<float>>(4, 4);
323 ones->Fill(1.0f);
324 auto input = AddParam(*ones, &b);
325 PaddingConfig padding_config = MakeNoPaddingConfig(2);
326 for (int dim : {0, 1}) {
327 padding_config.mutable_dimensions(dim)->set_edge_padding_low(
328 98 + 100 * (1 - dim));
329 padding_config.mutable_dimensions(dim)->set_edge_padding_high(58 +
330 100 * dim);
331 }
332 Pad(input, AddParam(LiteralUtil::CreateR0<float>(0.0f), &b), padding_config);
333
334 auto expected = ReferenceUtil::PadArray2D(*ones, padding_config, 0.0f);
335 ComputeAndCompareR2<float>(&b, *expected, {}, DefaultErrorSpec());
336 }
337
XLA_TEST_P(PadTestFloat,AllTypes2DPad)338 XLA_TEST_P(PadTestFloat, AllTypes2DPad) {
339 XlaBuilder b(TestName());
340
341 constexpr int64_t in_rows = 35;
342 constexpr int64_t in_cols = 35;
343 auto operand = std::make_unique<Array2D<float>>(in_rows, in_cols);
344 operand->FillUnique(0.0f);
345 auto input = AddParam(*operand, &b);
346
347 PaddingConfig padding_config = MakeNoPaddingConfig(2);
348 padding_config.mutable_dimensions(0)->set_edge_padding_low(7);
349 padding_config.mutable_dimensions(0)->set_edge_padding_high(5);
350 padding_config.mutable_dimensions(0)->set_interior_padding(3);
351 padding_config.mutable_dimensions(1)->set_edge_padding_low(6);
352 padding_config.mutable_dimensions(1)->set_edge_padding_high(4);
353 padding_config.mutable_dimensions(1)->set_interior_padding(2);
354 Pad(input, AddParam(LiteralUtil::CreateR0<float>(3.14f), &b), padding_config);
355
356 auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 3.14f);
357 ComputeAndCompareR2<float>(&b, *expected, {}, DefaultErrorSpec());
358 }
359
XLA_TEST_P(PadTestFloat,High2DPad)360 XLA_TEST_P(PadTestFloat, High2DPad) {
361 XlaBuilder b(TestName());
362
363 constexpr int64_t in_rows = 129;
364 constexpr int64_t in_cols = 129;
365 constexpr int64_t low_padding = 0;
366 int64_t high_padding[2] = {5, 7};
367 constexpr int64_t interior_padding = 0;
368 auto operand = std::make_unique<Array2D<float>>(in_rows, in_cols);
369 operand->FillUnique(1.0f);
370 auto input = AddParam(*operand, &b);
371 PaddingConfig padding_config = MakeNoPaddingConfig(2);
372 for (int dim : {0, 1}) {
373 padding_config.mutable_dimensions(dim)->set_edge_padding_low(low_padding);
374 padding_config.mutable_dimensions(dim)->set_edge_padding_high(
375 high_padding[dim]);
376 padding_config.mutable_dimensions(dim)->set_interior_padding(
377 interior_padding);
378 }
379 Pad(input, AddParam(LiteralUtil::CreateR0<float>(2.718f), &b),
380 padding_config);
381
382 auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f);
383
384 ComputeAndCompareR2<float>(&b, *expected, {}, DefaultErrorSpec());
385 }
386
XLA_TEST_P(PadTestFloat,NegativePadding2D)387 XLA_TEST_P(PadTestFloat, NegativePadding2D) {
388 XlaBuilder b(TestName());
389
390 constexpr int64_t in_rows = 129;
391 constexpr int64_t in_cols = 129;
392 int64_t low_padding[2] = {-1, -2};
393 int64_t high_padding[2] = {-3, 4};
394 constexpr int64_t interior_padding = 0;
395 auto operand = std::make_unique<Array2D<float>>(in_rows, in_cols);
396 operand->FillUnique(1.0f);
397 auto input = AddParam(*operand, &b);
398 PaddingConfig padding_config = MakeNoPaddingConfig(2);
399 for (int dim : {0, 1}) {
400 padding_config.mutable_dimensions(dim)->set_edge_padding_low(
401 low_padding[dim]);
402 padding_config.mutable_dimensions(dim)->set_edge_padding_high(
403 high_padding[dim]);
404 padding_config.mutable_dimensions(dim)->set_interior_padding(
405 interior_padding);
406 }
407 Pad(input, AddParam(LiteralUtil::CreateR0<float>(2.718f), &b),
408 padding_config);
409
410 auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f);
411
412 ComputeAndCompareR2<float>(&b, *expected, {}, DefaultErrorSpec());
413 }
414
XLA_TEST_P(PadTestFloat,NegativeAndInteriorPadding2D)415 XLA_TEST_P(PadTestFloat, NegativeAndInteriorPadding2D) {
416 XlaBuilder b(TestName());
417
418 constexpr int64_t in_rows = 8;
419 constexpr int64_t in_cols = 11;
420 int64_t low_padding[2] = {4, -1};
421 int64_t high_padding[2] = {-2, -4};
422 int64_t interior_padding[2] = {1, 2};
423 auto operand = std::make_unique<Array2D<float>>(in_rows, in_cols);
424 operand->FillUnique(1.0f);
425 auto input = AddParam(*operand, &b);
426 PaddingConfig padding_config = MakeNoPaddingConfig(2);
427 for (int dim : {0, 1}) {
428 padding_config.mutable_dimensions(dim)->set_edge_padding_low(
429 low_padding[dim]);
430 padding_config.mutable_dimensions(dim)->set_edge_padding_high(
431 high_padding[dim]);
432 padding_config.mutable_dimensions(dim)->set_interior_padding(
433 interior_padding[dim]);
434 }
435 Pad(input, AddParam(LiteralUtil::CreateR0<float>(2.718f), &b),
436 padding_config);
437
438 auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f);
439
440 ComputeAndCompareR2<float>(&b, *expected, {}, DefaultErrorSpec());
441 }
442
443 // Regression test for b/31827337.
XLA_TEST_P(PadTestFloat,ReducePad)444 XLA_TEST_P(PadTestFloat, ReducePad) {
445 XlaBuilder b(TestName());
446 auto ones = std::make_unique<Array4D<float>>(2, 2, 2, 2);
447 ones->Fill(1.0);
448 auto input = AddParam(*ones, &b);
449
450 XlaComputation add = CreateScalarAddComputation(FloatType(), &b);
451 auto reduce =
452 Reduce(input, AddParam(LiteralUtil::CreateR0<float>(0.0), &b), add, {0});
453
454 PaddingConfig padding_config = MakeNoPaddingConfig(3);
455 padding_config.mutable_dimensions(0)->set_edge_padding_low(1);
456 padding_config.mutable_dimensions(0)->set_edge_padding_high(1);
457 Pad(reduce, AddParam(LiteralUtil::CreateR0<float>(0.0f), &b), padding_config);
458
459 Array3D<float> expected({{{0.0, 0.0}, {0.0, 0.0}},
460 {{2.0, 2.0}, {2.0, 2.0}},
461 {{2.0, 2.0}, {2.0, 2.0}},
462 {{0.0, 0.0}, {0.0, 0.0}}});
463 ComputeAndCompareR3<float>(&b, expected, {}, DefaultErrorSpec());
464 }
465
466 INSTANTIATE_TEST_CASE_P(PadTestFloatInstantiation, PadTestFloat,
467 ::testing::ValuesIn(use_bfloat16_params));
468
469 } // namespace
470 } // namespace xla
471