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> 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 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 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 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 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 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 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 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 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 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
