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