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/array4d.h"
17
18#include <initializer_list>
19#include <numeric>
20
21#include "tensorflow/compiler/xla/test.h"
22#include "tensorflow/core/lib/gtl/array_slice.h"
23
24namespace xla {
25namespace {
26
27// Given an Array4D and a 4-tuple index, computes the linear index into the
28// array idx represents.
29template <typename T>
30int64 Array4DLinearIndex(const Array4D<T>& arr,
31                         tensorflow::gtl::ArraySlice<int64> idx) {
32  EXPECT_EQ(4, idx.size());
33  return (idx[3] + idx[2] * arr.n4() + idx[1] * arr.n3() * arr.n4() +
34          idx[0] * arr.n2() * arr.n3() * arr.n4());
35}
36
37TEST(Array4dTest, UninitializedDimsCtor) {
38  Array4D<int> empty(2, 3, 4, 5);
39  EXPECT_EQ(empty.n1(), 2);
40  EXPECT_EQ(empty.n2(), 3);
41  EXPECT_EQ(empty.n3(), 4);
42  EXPECT_EQ(empty.n4(), 5);
43  EXPECT_EQ(empty.num_elements(), 120);
44}
45
46TEST(Array4dTest, FillCtor) {
47  Array4D<int> fullof7(2, 3, 4, 5, 7);
48
49  EXPECT_EQ(fullof7.n1(), 2);
50  EXPECT_EQ(fullof7.n2(), 3);
51  EXPECT_EQ(fullof7.n3(), 4);
52  EXPECT_EQ(fullof7.n4(), 5);
53
54  fullof7.Each([](tensorflow::gtl::ArraySlice<int64> idx, int* cell) {
55    EXPECT_EQ(*cell, 7);
56  });
57}
58
59TEST(Array4dTest, ContainerCtor) {
60  // Fill an Array4D with a linear vector of [0..119] according to the default
61  // row-major ordering.
62  std::vector<int> filler(120);
63  std::iota(filler.begin(), filler.end(), 0);
64
65  Array4D<int> arr(2, 3, 4, 5, filler);
66
67  EXPECT_EQ(arr.n1(), 2);
68  EXPECT_EQ(arr.n2(), 3);
69  EXPECT_EQ(arr.n3(), 4);
70  EXPECT_EQ(arr.n4(), 5);
71
72  arr.Each([&arr](tensorflow::gtl::ArraySlice<int64> idx, int* cell) {
73    EXPECT_EQ(*cell, Array4DLinearIndex(arr, idx));
74  });
75}
76
77TEST(Array3dTest, InitializerListCtor) {
78  Array4D<int> arr = {{{{1}, {2}}, {{3}, {4}}, {{5}, {6}}, {{7}, {8}}},
79                      {{{9}, {10}}, {{11}, {12}}, {{13}, {14}}, {{15}, {16}}},
80                      {{{17}, {18}}, {{19}, {20}}, {{21}, {22}}, {{23}, {24}}}};
81
82  EXPECT_EQ(arr.n1(), 3);
83  EXPECT_EQ(arr.n2(), 4);
84  EXPECT_EQ(arr.n3(), 2);
85  EXPECT_EQ(arr.n4(), 1);
86  EXPECT_EQ(arr.num_elements(), 24);
87
88  EXPECT_EQ(arr(0, 0, 0, 0), 1);
89  EXPECT_EQ(arr(0, 0, 1, 0), 2);
90  EXPECT_EQ(arr(0, 1, 0, 0), 3);
91  EXPECT_EQ(arr(0, 3, 1, 0), 8);
92  EXPECT_EQ(arr(1, 0, 0, 0), 9);
93  EXPECT_EQ(arr(1, 1, 1, 0), 12);
94  EXPECT_EQ(arr(2, 0, 0, 0), 17);
95  EXPECT_EQ(arr(2, 1, 1, 0), 20);
96  EXPECT_EQ(arr(2, 2, 0, 0), 21);
97  EXPECT_EQ(arr(2, 3, 1, 0), 24);
98}
99
100TEST(Array3dTest, InitializerListCtorHalf) {
101  Array4D<Eigen::half> arr = {
102      {{{1.0f}, {2.0f}}, {{3.0f}, {4.0f}}, {{5.0f}, {6.0f}}, {{7.0f}, {8.0f}}},
103      {{{9.0f}, {10.0f}},
104       {{11.0f}, {12.0f}},
105       {{13.0f}, {14.0f}},
106       {{15.0f}, {16.0f}}},
107      {{{17.0f}, {18.0f}},
108       {{19.0f}, {20.0f}},
109       {{21.0f}, {22.0f}},
110       {{23.0f}, {24.0f}}}};
111
112  EXPECT_EQ(arr.n1(), 3);
113  EXPECT_EQ(arr.n2(), 4);
114  EXPECT_EQ(arr.n3(), 2);
115  EXPECT_EQ(arr.n4(), 1);
116  EXPECT_EQ(arr.num_elements(), 24);
117
118  EXPECT_EQ(arr(0, 0, 0, 0), static_cast<Eigen::half>(1));
119  EXPECT_EQ(arr(0, 0, 1, 0), static_cast<Eigen::half>(2));
120  EXPECT_EQ(arr(0, 1, 0, 0), static_cast<Eigen::half>(3));
121  EXPECT_EQ(arr(0, 3, 1, 0), static_cast<Eigen::half>(8));
122  EXPECT_EQ(arr(1, 0, 0, 0), static_cast<Eigen::half>(9));
123  EXPECT_EQ(arr(1, 1, 1, 0), static_cast<Eigen::half>(12));
124  EXPECT_EQ(arr(2, 0, 0, 0), static_cast<Eigen::half>(17));
125  EXPECT_EQ(arr(2, 1, 1, 0), static_cast<Eigen::half>(20));
126  EXPECT_EQ(arr(2, 2, 0, 0), static_cast<Eigen::half>(21));
127  EXPECT_EQ(arr(2, 3, 1, 0), static_cast<Eigen::half>(24));
128}
129
130TEST(Array4dTest, Fill) {
131  Array4D<int> fullof7(2, 3, 4, 5, 7);
132  fullof7.Each([](tensorflow::gtl::ArraySlice<int64> idx, int* cell) {
133    EXPECT_EQ(*cell, 7);
134  });
135
136  fullof7.Fill(11);
137  fullof7.Each([](tensorflow::gtl::ArraySlice<int64> idx, int* cell) {
138    EXPECT_EQ(*cell, 11);
139  });
140}
141
142TEST(Array4dTest, FillWithMultiples) {
143  Array4D<float> arr(2, 3, 4, 5);
144  arr.FillWithMultiples(2.0f);
145
146  arr.Each([&arr](tensorflow::gtl::ArraySlice<int64> idx, float* cell) {
147    EXPECT_EQ(*cell, 2.0f * Array4DLinearIndex(arr, idx));
148  });
149}
150
151TEST(Array4dTest, FillRasterDimensionDepthOne) {
152  Array4D<float> array(1, 1, 128, 128);
153  Array2D<float> raster(128, 128);
154  for (int row = 0; row < 128; ++row) {
155    for (int col = 0; col < 128; ++col) {
156      raster(row, col) = row * 1000.0 + col;
157    }
158  }
159
160  array.FillWithYX(raster);
161
162  VLOG(1) << array.ToString();
163
164  EXPECT_FLOAT_EQ(raster(0, 0), array(0, 0, 0, 0));
165  EXPECT_FLOAT_EQ(raster(0, 1), array(0, 0, 0, 1));
166  EXPECT_FLOAT_EQ(raster(1, 0), array(0, 0, 1, 0));
167  EXPECT_FLOAT_EQ(raster(1, 1), array(0, 0, 1, 1));
168  EXPECT_FLOAT_EQ(raster(2, 0), array(0, 0, 2, 0));
169  EXPECT_FLOAT_EQ(raster(127, 127), array(0, 0, 127, 127));
170
171  EXPECT_FLOAT_EQ(0, array(0, 0, 0, 0));
172  EXPECT_FLOAT_EQ(1, array(0, 0, 0, 1));
173  EXPECT_FLOAT_EQ(2, array(0, 0, 0, 2));
174
175  EXPECT_FLOAT_EQ(1001, array(0, 0, 1, 1));
176  EXPECT_FLOAT_EQ(2001, array(0, 0, 2, 1));
177  EXPECT_FLOAT_EQ(127000, array(0, 0, 127, 0));
178  EXPECT_FLOAT_EQ(127127, array(0, 0, 127, 127));
179}
180
181TEST(Array4dTest, FillWithPzTestDepthOne) {
182  Array2D<float> matrix(3, 2);
183  std::initializer_list<std::initializer_list<float>> values = {
184      {-3.f, -0.1f}, {0.f, -0.1f}, {3.f, 0.2f},
185  };
186  int rowno = 0;
187  for (auto row : values) {
188    int colno = 0;
189    for (float f : row) {
190      matrix(rowno, colno) = f;
191      colno++;
192    }
193    rowno++;
194  }
195
196  Array4D<float> actual(3, 2, 1, 1);
197  actual.FillWithPZ(matrix);
198
199  EXPECT_FLOAT_EQ(-3, actual(0, 0, 0, 0));
200  EXPECT_FLOAT_EQ(-0.1, actual(0, 1, 0, 0));
201
202  EXPECT_FLOAT_EQ(0, actual(1, 0, 0, 0));
203  EXPECT_FLOAT_EQ(-0.1, actual(1, 1, 0, 0));
204
205  EXPECT_FLOAT_EQ(3, actual(2, 0, 0, 0));
206  EXPECT_FLOAT_EQ(0.2, actual(2, 1, 0, 0));
207}
208
209}  // namespace
210}  // namespace xla
211