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