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/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 31 namespace xla { 32 namespace { 33 34 // Tests linear algebra routines implemented in ReferenceUtil class. 35 // TODO(b/23829238): Currently missing tests for the convolution routine. 36 class 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 54 TEST_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 61 TEST_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 73 TEST_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 81 TEST_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 89 TEST_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 96 TEST_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 104 TEST_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 114 TEST_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 128 TEST_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 145 TEST_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 153 TEST_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 161 TEST_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 174 TEST_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 187 TEST_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 200 TEST_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 214 TEST_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 227 TEST_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 240 TEST_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 275 TEST_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 309 TEST_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 365 TEST_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 418 TEST_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
