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