1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/compiler/xla/reference_util.h"
17
18#include <cmath>
19#include <memory>
20
21#include "tensorflow/compiler/xla/array2d.h"
22#include "tensorflow/compiler/xla/array3d.h"
23#include "tensorflow/compiler/xla/array4d.h"
24#include "tensorflow/compiler/xla/client/padding.h"
25#include "tensorflow/compiler/xla/literal_util.h"
26#include "tensorflow/compiler/xla/ptr_util.h"
27#include "tensorflow/compiler/xla/test.h"
28#include "tensorflow/compiler/xla/tests/literal_test_util.h"
29#include "tensorflow/compiler/xla/xla_data.pb.h"
30
31namespace xla {
32namespace {
33
34// Tests linear algebra routines implemented in ReferenceUtil class.
35// TODO(b/23829238): Currently missing tests for the convolution routine.
36class ReferenceUtilTest : public ::testing::Test {
37 protected:
38  ReferenceUtilTest() {
39    matrix_ = MakeUnique<Array2D<float>>(rows_, cols_);
40    // [1.f  2.f  3.f]
41    // [4.f  5.f  6.f]
42    for (int64 i = 0; i < rows_; ++i) {
43      for (int64 j = 0; j < cols_; ++j) {
44        (*matrix_)(i, j) = i * cols_ + j + 1;
45      }
46    }
47  }
48
49  const int64 rows_ = 2;
50  const int64 cols_ = 3;
51  std::unique_ptr<Array2D<float>> matrix_;
52};
53
54TEST_F(ReferenceUtilTest, TransposeArray2D) {
55  auto result = ReferenceUtil::TransposeArray2D(*matrix_);
56  auto actual_literal = Literal::CreateR2FromArray2D(*result);
57  LiteralTestUtil::ExpectR2Near<float>({{1.f, 4.f}, {2.f, 5.f}, {3.f, 6.f}},
58                                       *actual_literal, ErrorSpec(0.0001));
59}
60
61TEST_F(ReferenceUtilTest, MatmulArray2D) {
62  Array2D<float> rhs({
63      {7.f, 8.f},
64      {9.f, 10.f},
65      {11.f, 12.f},
66  });
67  auto result = ReferenceUtil::MatmulArray2D(*matrix_, rhs);
68  auto actual_literal = Literal::CreateR2FromArray2D(*result);
69  LiteralTestUtil::ExpectR2Near<float>({{58.f, 64.f}, {139.f, 154.f}},
70                                       *actual_literal, ErrorSpec(0.0001));
71}
72
73TEST_F(ReferenceUtilTest, ReduceToColArray2D) {
74  auto add = [](float lhs, float rhs) { return lhs + rhs; };
75  auto result = ReferenceUtil::ReduceToColArray2D(*matrix_, 0.0f, add);
76  auto actual_literal = Literal::CreateR1<float>(*result);
77  LiteralTestUtil::ExpectR1Near<float>({6.f, 15.f}, *actual_literal,
78                                       ErrorSpec(0.0001));
79}
80
81TEST_F(ReferenceUtilTest, ReduceToRowArray2D) {
82  auto add = [](float lhs, float rhs) { return lhs + rhs; };
83  auto result = ReferenceUtil::ReduceToRowArray2D(*matrix_, 0.0f, add);
84  auto actual_literal = Literal::CreateR1<float>(*result);
85  LiteralTestUtil::ExpectR1Near<float>({5.f, 7.f, 9.f}, *actual_literal,
86                                       ErrorSpec(0.0001));
87}
88
89TEST_F(ReferenceUtilTest, Reduce4Dto1DZeroSizedArray) {
90  auto result = Literal::CreateR1<float>(ReferenceUtil::Reduce4DTo1D(
91      Array4D<float>(1, 0, 1, 1), /*init=*/0, /*dims=*/{0, 1, 2},
92      [](float a, float b) { return a + b; }));
93  LiteralTestUtil::ExpectR1Equal<float>({0}, *result);
94}
95
96TEST_F(ReferenceUtilTest, MapArray2D) {
97  auto identity = [](float value) { return log(exp(value)); };
98  auto result = ReferenceUtil::MapArray2D(*matrix_, identity);
99  auto actual_literal = Literal::CreateR2FromArray2D(*result);
100  LiteralTestUtil::ExpectR2NearArray2D(*matrix_, *actual_literal,
101                                       ErrorSpec(0.0001));
102}
103
104TEST_F(ReferenceUtilTest, MapWithIndexArray2D) {
105  auto add_index = [](float value, int64 row, int64 col) {
106    return value + row + col;
107  };
108  auto result = ReferenceUtil::MapWithIndexArray2D(*matrix_, add_index);
109  auto actual_literal = Literal::CreateR2FromArray2D(*result);
110  LiteralTestUtil::ExpectR2Near<float>({{1.f, 3.f, 5.f}, {5.f, 7.f, 9.f}},
111                                       *actual_literal, ErrorSpec(0.0001));
112}
113
114TEST_F(ReferenceUtilTest, MapArray4D) {
115  auto input = MakeUnique<Array4D<float>>(/*planes=*/2, /*depth=*/3,
116                                          /*height=*/4, /*width=*/5);
117  input->FillWithMultiples(1.0f);
118  auto multiply_by_two = [](float value) { return 2 * value; };
119  auto result = ReferenceUtil::MapArray4D(*input, multiply_by_two);
120  auto actual_literal = Literal::CreateR4FromArray4D(*result);
121
122  Array4D<float> expected(/*planes=*/2, /*depth=*/3, /*height=*/4, /*width=*/5);
123  expected.FillWithMultiples(2.0f);
124  LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal,
125                                       ErrorSpec(0.0001));
126}
127
128TEST_F(ReferenceUtilTest, MapWithIndexArray4D) {
129  auto input = MakeUnique<Array4D<float>>(/*planes=*/2, /*depth=*/3,
130                                          /*height=*/4, /*width=*/5);
131  input->FillWithMultiples(1.0f);
132  auto subtract_index = [](float value, int64 plane, int64 depth, int64 height,
133                           int64 width) {
134    return value - (3 * 4 * 5 * plane + 4 * 5 * depth + 5 * height + width);
135  };
136  auto result = ReferenceUtil::MapWithIndexArray4D(*input, subtract_index);
137  auto actual_literal = Literal::CreateR4FromArray4D(*result);
138
139  Array4D<float> expected(/*planes=*/2, /*depth=*/3, /*height=*/4, /*width=*/5);
140  expected.Fill(0.0f);
141  LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal,
142                                       ErrorSpec(0.0001));
143}
144
145TEST_F(ReferenceUtilTest, SliceArray2D) {
146  auto result = ReferenceUtil::Slice2D(*matrix_, {{0, 0}}, {{2, 2}}, {{1, 1}});
147  auto actual_literal = Literal::CreateR2FromArray2D(*result);
148
149  LiteralTestUtil::ExpectR2Near<float>({{1.f, 2.f}, {4.f, 5.f}},
150                                       *actual_literal, ErrorSpec(0.0001));
151}
152
153TEST_F(ReferenceUtilTest, SliceStridedArray2D) {
154  auto result = ReferenceUtil::Slice2D(*matrix_, {{0, 0}}, {{2, 3}}, {{1, 2}});
155  auto actual_literal = Literal::CreateR2FromArray2D(*result);
156
157  LiteralTestUtil::ExpectR2Near<float>({{1.f, 3.f}, {4.f, 6.f}},
158                                       *actual_literal, ErrorSpec(0.0001));
159}
160
161TEST_F(ReferenceUtilTest, SliceArray3D) {
162  Array3D<float> input(2, 3, 4);
163  input.FillIota(0);
164
165  auto result =
166      ReferenceUtil::Slice3D(input, {{0, 0, 0}}, {{2, 2, 2}}, {{1, 1, 1}});
167  auto actual_literal = Literal::CreateR3FromArray3D(*result);
168
169  LiteralTestUtil::ExpectR3Near<float>(
170      {{{0.f, 1.f}, {4.f, 5.f}}, {{12.f, 13.f}, {16.f, 17.f}}}, *actual_literal,
171      ErrorSpec(0.0001));
172}
173
174TEST_F(ReferenceUtilTest, SliceStridedArray3D) {
175  Array3D<float> input(2, 3, 4);
176  input.FillIota(0);
177
178  auto result =
179      ReferenceUtil::Slice3D(input, {{0, 0, 0}}, {{2, 3, 4}}, {{1, 2, 2}});
180  auto actual_literal = Literal::CreateR3FromArray3D(*result);
181
182  LiteralTestUtil::ExpectR3Near<float>(
183      {{{0.f, 2.f}, {8.f, 10.f}}, {{12.f, 14.f}, {20.f, 22.f}}},
184      *actual_literal, ErrorSpec(0.0001));
185}
186
187TEST_F(ReferenceUtilTest, SliceArray4D) {
188  Array4D<float> input(2, 3, 4, 5);
189  input.FillIota(0);
190
191  auto result = ReferenceUtil::Slice4D(input, {{1, 0, 0, 0}}, {{2, 2, 2, 2}},
192                                       {{1, 1, 1, 1}});
193  auto actual_literal = Literal::CreateR4FromArray4D(*result);
194
195  LiteralTestUtil::ExpectR4Near<float>(
196      {{{{60.f, 61.f}, {65.f, 66.f}}, {{80.f, 81.f}, {85.f, 86.f}}}},
197      *actual_literal, ErrorSpec(0.0001));
198}
199
200TEST_F(ReferenceUtilTest, SliceStridedArray4D) {
201  Array4D<float> input(2, 3, 4, 5);
202  input.FillIota(0);
203
204  auto result = ReferenceUtil::Slice4D(input, {{1, 0, 0, 0}}, {{2, 3, 4, 5}},
205                                       {{1, 2, 2, 2}});
206  auto actual_literal = Literal::CreateR4FromArray4D(*result);
207
208  LiteralTestUtil::ExpectR4Near<float>(
209      {{{{60.f, 62.f, 64.f}, {70.f, 72.f, 74.f}},
210        {{100.f, 102.f, 104.f}, {110.f, 112.f, 114.f}}}},
211      *actual_literal, ErrorSpec(0.0001));
212}
213
214TEST_F(ReferenceUtilTest, ConvArray3DWithSamePadding) {
215  Array3D<float> input = {{{1, 2, 3, 4}}};
216  Array3D<float> weights = {{{5, 6}}};
217  std::unique_ptr<Array3D<float>> actual =
218      ReferenceUtil::ConvArray3D(input, weights, 1, Padding::kSame);
219  Array3D<float> expected = {{{17, 28, 39, 20}}};
220
221  auto actual_literal = Literal::CreateR3FromArray3D(*actual);
222
223  LiteralTestUtil::ExpectR3NearArray3D<float>(expected, *actual_literal,
224                                              ErrorSpec(0.0001));
225}
226
227TEST_F(ReferenceUtilTest, ConvArray3DWithValidPadding) {
228  Array3D<float> input = {{{1, 2, 3, 4}}};
229  Array3D<float> weights = {{{5, 6}}};
230  std::unique_ptr<Array3D<float>> actual =
231      ReferenceUtil::ConvArray3D(input, weights, 1, Padding::kValid);
232  Array3D<float> expected = {{{17, 28, 39}}};
233
234  auto actual_literal = Literal::CreateR3FromArray3D(*actual);
235
236  LiteralTestUtil::ExpectR3NearArray3D<float>(expected, *actual_literal,
237                                              ErrorSpec(0.0001));
238}
239
240TEST_F(ReferenceUtilTest, ConvWithSamePadding) {
241  Array4D<float> input(1, 1, 4, 4);
242  // clang-format off
243  input.FillWithYX(Array2D<float>({
244    {1,  2,  3,  4 },
245    {5,  6,  7,  8 },
246    {9,  10, 11, 12},
247    {13, 14, 15, 16},
248  }));
249  // clang-format on
250  Array4D<float> weights(1, 1, 2, 2);
251  // clang-format off
252  weights.FillWithYX(Array2D<float>({
253    {5, 6},
254    {7, 8},
255  }));
256  // clang-format on
257  std::unique_ptr<Array4D<float>> actual =
258      ReferenceUtil::ConvArray4D(input, weights, {1, 1}, Padding::kSame);
259  Array4D<float> expected(1, 1, 4, 4);
260  // clang-format off
261  expected.FillWithYX(Array2D<float>({
262    {100, 126, 152,  76},
263    {204, 230, 256, 124},
264    {308, 334, 360, 172},
265    {149, 160, 171,  80},
266  }));
267  // clang-format on
268
269  auto actual_literal = Literal::CreateR4FromArray4D(*actual);
270
271  LiteralTestUtil::ExpectR4NearArray4D<float>(expected, *actual_literal,
272                                              ErrorSpec(0.0001));
273}
274
275TEST_F(ReferenceUtilTest, ConvWithValidPadding) {
276  Array4D<float> input(1, 1, 4, 4);
277  // clang-format off
278  input.FillWithYX(Array2D<float>({
279    {1,  2,  3,  4 },
280    {5,  6,  7,  8 },
281    {9,  10, 11, 12},
282    {13, 14, 15, 16},
283  }));
284  // clang-format on
285  Array4D<float> weights(1, 1, 2, 2);
286  // clang-format off
287  weights.FillWithYX(Array2D<float>({
288    {5, 6},
289    {7, 8},
290  }));
291  // clang-format on
292  std::unique_ptr<Array4D<float>> actual =
293      ReferenceUtil::ConvArray4D(input, weights, {1, 1}, Padding::kValid);
294  Array4D<float> expected(1, 1, 3, 3);
295  // clang-format off
296  expected.FillWithYX(Array2D<float>({
297    {1*5+2*6+5*7+6*8, 126, 152},
298    {204, 230, 256},
299    {308, 334, 11*5+12*6+15*7+16*8},
300  }));
301  // clang-format on
302
303  auto actual_literal = Literal::CreateR4FromArray4D(*actual);
304
305  LiteralTestUtil::ExpectR4NearArray4D<float>(expected, *actual_literal,
306                                              ErrorSpec(0.0001));
307}
308
309TEST_F(ReferenceUtilTest, ConvGeneralDimensionsWithSamePadding) {
310  // clang-format off
311  // Input dimensions: [feature=2, height=3, batch=1, width=4]
312  Array4D<float> input({
313    {{{1, 2, 3, 4}},
314     {{5, 6, 7, 8}},
315     {{9, 10, 11, 12}}},
316    {{{13, 14, 15, 16}},
317     {{17, 18, 19, 20}},
318     {{21, 22, 23, 24}}}
319  });
320  // Weight dimensions:
321  // [kernel_output_feature=1, height=3, kernel_input_feature=2, width=3]
322  Array4D<float> weight({{
323    {{1, 2, 3},
324     {4, 5, 6}},
325    {{7, 8, 9},
326     {10, 11, 12}},
327    {{13, 14, 15},
328     {16, 17, 18}}
329  }});
330  // clang-format on
331
332  // Set the convolution dimension numbers.
333  ConvolutionDimensionNumbers dimension_numbers;
334  dimension_numbers.set_input_batch_dimension(2);
335  dimension_numbers.set_input_feature_dimension(0);
336  dimension_numbers.set_output_batch_dimension(2);
337  dimension_numbers.set_output_feature_dimension(0);
338  dimension_numbers.add_input_spatial_dimensions(1);
339  dimension_numbers.add_output_spatial_dimensions(1);
340  dimension_numbers.add_input_spatial_dimensions(3);
341  dimension_numbers.add_output_spatial_dimensions(3);
342  dimension_numbers.set_kernel_output_feature_dimension(0);
343  dimension_numbers.set_kernel_input_feature_dimension(2);
344  dimension_numbers.add_kernel_spatial_dimensions(1);
345  dimension_numbers.add_kernel_spatial_dimensions(3);
346
347  std::unique_ptr<Array4D<float>> actual =
348      ReferenceUtil::ConvArray4DGeneralDimensions(
349          input, weight, {1, 1}, Padding::kSame, dimension_numbers);
350  // clang-format off
351  // Result dimensions: [feature=1, height=3, batch=1, width=4]
352  Array4D<float> expected({{
353    {{1110, 1688, 1838, 1226}},
354    {{1683, 2514, 2685, 1761}},
355    {{878, 1280, 1358, 866}}
356  }});
357  // clang-format on
358
359  auto actual_literal = Literal::CreateR4FromArray4D(*actual);
360
361  LiteralTestUtil::ExpectR4NearArray4D<float>(expected, *actual_literal,
362                                              ErrorSpec(0.0001));
363}
364
365TEST_F(ReferenceUtilTest, ConvGeneralDimensionsWithValidPadding) {
366  // clang-format off
367  // Input dimensions: [feature=2, height=3, batch=1, width=4]
368  Array4D<float> input({
369    {{{1, 2, 3, 4}},
370     {{5, 6, 7, 8}},
371     {{9, 10, 11, 12}}},
372    {{{13, 14, 15, 16}},
373     {{17, 18, 19, 20}},
374     {{21, 22, 23, 24}}}
375  });
376  // Weight dimensions:
377  // [kernel_output_feature=1, width=3, kernel_input_feature=2, height=3]
378  Array4D<float> weight({{
379    {{1, 7, 13},
380     {4, 10, 16}},
381    {{2, 8, 14},
382     {5, 11, 17}},
383    {{3, 9, 15},
384     {6, 12, 18}}
385  }});
386  // clang-format on
387
388  // Set the convolution dimension numbers.
389  ConvolutionDimensionNumbers dimension_numbers;
390  dimension_numbers.set_input_batch_dimension(2);
391  dimension_numbers.set_input_feature_dimension(0);
392  dimension_numbers.set_output_batch_dimension(2);
393  dimension_numbers.set_output_feature_dimension(0);
394  dimension_numbers.add_input_spatial_dimensions(1);
395  dimension_numbers.add_output_spatial_dimensions(1);
396  dimension_numbers.add_input_spatial_dimensions(3);
397  dimension_numbers.add_output_spatial_dimensions(3);
398
399  dimension_numbers.set_kernel_output_feature_dimension(0);
400  dimension_numbers.set_kernel_input_feature_dimension(2);
401  dimension_numbers.add_kernel_spatial_dimensions(3);
402  dimension_numbers.add_kernel_spatial_dimensions(1);
403
404  std::unique_ptr<Array4D<float>> actual =
405      ReferenceUtil::ConvArray4DGeneralDimensions(
406          input, weight, {1, 1}, Padding::kValid, dimension_numbers);
407  // clang-format off
408  // Result dimensions: [feature=1, height=1, batch=1, width=2]
409  Array4D<float> expected({{{{2514, 2685}}}});
410  // clang-format on
411
412  auto actual_literal = Literal::CreateR4FromArray4D(*actual);
413
414  LiteralTestUtil::ExpectR4NearArray4D<float>(expected, *actual_literal,
415                                              ErrorSpec(0.0001));
416}
417
418TEST_F(ReferenceUtilTest, ApplyElementwise2D) {
419  Array2D<float> a({{1, 2}, {3, 4}});
420  Array2D<float> b({{10, 20}, {30, 40}});
421  Array2D<float> c({{100, 200}, {300, 400}});
422
423  auto actual = ReferenceUtil::ApplyElementwise2D(
424      [](float x, float y, float z) { return 100 * x + 10 * y + z; }, a, b, c);
425  auto actual_literal = Literal::CreateR2FromArray2D(*actual);
426  LiteralTestUtil::ExpectR2Near({{300.f, 600.f}, {900.f, 1200.f}},
427                                *actual_literal, ErrorSpec(0.0001));
428}
429
430}  // namespace
431}  // namespace xla
432