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 // Tests of convolution variants -- kernel sizes, padding, and strides -- 17 // in small sized data. 18 19 #include <algorithm> 20 #include <initializer_list> 21 #include <memory> 22 #include <numeric> 23 #include <random> 24 #include <vector> 25 26 #include "tensorflow/compiler/xla/array3d.h" 27 #include "tensorflow/compiler/xla/array4d.h" 28 #include "tensorflow/compiler/xla/client/computation_builder.h" 29 #include "tensorflow/compiler/xla/client/local_client.h" 30 #include "tensorflow/compiler/xla/client/padding.h" 31 #include "tensorflow/compiler/xla/literal_util.h" 32 #include "tensorflow/compiler/xla/reference_util.h" 33 #include "tensorflow/compiler/xla/tests/client_library_test_base.h" 34 #include "tensorflow/compiler/xla/tests/literal_test_util.h" 35 #include "tensorflow/compiler/xla/tests/test_macros.h" 36 #include "tensorflow/compiler/xla/xla_data.pb.h" 37 #include "tensorflow/core/platform/test.h" 38 #include "tensorflow/core/platform/types.h" 39 40 namespace xla { 41 namespace { 42 43 class ConvolutionVariantsTest : public ClientLibraryTestBase { 44 protected: 45 #if XLA_TEST_BACKEND_GPU 46 // XLA:GPU sometimes uses FFT convolution which isn't as precise as spatial 47 // convolution. So relax the absolute error threshold. 48 ErrorSpec error_spec_ = ErrorSpec(1e-1, 1e-5); 49 #else 50 ErrorSpec error_spec_ = ErrorSpec(1e-4, 1e-2); 51 #endif 52 }; 53 54 XLA_TEST_F(ConvolutionVariantsTest, Minimal) { 55 ComputationBuilder builder(client_, TestName()); 56 57 const Array4D<float> input_array(1, 1, 1, 1, {2}); 58 auto input = builder.ConstantR4FromArray4D<float>(input_array); 59 60 const Array4D<float> filter_array(1, 1, 1, 1, {3}); 61 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 62 63 builder.Conv(input, filter, {1, 1}, Padding::kValid); 64 65 const Array4D<float> expected(1, 1, 1, 1, {6}); 66 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 67 } 68 69 XLA_TEST_F(ConvolutionVariantsTest, MinimalWithBatch) { 70 ComputationBuilder builder(client_, TestName()); 71 72 const Array4D<float> input_array(5, 1, 1, 1, {1, 2, 3, 4, 5}); 73 auto input = builder.ConstantR4FromArray4D<float>(input_array); 74 75 const Array4D<float> filter_array(1, 1, 1, 1, {2}); 76 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 77 78 builder.Conv(input, filter, {1, 1}, Padding::kValid); 79 80 const Array4D<float> expected(5, 1, 1, 1, {2, 4, 6, 8, 10}); 81 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 82 } 83 84 XLA_TEST_F(ConvolutionVariantsTest, Flat1x1) { 85 ComputationBuilder builder(client_, TestName()); 86 87 Array4D<float> input_array(2, 1, 3, 4); 88 input_array.FillWithMultiples(1); 89 auto input = builder.ConstantR4FromArray4D<float>(input_array); 90 91 const Array4D<float> filter_array(1, 1, 1, 1, {2.3}); 92 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 93 94 builder.Conv(input, filter, {1, 1}, Padding::kValid); 95 96 Array4D<float> expected(2, 1, 3, 4); 97 expected.FillWithMultiples(2.3); 98 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 99 } 100 101 XLA_TEST_F(ConvolutionVariantsTest, Deep1x1) { 102 ComputationBuilder builder(client_, TestName()); 103 104 Array4D<float> input_array(1, 2, 1, 1, {10, 1}); 105 auto input = builder.ConstantR4FromArray4D<float>(input_array); 106 107 const Array4D<float> filter_array(3, 2, 1, 1, {1, 2, 3, 4, 5, 6}); 108 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 109 110 builder.Conv(input, filter, {1, 1}, Padding::kValid); 111 112 Array4D<float> expected(1, 3, 1, 1, {12, 34, 56}); 113 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 114 } 115 116 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in1x2) { 117 ComputationBuilder builder(client_, TestName()); 118 119 Array4D<float> input_array(1, 1, 1, 2, {1, 2}); 120 auto input = builder.ConstantR4FromArray4D<float>(input_array); 121 122 const Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 123 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 124 125 builder.Conv(input, filter, {1, 1}, Padding::kValid); 126 127 Array4D<float> expected(1, 1, 1, 1, {12}); 128 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 129 } 130 131 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in1x3) { 132 ComputationBuilder builder(client_, TestName()); 133 134 Array4D<float> input_array(1, 1, 1, 3, {1, 2, 3}); 135 auto input = builder.ConstantR4FromArray4D<float>(input_array); 136 137 const Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 138 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 139 140 builder.Conv(input, filter, {1, 1}, Padding::kValid); 141 142 Array4D<float> expected(1, 1, 1, 2, {12, 23}); 143 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 144 } 145 146 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in2x2) { 147 ComputationBuilder builder(client_, TestName()); 148 149 Array4D<float> input_array(1, 1, 2, 2, {1, 2, 3, 4}); 150 auto input = builder.ConstantR4FromArray4D<float>(input_array); 151 152 const Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 153 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 154 155 builder.Conv(input, filter, {1, 1}, Padding::kValid); 156 157 Array4D<float> expected(1, 1, 2, 1, {12, 34}); 158 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 159 } 160 161 XLA_TEST_F(ConvolutionVariantsTest, Filter2x1in2x2) { 162 ComputationBuilder builder(client_, TestName()); 163 164 Array4D<float> input_array(1, 1, 2, 2, {1, 2, 3, 4}); 165 auto input = builder.ConstantR4FromArray4D<float>(input_array); 166 167 const Array4D<float> filter_array(1, 1, 2, 1, {10, 1}); 168 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 169 170 builder.Conv(input, filter, {1, 1}, Padding::kValid); 171 172 Array4D<float> expected(1, 1, 1, 2, {13, 24}); 173 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 174 } 175 176 XLA_TEST_F(ConvolutionVariantsTest, Filter2x2in2x2) { 177 ComputationBuilder builder(client_, TestName()); 178 179 Array4D<float> input_array(1, 1, 2, 2, {1, 2, 3, 4}); 180 auto input = builder.ConstantR4FromArray4D<float>(input_array); 181 182 const Array4D<float> filter_array(1, 1, 2, 2, {1000, 100, 10, 1}); 183 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 184 185 builder.Conv(input, filter, {1, 1}, Padding::kValid); 186 187 Array4D<float> expected(1, 1, 1, 1, {1234}); 188 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 189 } 190 191 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in2x3WithDepthAndBatch) { 192 ComputationBuilder builder(client_, TestName()); 193 194 Array4D<float> input_array( 195 2, 2, 2, 3, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, // plane 0 196 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 0, 0}); // plane 1 197 auto input = builder.ConstantR4FromArray4D<float>(input_array); 198 199 const Array4D<float> filter_array( 200 2, 2, 1, 2, {1000, 100, 10, 1, 0.1, 0.01, 0.001, 0.0001}); 201 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 202 203 builder.Conv(input, filter, {1, 1}, Padding::kValid); 204 205 Array4D<float> expected( 206 2, 2, 2, 2, 207 {167, 1278, 3490, 4500, 0.0167, 0.1278, 0.3490, 0.4500, // plane 0 208 334, 2556, 6980, 9000, 0.0334, 0.2556, 0.6980, 0.9000}); // plane 1 209 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 210 } 211 212 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x4) { 213 ComputationBuilder builder(client_, TestName()); 214 215 Array4D<float> input_array(1, 1, 1, 4, {1, 2, 3, 4}); 216 auto input = builder.ConstantR4FromArray4D<float>(input_array); 217 218 const Array4D<float> filter_array(1, 1, 1, 1, {10}); 219 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 220 221 builder.Conv(input, filter, {1, 2}, Padding::kValid); 222 223 Array4D<float> expected(1, 1, 1, 2, {10, 30}); 224 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 225 } 226 227 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x5) { 228 ComputationBuilder builder(client_, TestName()); 229 230 Array4D<float> input_array(1, 1, 1, 5, {1, 2, 3, 4, 5}); 231 auto input = builder.ConstantR4FromArray4D<float>(input_array); 232 233 const Array4D<float> filter_array(1, 1, 1, 1, {10}); 234 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 235 236 builder.Conv(input, filter, {1, 2}, Padding::kValid); 237 238 Array4D<float> expected(1, 1, 1, 3, {10, 30, 50}); 239 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 240 } 241 242 XLA_TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x4) { 243 ComputationBuilder builder(client_, TestName()); 244 245 Array4D<float> input_array(1, 1, 1, 4, {1, 2, 3, 4}); 246 auto input = builder.ConstantR4FromArray4D<float>(input_array); 247 248 const Array4D<float> filter_array(1, 1, 1, 3, {100, 10, 1}); 249 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 250 251 builder.Conv(input, filter, {1, 2}, Padding::kValid); 252 253 Array4D<float> expected(1, 1, 1, 1, {123}); 254 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 255 } 256 257 XLA_TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x5) { 258 ComputationBuilder builder(client_, TestName()); 259 260 Array4D<float> input_array(1, 1, 1, 5, {1, 2, 3, 4, 5}); 261 auto input = builder.ConstantR4FromArray4D<float>(input_array); 262 263 const Array4D<float> filter_array(1, 1, 1, 3, {100, 10, 1}); 264 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 265 266 builder.Conv(input, filter, {1, 2}, Padding::kValid); 267 268 Array4D<float> expected(1, 1, 1, 2, {123, 345}); 269 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 270 } 271 272 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride2x2in3x3) { 273 ComputationBuilder builder(client_, TestName()); 274 275 Array4D<float> input_array(1, 1, 3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9}); 276 auto input = builder.ConstantR4FromArray4D<float>(input_array); 277 278 const Array4D<float> filter_array(1, 1, 1, 1, {10}); 279 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 280 281 builder.Conv(input, filter, {2, 2}, Padding::kValid); 282 283 Array4D<float> expected(1, 1, 2, 2, {10, 30, 70, 90}); 284 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 285 } 286 287 XLA_TEST_F(ConvolutionVariantsTest, Filter3x1in1x1Padded) { 288 ComputationBuilder builder(client_, TestName()); 289 290 Array4D<float> input_array(1, 1, 1, 1, {1}); 291 auto input = builder.ConstantR4FromArray4D<float>(input_array); 292 293 const Array4D<float> filter_array(1, 1, 1, 3, {10, 20, 30}); 294 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 295 296 builder.Conv(input, filter, {1, 1}, Padding::kSame); 297 298 Array4D<float> expected(1, 1, 1, 1, {20}); 299 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 300 } 301 302 XLA_TEST_F(ConvolutionVariantsTest, Filter5x1in3x1Padded) { 303 ComputationBuilder builder(client_, TestName()); 304 305 Array4D<float> input_array(1, 1, 1, 3, {1, 2, 3}); 306 auto input = builder.ConstantR4FromArray4D<float>(input_array); 307 308 const Array4D<float> filter_array(1, 1, 1, 5, {10000, 1000, 100, 10, 1}); 309 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 310 311 builder.Conv(input, filter, {1, 1}, Padding::kSame); 312 313 Array4D<float> expected(1, 1, 1, 3, {123, 1230, 12300}); 314 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 315 } 316 317 XLA_TEST_F(ConvolutionVariantsTest, Filter3x3in2x2Padded) { 318 ComputationBuilder builder(client_, TestName()); 319 320 Array4D<float> input_array(1, 1, 2, 2, {1, 2, 3, 4}); 321 auto input = builder.ConstantR4FromArray4D<float>(input_array); 322 323 const Array4D<float> filter_array(1, 1, 3, 3, 324 {10000, 0, 1000, // row 0 325 0, 100, 0, // row 1 326 10, 0, 1}); // row 2 327 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 328 329 builder.Conv(input, filter, {1, 1}, Padding::kSame); 330 331 Array4D<float> expected(1, 1, 2, 2, {104, 230, 2300, 10400}); 332 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 333 } 334 335 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1in2x1WithPaddingAndDepth) { 336 ComputationBuilder builder(client_, TestName()); 337 338 Array4D<float> input_array(1, 2, 1, 2, {1, 2, 3, 4}); 339 auto input = builder.ConstantR4FromArray4D<float>(input_array); 340 341 const Array4D<float> filter_array(1, 2, 1, 1, {10, 1}); 342 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 343 344 builder.Conv(input, filter, {1, 1}, Padding::kSame); 345 346 Array4D<float> expected(1, 1, 1, 2, {13, 24}); 347 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 348 } 349 350 XLA_TEST_F(ConvolutionVariantsTest, Filter2x2Stride1x1Input3x3) { 351 ComputationBuilder builder(client_, TestName()); 352 353 Array4D<float> input_array(1, 1, 3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9}); 354 auto input = builder.ConstantR4FromArray4D<float>(input_array); 355 356 const Array4D<float> filter_array(1, 1, 2, 2, {7, 13, 17, 23}); 357 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 358 359 builder.Conv(input, filter, {1, 1}, Padding::kValid); 360 361 Array4D<float> expected(1, 1, 2, 2, {216, 276, 396, 456}); 362 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 363 } 364 365 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2Stride1x1Input1x3) { 366 ComputationBuilder builder(client_, TestName()); 367 368 Array4D<float> input_array(1, 1, 1, 3, {1, 2, 3}); 369 auto input = builder.ConstantR4FromArray4D<float>(input_array); 370 371 const Array4D<float> filter_array(1, 1, 1, 2, {7, 13}); 372 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 373 374 builder.Conv(input, filter, {1, 1}, Padding::kValid); 375 376 Array4D<float> expected(1, 1, 1, 2, {33, 53}); 377 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 378 } 379 380 XLA_TEST_F(ConvolutionVariantsTest, Filter2x1x8x8Input1x1x8x8) { 381 ComputationBuilder builder(client_, TestName()); 382 383 std::vector<float> input_data(64); 384 std::iota(input_data.begin(), input_data.end(), 0.0); 385 Array4D<float> input_array(1, 1, 8, 8, input_data); 386 auto input = builder.ConstantR4FromArray4D<float>(input_array); 387 388 std::vector<float> filter_data(128); 389 std::fill(filter_data.begin(), filter_data.begin() + 64, 1.0); 390 std::fill(filter_data.begin() + 64, filter_data.begin() + 128, 2.0); 391 const Array4D<float> filter_array(2, 1, 8, 8, filter_data); 392 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 393 394 builder.Conv(input, filter, {1, 1}, Padding::kValid); 395 396 Array4D<float> expected(1, 2, 1, 1, {2016, 4032}); 397 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 398 } 399 400 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input16x1x1x1) { 401 ComputationBuilder builder(client_, TestName()); 402 403 std::vector<float> input_data(16 * 1 * 1 * 1); 404 std::iota(input_data.begin(), input_data.end(), 1.0); 405 Array4D<float> input_array(16, 1, 1, 1, input_data); 406 auto input = builder.ConstantR4FromArray4D<float>(input_array); 407 408 std::vector<float> filter_data(1 * 1 * 1 * 1); 409 std::iota(filter_data.begin(), filter_data.end(), 1.0); 410 const Array4D<float> filter_array(1, 1, 1, 1, filter_data); 411 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 412 413 builder.Conv(input, filter, {1, 1}, Padding::kValid); 414 415 std::vector<float> expected_data = {1, 2, 3, 4, 5, 6, 7, 8, 416 9, 10, 11, 12, 13, 14, 15, 16}; 417 Array4D<float> expected(16, 1, 1, 1, expected_data); 418 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 419 } 420 421 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input16x1x2x2) { 422 ComputationBuilder builder(client_, TestName()); 423 424 constexpr int bs = 16; 425 constexpr int kx = 2; 426 constexpr int ky = 2; 427 Array4D<float> input_array(bs, 1, ky, kx); 428 for (int i0 = 0; i0 < bs; ++i0) { 429 for (int i2 = 0; i2 < ky; ++i2) { 430 for (int i3 = 0; i3 < kx; ++i3) { 431 input_array(i0, 0, i2, i3) = i0 + 1; 432 } 433 } 434 } 435 auto input = builder.ConstantR4FromArray4D<float>(input_array); 436 437 std::vector<float> filter_data(1 * 1 * ky * kx); 438 std::iota(filter_data.begin(), filter_data.end(), 1.0); 439 const Array4D<float> filter_array(1, 1, ky, kx, filter_data); 440 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 441 442 builder.Conv(input, filter, {1, 1}, Padding::kValid); 443 444 std::vector<float> expected_data(bs); 445 for (int i = 0; i < bs; ++i) { 446 expected_data[i] = 10 * (i + 1); 447 } 448 Array4D<float> expected(bs, 1, 1, 1, expected_data); 449 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 450 } 451 452 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input3x1x2x2) { 453 ComputationBuilder builder(client_, TestName()); 454 455 constexpr int kx = 2; 456 constexpr int ky = 2; 457 constexpr int bs = 3; 458 Array4D<float> input_array(bs, 1, ky, kx); 459 for (int i0 = 0; i0 < bs; ++i0) { 460 for (int i2 = 0; i2 < ky; ++i2) { 461 for (int i3 = 0; i3 < kx; ++i3) { 462 input_array(i0, 0, i2, i3) = i0 + i2 + i3 + 1; 463 } 464 } 465 } 466 auto input = builder.ConstantR4FromArray4D<float>(input_array); 467 468 std::vector<float> filter_data(1 * 1 * ky * kx); 469 std::iota(filter_data.begin(), filter_data.end(), 1.0); 470 const Array4D<float> filter_array(1, 1, ky, kx, filter_data); 471 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 472 473 builder.Conv(input, filter, {1, 1}, Padding::kValid); 474 475 std::vector<float> expected_data = { 476 23, 477 33, 478 43, 479 }; 480 Array4D<float> expected(bs, 1, 1, 1, expected_data); 481 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 482 } 483 484 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x8x8Input16x1x8x8) { 485 ComputationBuilder builder(client_, TestName()); 486 487 Array4D<float> input_array(16, 1, 8, 8); 488 for (int i0 = 0; i0 < 16; ++i0) { 489 for (int i2 = 0; i2 < 8; ++i2) { 490 for (int i3 = 0; i3 < 8; ++i3) { 491 input_array(i0, 0, i2, i3) = i0 + i2 + i3 + 1; 492 } 493 } 494 } 495 auto input = builder.ConstantR4FromArray4D<float>(input_array); 496 497 std::vector<float> filter_data(1 * 1 * 8 * 8); 498 std::iota(filter_data.begin(), filter_data.end(), 1.0); 499 const Array4D<float> filter_array(1, 1, 8, 8, filter_data); 500 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 501 502 builder.Conv(input, filter, {1, 1}, Padding::kValid); 503 504 std::vector<float> expected_data = { 505 19664, 21744, 23824, 25904, 27984, 30064, 32144, 34224, 506 36304, 38384, 40464, 42544, 44624, 46704, 48784, 50864, 507 }; 508 Array4D<float> expected(16, 1, 1, 1, expected_data); 509 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 510 } 511 512 XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input1x2x8x8) { 513 ComputationBuilder builder(client_, TestName()); 514 515 std::vector<float> input_data(2 * 8 * 8); 516 std::iota(input_data.begin(), input_data.end(), 0.0); 517 Array4D<float> input_array(1, 2, 8, 8, input_data); 518 auto input = builder.ConstantR4FromArray4D<float>(input_array); 519 520 std::vector<float> filter_data(2 * 2 * 8 * 8); 521 std::fill(filter_data.begin(), filter_data.begin() + filter_data.size() / 4, 522 1.0); 523 std::fill(filter_data.begin() + filter_data.size() / 4, 524 filter_data.begin() + filter_data.size() / 2, 2.0); 525 std::fill(filter_data.begin() + filter_data.size() / 2, 526 filter_data.begin() + 3 * filter_data.size() / 4, 3.0); 527 std::fill(filter_data.begin() + 3 * filter_data.size() / 4, filter_data.end(), 528 4.0); 529 const Array4D<float> filter_array(2, 2, 8, 8, filter_data); 530 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 531 532 builder.Conv(input, filter, {1, 1}, Padding::kValid); 533 534 Array4D<float> expected(1, 2, 1, 1, {14240, 30496}); 535 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 536 } 537 538 XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input2x2x8x8) { 539 ComputationBuilder builder(client_, TestName()); 540 541 std::vector<float> input_data(2 * 2 * 8 * 8); 542 std::iota(input_data.begin(), input_data.end(), 0.0); 543 Array4D<float> input_array(2, 2, 8, 8, input_data); 544 auto input = builder.ConstantR4FromArray4D<float>(input_array); 545 546 std::vector<float> filter_data(2 * 2 * 8 * 8); 547 std::fill(filter_data.begin(), filter_data.begin() + filter_data.size() / 4, 548 1.0); 549 std::fill(filter_data.begin() + filter_data.size() / 4, 550 filter_data.begin() + filter_data.size() / 2, 2.0); 551 std::fill(filter_data.begin() + filter_data.size() / 2, 552 filter_data.begin() + 3 * filter_data.size() / 4, 3.0); 553 std::fill(filter_data.begin() + 3 * filter_data.size() / 4, filter_data.end(), 554 4.0); 555 const Array4D<float> filter_array(2, 2, 8, 8, filter_data); 556 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 557 558 builder.Conv(input, filter, {1, 1}, Padding::kValid); 559 560 Array4D<float> expected(2, 2, 1, 1, {14240, 30496, 38816, 87840}); 561 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 562 } 563 564 XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input32x2x8x8) { 565 ComputationBuilder builder(client_, TestName()); 566 567 std::vector<float> input_data(32 * 2 * 8 * 8); 568 std::iota(input_data.begin(), input_data.end(), 0.0); 569 Array4D<float> input_array(32, 2, 8, 8, input_data); 570 auto input = builder.ConstantR4FromArray4D<float>(input_array); 571 572 std::vector<float> filter_data(2 * 2 * 8 * 8); 573 std::fill(filter_data.begin(), filter_data.begin() + filter_data.size() / 4, 574 1.0); 575 std::fill(filter_data.begin() + filter_data.size() / 4, 576 filter_data.begin() + filter_data.size() / 2, 2.0); 577 std::fill(filter_data.begin() + filter_data.size() / 2, 578 filter_data.begin() + 3 * filter_data.size() / 4, 3.0); 579 std::fill(filter_data.begin() + 3 * filter_data.size() / 4, filter_data.end(), 580 4.0); 581 const Array4D<float> filter_array(2, 2, 8, 8, filter_data); 582 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 583 584 builder.Conv(input, filter, {1, 1}, Padding::kValid); 585 586 std::vector<float> expected_data = { 587 14240, 30496, 38816, 87840, 63392, 145184, 87968, 588 202528, 112544, 259872, 137120, 317216, 161696, 374560, 589 186272, 431904, 210848, 489248, 235424, 546592, 260000, 590 603936, 284576, 661280, 309152, 718624, 333728, 775968, 591 358304, 833312, 382880, 890656, 407456, 948000, 432032, 592 1005344, 456608, 1062688, 481184, 1120032, 505760, 1177376, 593 530336, 1.23472e+06, 554912, 1292064, 579488, 1349408, 604064, 594 1406752, 628640, 1464096, 653216, 1.52144e+06, 677792, 1578784, 595 702368, 1636128, 726944, 1693472, 751520, 1750816, 776096, 596 1.80816e+06, 597 }; 598 Array4D<float> expected(32, 2, 1, 1, expected_data); 599 // The output elements can be larger than 1e+5, making the absolute error 600 // large sometimes. So, we focus on relative errors for this test case. 601 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 602 } 603 604 XLA_TEST_F(ConvolutionVariantsTest, Filter16x16x1x1Input16x16x1x1) { 605 ComputationBuilder builder(client_, TestName()); 606 607 Array4D<float> input_array(16, 16, 1, 1); 608 Array4D<float> filter_array(16, 16, 1, 1); 609 for (int i0 = 0; i0 < 16; ++i0) { 610 for (int i1 = 0; i1 < 16; ++i1) { 611 input_array(i0, i1, 0, 0) = 1000 * i0 + i1; 612 filter_array(i0, i1, 0, 0) = 1; 613 } 614 } 615 616 auto input = builder.ConstantR4FromArray4D<float>(input_array); 617 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 618 builder.Conv(input, filter, {1, 1}, Padding::kValid); 619 620 Array4D<float> expected(16, 16, 1, 1); 621 for (int i0 = 0; i0 < 16; ++i0) { 622 for (int i1 = 0; i1 < 16; ++i1) { 623 expected(i0, i1, 0, 0) = 16000 * i0 + 120; 624 } 625 } 626 627 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 628 } 629 630 XLA_TEST_F(ConvolutionVariantsTest, FlatRhsDilation) { 631 ComputationBuilder builder(client_, TestName()); 632 633 std::vector<float> input_data(1 * 1 * 4 * 6); 634 std::iota(input_data.begin(), input_data.end(), 0.0); 635 Array4D<float> input_array(1, 1, 4, 6, input_data); 636 637 Array4D<float> filter_array(1, 1, 2, 3, {1, 10, 100, 2, 20, 200}); 638 auto input = builder.ConstantR4FromArray4D<float>(input_array); 639 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 640 builder.ConvGeneralDilated( 641 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, /*padding=*/{}, 642 /*lhs_dilation=*/{}, /*rhs_dilation=*/{2, 2}, 643 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 644 645 Array4D<float> expected(1, 1, 2, 2, {3924, 4257, 5922, 6255}); 646 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 647 } 648 649 XLA_TEST_F(ConvolutionVariantsTest, FlatLhsDilation1D) { 650 ComputationBuilder builder(client_, TestName()); 651 652 std::vector<float> input_data(1 * 1 * 1 * 5); 653 std::iota(input_data.begin(), input_data.end(), 1.0); 654 Array4D<float> input_array(1, 1, 1, 5, input_data); 655 656 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 657 auto input = builder.ConstantR4FromArray4D<float>(input_array); 658 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 659 builder.ConvGeneralDilated( 660 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, /*padding=*/{}, 661 /*lhs_dilation=*/{1, 2}, /*rhs_dilation=*/{}, 662 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 663 664 Array4D<float> expected(1, 1, 1, 8, {10, 2, 20, 3, 30, 4, 40, 5}); 665 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 666 } 667 668 XLA_TEST_F(ConvolutionVariantsTest, FlatLhsDilation) { 669 ComputationBuilder builder(client_, TestName()); 670 671 std::vector<float> input_data(1 * 1 * 3 * 4); 672 std::iota(input_data.begin(), input_data.end(), 1.0); 673 Array4D<float> input_array(1, 1, 3, 4, input_data); 674 675 Array4D<float> filter_array(1, 1, 4, 3, 676 {100, 10, 1, // 677 200, 20, 2, // 678 300, 30, 3, // 679 400, 40, 4}); 680 auto input = builder.ConstantR4FromArray4D<float>(input_array); 681 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 682 builder.ConvGeneralDilated( 683 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{2, 1}, 684 /*padding=*/{{1, 0}, {0, 0}}, /*lhs_dilation=*/{3, 2}, 685 /*rhs_dilation=*/{}, 686 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 687 688 Array4D<float> expected(1, 1, 3, 5, 689 {204, 40, 406, 60, 608, // 690 1518, 180, 1821, 210, 2124, // 691 4146, 460, 4651, 510, 5156}); 692 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 693 } 694 695 XLA_TEST_F(ConvolutionVariantsTest, NegativePaddingOnBothEnds) { 696 ComputationBuilder builder(client_, TestName()); 697 698 std::vector<float> input_data(1 * 1 * 1 * 5); 699 std::iota(input_data.begin(), input_data.end(), 1.0); 700 Array4D<float> input_array(1, 1, 1, 5, input_data); 701 702 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 703 auto input = builder.ConstantR4FromArray4D<float>(input_array); 704 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 705 builder.ConvGeneral( 706 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, 707 /*padding=*/{{0, 0}, {-1, -1}}, 708 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 709 710 Array4D<float> expected(1, 1, 1, 2, {23, 34}); 711 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 712 } 713 714 XLA_TEST_F(ConvolutionVariantsTest, NegativePaddingLowAndPositivePaddingHigh) { 715 ComputationBuilder builder(client_, TestName()); 716 717 std::vector<float> input_data(1 * 1 * 1 * 5); 718 std::iota(input_data.begin(), input_data.end(), 1.0); 719 Array4D<float> input_array(1, 1, 1, 5, input_data); 720 721 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 722 auto input = builder.ConstantR4FromArray4D<float>(input_array); 723 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 724 builder.ConvGeneral( 725 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, 726 /*padding=*/{{0, 0}, {-1, 2}}, 727 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 728 729 Array4D<float> expected(1, 1, 1, 5, {23, 34, 45, 50, 0}); 730 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 731 } 732 733 XLA_TEST_F(ConvolutionVariantsTest, PositivePaddingLowAndNegativePaddingHigh) { 734 ComputationBuilder builder(client_, TestName()); 735 736 std::vector<float> input_data(1 * 1 * 1 * 5); 737 std::iota(input_data.begin(), input_data.end(), 1.0); 738 Array4D<float> input_array(1, 1, 1, 5, input_data); 739 740 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 741 auto input = builder.ConstantR4FromArray4D<float>(input_array); 742 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 743 builder.ConvGeneral( 744 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, 745 /*padding=*/{{0, 0}, {2, -1}}, 746 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 747 748 Array4D<float> expected(1, 1, 1, 5, {0, 1, 12, 23, 34}); 749 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 750 } 751 752 XLA_TEST_F(ConvolutionVariantsTest, PositivePaddingAndDilation) { 753 ComputationBuilder builder(client_, TestName()); 754 755 std::vector<float> input_data(1 * 1 * 1 * 5); 756 std::iota(input_data.begin(), input_data.end(), 1.0); 757 Array4D<float> input_array(1, 1, 1, 5, input_data); 758 759 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 760 auto input = builder.ConstantR4FromArray4D<float>(input_array); 761 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 762 builder.ConvGeneralDilated( 763 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, 764 /*padding=*/{{0, 0}, {3, 2}}, 765 /*lhs_dilation=*/{1, 2}, /*rhs_dilation=*/{1, 2}, 766 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 767 768 // input: 769 // [1, 2, 3, 4, 5] --dilate-> [1, 0, 2, 0, 3, 0, 4, 0, 5] 770 // ---pad---> [0, 0, 0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 0] 771 // filter: 772 // [10, 1] --dilate-> [10, 0, 1] 773 Array4D<float> expected(1, 1, 1, 12, 774 {0, 1, 0, 12, 0, 23, 0, 34, 0, 45, 0, 50}); 775 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 776 } 777 XLA_TEST_F(ConvolutionVariantsTest, NegativePaddingAndDilation) { 778 ComputationBuilder builder(client_, TestName()); 779 780 std::vector<float> input_data(1 * 1 * 1 * 5); 781 std::iota(input_data.begin(), input_data.end(), 1.0); 782 Array4D<float> input_array(1, 1, 1, 5, input_data); 783 784 Array4D<float> filter_array(1, 1, 1, 2, {10, 1}); 785 auto input = builder.ConstantR4FromArray4D<float>(input_array); 786 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 787 builder.ConvGeneralDilated( 788 /*lhs=*/input, /*rhs=*/filter, /*window_strides=*/{}, 789 /*padding=*/{{0, 0}, {-3, -2}}, 790 /*lhs_dilation=*/{1, 2}, /*rhs_dilation=*/{1, 2}, 791 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 792 793 // input: 794 // [1, 2, 3, 4, 5] --dilate-> [1, 0, 2, 0, 3, 0, 4, 0, 5] 795 // ---pad---> [0, 3, 0, 4] 796 // filter: 797 // [10, 1] --dilate-> [10, 0, 1] 798 Array4D<float> expected(1, 1, 1, 2, {0, 34}); 799 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 800 } 801 802 XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input1x1x2x3_Filter2x1x1x2) { 803 constexpr int bs = 1; 804 constexpr int iz = 1; 805 constexpr int oz = 2; 806 constexpr int iy = 2; 807 constexpr int ix = 3; 808 constexpr int ky = 1; 809 constexpr int kx = 2; 810 std::mt19937 rng; 811 std::uniform_real_distribution<float> distribution; 812 std::vector<float> input_data(bs * iz * iy * ix); 813 for (float& f : input_data) { 814 f = distribution(rng); 815 } 816 std::vector<float> kernel_data(oz * iz * ky * kx); 817 for (float& f : kernel_data) { 818 f = distribution(rng); 819 } 820 821 Array4D<float> input_array(bs, iz, iy, ix, input_data); 822 Array4D<float> filter_array(oz, iz, ky, kx, kernel_data); 823 824 ComputationBuilder builder(client_, TestName()); 825 auto input = builder.ConstantR4FromArray4D<float>(input_array); 826 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 827 builder.Conv(input, filter, {1, 1}, Padding::kValid); 828 829 std::unique_ptr<Array4D<float>> expected = ReferenceUtil::ConvArray4D( 830 input_array, filter_array, {1, 1}, Padding::kValid); 831 832 ComputeAndCompareR4<float>(&builder, *expected, {}, error_spec_); 833 } 834 835 XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input1x16x1x1_Filter1x16x1x1) { 836 constexpr int bs = 1; 837 constexpr int iz = 16; 838 constexpr int oz = 1; 839 constexpr int iy = 1; 840 constexpr int ix = 1; 841 constexpr int ky = 1; 842 constexpr int kx = 1; 843 std::mt19937 rng; 844 std::uniform_real_distribution<float> distribution; 845 std::vector<float> input_data(bs * iz * iy * ix); 846 for (float& f : input_data) { 847 f = distribution(rng); 848 } 849 std::vector<float> kernel_data(oz * iz * ky * kx); 850 for (float& f : kernel_data) { 851 f = distribution(rng); 852 } 853 854 Array4D<float> input_array(bs, iz, iy, ix, input_data); 855 Array4D<float> filter_array(oz, iz, ky, kx, kernel_data); 856 857 ComputationBuilder builder(client_, TestName()); 858 auto input = builder.ConstantR4FromArray4D<float>(input_array); 859 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 860 builder.Conv(input, filter, {1, 1}, Padding::kValid); 861 862 std::unique_ptr<Array4D<float>> expected = ReferenceUtil::ConvArray4D( 863 input_array, filter_array, {1, 1}, Padding::kValid); 864 865 ComputeAndCompareR4<float>(&builder, *expected, {}, error_spec_); 866 } 867 868 XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter1x16x1x1) { 869 constexpr int bs = 16; 870 constexpr int iz = 16; 871 constexpr int oz = 1; 872 constexpr int iy = 1; 873 constexpr int ix = 1; 874 constexpr int ky = 1; 875 constexpr int kx = 1; 876 std::mt19937 rng; 877 std::uniform_real_distribution<float> distribution; 878 std::vector<float> input_data(bs * iz * iy * ix); 879 for (float& f : input_data) { 880 f = distribution(rng); 881 } 882 std::vector<float> kernel_data(oz * iz * ky * kx); 883 for (float& f : kernel_data) { 884 f = distribution(rng); 885 } 886 887 Array4D<float> input_array(bs, iz, iy, ix, input_data); 888 Array4D<float> filter_array(oz, iz, ky, kx, kernel_data); 889 890 ComputationBuilder builder(client_, TestName()); 891 auto input = builder.ConstantR4FromArray4D<float>(input_array); 892 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 893 builder.Conv(input, filter, {1, 1}, Padding::kValid); 894 895 std::unique_ptr<Array4D<float>> expected = ReferenceUtil::ConvArray4D( 896 input_array, filter_array, {1, 1}, Padding::kValid); 897 898 ComputeAndCompareR4<float>(&builder, *expected, {}, error_spec_); 899 } 900 901 XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter16x16x1x1) { 902 constexpr int bs = 16; 903 constexpr int iz = 16; 904 constexpr int oz = 16; 905 constexpr int iy = 1; 906 constexpr int ix = 1; 907 constexpr int ky = 1; 908 constexpr int kx = 1; 909 std::mt19937 rng; 910 std::uniform_real_distribution<float> distribution; 911 std::vector<float> input_data(bs * iz * iy * ix); 912 for (float& f : input_data) { 913 f = distribution(rng); 914 } 915 std::vector<float> kernel_data(oz * iz * ky * kx); 916 for (float& f : kernel_data) { 917 f = distribution(rng); 918 } 919 920 Array4D<float> input_array(bs, iz, iy, ix, input_data); 921 Array4D<float> filter_array(oz, iz, ky, kx, kernel_data); 922 923 ComputationBuilder builder(client_, TestName()); 924 auto input = builder.ConstantR4FromArray4D<float>(input_array); 925 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 926 builder.Conv(input, filter, {1, 1}, Padding::kValid); 927 928 std::unique_ptr<Array4D<float>> expected = ReferenceUtil::ConvArray4D( 929 input_array, filter_array, {1, 1}, Padding::kValid); 930 931 ComputeAndCompareR4<float>(&builder, *expected, {}, error_spec_); 932 } 933 934 XLA_TEST_F(ConvolutionVariantsTest, 935 RandomData_Input16x16x16x16_Filter16x16x16x16) { 936 constexpr int bs = 16; 937 constexpr int iz = 16; 938 constexpr int oz = 16; 939 constexpr int iy = 16; 940 constexpr int ix = 16; 941 constexpr int ky = 16; 942 constexpr int kx = 16; 943 std::mt19937 rng; 944 std::uniform_real_distribution<float> distribution; 945 std::vector<float> input_data(bs * iz * iy * ix); 946 for (float& f : input_data) { 947 f = distribution(rng); 948 } 949 std::vector<float> kernel_data(oz * iz * ky * kx); 950 for (float& f : kernel_data) { 951 f = distribution(rng); 952 } 953 954 Array4D<float> input_array(bs, iz, iy, ix, input_data); 955 Array4D<float> filter_array(oz, iz, ky, kx, kernel_data); 956 957 ComputationBuilder builder(client_, TestName()); 958 auto input = builder.ConstantR4FromArray4D<float>(input_array); 959 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 960 builder.Conv(input, filter, {1, 1}, Padding::kValid); 961 962 std::unique_ptr<Array4D<float>> expected = ReferenceUtil::ConvArray4D( 963 input_array, filter_array, {1, 1}, Padding::kValid); 964 965 ComputeAndCompareR4<float>(&builder, *expected, {}, error_spec_); 966 } 967 968 XLA_TEST_F(ConvolutionVariantsTest, Filter1x2x1x1Input1x2x3x1GeneralPadding) { 969 ComputationBuilder builder(client_, TestName()); 970 971 std::vector<float> input_data(1 * 2 * 3 * 1); 972 std::iota(input_data.begin(), input_data.end(), 1.0); 973 Array4D<float> input_array(1, 2, 3, 1, input_data); 974 auto input = builder.ConstantR4FromArray4D<float>(input_array); 975 976 std::vector<float> filter_data(1 * 2 * 1 * 1); 977 std::iota(filter_data.begin(), filter_data.end(), 1.0); 978 Array4D<float> filter_array(1, 2, 1, 1, filter_data); 979 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 980 981 ConvolutionDimensionNumbers dnums; 982 // NHWC input format. 983 dnums.set_input_batch_dimension(0); 984 dnums.set_output_batch_dimension(0); 985 dnums.add_input_spatial_dimensions(1); 986 dnums.add_output_spatial_dimensions(1); 987 dnums.add_input_spatial_dimensions(2); 988 dnums.add_output_spatial_dimensions(2); 989 dnums.set_input_feature_dimension(3); 990 dnums.set_output_feature_dimension(3); 991 992 // Tensorflow filter shape: [ H, W, inC, outC ] 993 dnums.add_kernel_spatial_dimensions(0); 994 dnums.add_kernel_spatial_dimensions(1); 995 dnums.set_kernel_input_feature_dimension(2); 996 dnums.set_kernel_output_feature_dimension(3); 997 998 // Tests padding sizes that don't correspond either to SAME or VALID padding. 999 builder.ConvGeneral(input, filter, {1, 1}, {{2, 1}, {2, 3}}, dnums); 1000 1001 std::vector<float> expected_data = { 1002 0, 0, 0, 0, 0, 0, 0, // 1003 0, 0, 0, 0, 0, 0, 0, // 1004 0, 2, 5, 8, 3, 0, 0, // 1005 0, 8, 14, 17, 6, 0, 0, // 1006 0, 0, 0, 0, 0, 0, 0 // 1007 }; 1008 Array4D<float> expected(1, 5, 7, 1, expected_data); 1009 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 1010 } 1011 1012 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1GeneralPadding) { 1013 ComputationBuilder builder(client_, TestName()); 1014 1015 std::vector<float> input_data(1 * 2 * 3 * 1); 1016 std::iota(input_data.begin(), input_data.end(), 1.0); 1017 Array4D<float> input_array(1, 2, 3, 1, input_data); 1018 auto input = builder.ConstantR4FromArray4D<float>(input_array); 1019 1020 std::vector<float> filter_data(1 * 1 * 1 * 1); 1021 std::iota(filter_data.begin(), filter_data.end(), 2.0); 1022 Array4D<float> filter_array(1, 1, 1, 1, filter_data); 1023 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 1024 1025 ConvolutionDimensionNumbers dnums; 1026 // NHWC input format. 1027 dnums.set_input_batch_dimension(0); 1028 dnums.set_output_batch_dimension(0); 1029 dnums.add_input_spatial_dimensions(1); 1030 dnums.add_output_spatial_dimensions(1); 1031 dnums.add_input_spatial_dimensions(2); 1032 dnums.add_output_spatial_dimensions(2); 1033 dnums.set_input_feature_dimension(3); 1034 dnums.set_output_feature_dimension(3); 1035 1036 // Tensorflow filter shape: [ H, W, inC, outC ] 1037 dnums.add_kernel_spatial_dimensions(0); 1038 dnums.add_kernel_spatial_dimensions(1); 1039 dnums.set_kernel_input_feature_dimension(2); 1040 dnums.set_kernel_output_feature_dimension(3); 1041 1042 // Tests padding sizes that don't correspond either to SAME or VALID padding. 1043 builder.ConvGeneral(input, filter, {1, 1}, {{2, 1}, {2, 3}}, dnums); 1044 1045 std::vector<float> expected_data = { 1046 0, 0, 0, 0, 0, 0, 0, 0, // 1047 0, 0, 0, 0, 0, 0, 0, 0, // 1048 0, 0, 2, 4, 6, 0, 0, 0, // 1049 0, 0, 8, 10, 12, 0, 0, 0, // 1050 0, 0, 0, 0, 0, 0, 0, 0 // 1051 }; 1052 Array4D<float> expected(1, 5, 8, 1, expected_data); 1053 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 1054 } 1055 1056 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1NoPadding) { 1057 ComputationBuilder builder(client_, TestName()); 1058 1059 std::vector<float> input_data(1 * 2 * 3 * 1); 1060 std::iota(input_data.begin(), input_data.end(), 1.0); 1061 Array4D<float> input_array(1, 2, 3, 1, input_data); 1062 auto input = builder.ConstantR4FromArray4D<float>(input_array); 1063 1064 std::vector<float> filter_data(1 * 1 * 1 * 1); 1065 std::iota(filter_data.begin(), filter_data.end(), 2.0); 1066 Array4D<float> filter_array(1, 1, 1, 1, filter_data); 1067 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 1068 1069 ConvolutionDimensionNumbers dnums; 1070 // NHWC input format. 1071 dnums.set_input_batch_dimension(0); 1072 dnums.set_output_batch_dimension(0); 1073 dnums.add_input_spatial_dimensions(1); 1074 dnums.add_output_spatial_dimensions(1); 1075 dnums.add_input_spatial_dimensions(2); 1076 dnums.add_output_spatial_dimensions(2); 1077 dnums.set_input_feature_dimension(3); 1078 dnums.set_output_feature_dimension(3); 1079 1080 // Tensorflow filter shape: [ H, W, inC, outC ] 1081 dnums.add_kernel_spatial_dimensions(0); 1082 dnums.add_kernel_spatial_dimensions(1); 1083 dnums.set_kernel_input_feature_dimension(2); 1084 dnums.set_kernel_output_feature_dimension(3); 1085 1086 // Tests zero padding sizes. This can use matmul for computation. 1087 builder.ConvGeneral(input, filter, {1, 1}, {{0, 0}, {0, 0}}, dnums); 1088 1089 std::vector<float> expected_data = { 1090 2, 4, 6, // 1091 8, 10, 12, 1092 }; 1093 Array4D<float> expected(1, 2, 3, 1, expected_data); 1094 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 1095 } 1096 1097 XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x3Input1x2x3x2NoPadding) { 1098 ComputationBuilder builder(client_, TestName()); 1099 1100 std::vector<float> input_data(1 * 2 * 3 * 2); 1101 std::iota(input_data.begin(), input_data.end(), 1.0); 1102 Array4D<float> input_array(1, 2, 3, 2, input_data); 1103 auto input = builder.ConstantR4FromArray4D<float>(input_array); 1104 1105 std::vector<float> filter_data(1 * 1 * 2 * 3); 1106 std::iota(filter_data.begin(), filter_data.end(), 2.0); 1107 Array4D<float> filter_array(1, 1, 2, 3, filter_data); 1108 auto filter = builder.ConstantR4FromArray4D<float>(filter_array); 1109 1110 ConvolutionDimensionNumbers dnums; 1111 // NHWC input format. 1112 dnums.set_input_batch_dimension(0); 1113 dnums.set_output_batch_dimension(0); 1114 dnums.add_input_spatial_dimensions(1); 1115 dnums.add_output_spatial_dimensions(1); 1116 dnums.add_input_spatial_dimensions(2); 1117 dnums.add_output_spatial_dimensions(2); 1118 dnums.set_input_feature_dimension(3); 1119 dnums.set_output_feature_dimension(3); 1120 1121 // Tensorflow filter shape: [ H, W, inC, outC ] 1122 dnums.add_kernel_spatial_dimensions(0); 1123 dnums.add_kernel_spatial_dimensions(1); 1124 dnums.set_kernel_input_feature_dimension(2); 1125 dnums.set_kernel_output_feature_dimension(3); 1126 1127 // Tests zero padding sizes. This can use matmul for computation. 1128 builder.ConvGeneral(input, filter, {1, 1}, {{0, 0}, {0, 0}}, dnums); 1129 1130 std::vector<float> expected_data = { 1131 12, 15, 18, // 1132 26, 33, 40, // 1133 40, 51, 62, // 1134 54, 69, 84, // 1135 68, 87, 106, // 1136 82, 105, 128, // 1137 }; 1138 Array4D<float> expected(1, 2, 3, 3, expected_data); 1139 ComputeAndCompareR4<float>(&builder, expected, {}, error_spec_); 1140 } 1141 1142 // Regression test for b/32034796. 1143 // 1144 // XLA:GPU fuses 1145 // Conv([1,2,3], Reverse([5,6]), padding_low=1) 1146 // into 1147 // BackwardInputConv([1,2,3], [5,6], padding_low=0, padding_high=1) 1148 XLA_TEST_F(ConvolutionVariantsTest, 1149 BackwardInputLowPaddingLessThanHighPadding) { 1150 ComputationBuilder builder(client_, TestName()); 1151 1152 auto gradients = builder.ConstantR4FromArray4D<float>( 1153 Array4D<float>(1, 1, 1, 3, /*values=*/{1, 2, 3})); 1154 auto weights = builder.ConstantR4FromArray4D<float>( 1155 Array4D<float>(1, 1, 1, 2, /*values=*/{5, 6})); 1156 auto mirrored_weights = builder.Rev(weights, {2, 3}); 1157 builder.ConvWithGeneralPadding(gradients, mirrored_weights, 1158 /*window_strides=*/{1, 1}, 1159 /*padding=*/{{0, 0}, {1, 0}}); 1160 ComputeAndCompareR4<float>(&builder, {{{{5, 16, 27}}}}, {}, error_spec_); 1161 } 1162 1163 // XLA:GPU fuses 1164 // Conv([1], Reverse([1,10,100]), padding_high=3, base_dilation=3) 1165 // into 1166 // BackwardInputConv([1], [1,10,100], stride=3, padding=(2,1)) 1167 XLA_TEST_F(ConvolutionVariantsTest, 1168 BackwardInputLowPaddingGreaterThanHighPadding) { 1169 ComputationBuilder builder(client_, TestName()); 1170 1171 auto gradients = builder.ConstantR4FromArray4D<float>( 1172 Array4D<float>(1, 1, 1, 1, /*values=*/{1})); 1173 auto weights = builder.ConstantR4FromArray4D<float>( 1174 Array4D<float>(1, 1, 1, 3, /*values=*/{1, 10, 100})); 1175 auto mirrored_weights = builder.Rev(weights, {2, 3}); 1176 builder.ConvGeneralDilated( 1177 gradients, mirrored_weights, 1178 /*window_strides=*/{1, 1}, 1179 /*padding=*/{{0, 0}, {0, 3}}, 1180 /*lhs_dilation=*/{1, 3}, /*rhs_dilation=*/{}, 1181 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 1182 ComputeAndCompareR4<float>(&builder, {{{{100, 0}}}}, {}, error_spec_); 1183 } 1184 1185 // XLA:GPU fuses 1186 // Conv([1], Reverse([1,10,100]), padding=(1,1)) 1187 // into 1188 // BackwardInputConv([1], [1,10,100], padding=(1,1)) 1189 XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding) { 1190 ComputationBuilder builder(client_, TestName()); 1191 1192 auto gradients = builder.ConstantR4FromArray4D<float>( 1193 Array4D<float>(1, 1, 1, 1, /*values=*/{1})); 1194 auto weights = builder.ConstantR4FromArray4D<float>( 1195 Array4D<float>(1, 1, 1, 3, /*values=*/{1, 10, 100})); 1196 auto mirrored_weights = builder.Rev(weights, {2, 3}); 1197 builder.ConvWithGeneralPadding(gradients, mirrored_weights, 1198 /*window_strides=*/{1, 1}, 1199 /*padding=*/{{0, 0}, {1, 1}}); 1200 ComputeAndCompareR4<float>(&builder, {{{{10}}}}, {}, error_spec_); 1201 } 1202 1203 // HLO pattern 1204 // Conv([1,2,3], Reverse([1,10], padding_high=2) 1205 // could be fused to 1206 // BackwardInputConv([1,2,3], [1,10], padding_low=1, padding_high=-1) 1207 // 1208 // However, XLA:GPU doesn't actually fuse it because PadInsertion doesn't 1209 // support negative padding on backward convolution yet (b/32744257). 1210 XLA_TEST_F(ConvolutionVariantsTest, BackwardInputWithNegativePaddingHigh) { 1211 ComputationBuilder builder(client_, TestName()); 1212 1213 auto gradients = builder.ConstantR4FromArray4D<float>( 1214 Array4D<float>(1, 1, 1, 3, /*values=*/{1, 2, 3})); 1215 auto weights = builder.ConstantR4FromArray4D<float>( 1216 Array4D<float>(1, 1, 1, 2, /*values=*/{1, 10})); 1217 auto mirrored_weights = builder.Rev(weights, {2, 3}); 1218 builder.ConvWithGeneralPadding(gradients, mirrored_weights, 1219 /*window_strides=*/{1, 1}, 1220 /*padding=*/{{0, 0}, {0, 2}}); 1221 1222 ComputeAndCompareR4<float>(&builder, {{{{12, 23, 30, 0}}}}, {}, error_spec_); 1223 } 1224 1225 XLA_TEST_F(ConvolutionVariantsTest, 1226 BackwardFilterLowPaddingLessThanHighPadding) { 1227 ComputationBuilder builder(client_, TestName()); 1228 1229 // activations: 1,2,3,4 ---pad--> 0,1,2,3,4,0,0 1230 // gradients: 100,10,1 -dilate-> 100,0,10,0,1 1231 // weight gradients: 24,130,240 1232 // 1233 // This pattern will be fused to backward convolution with padding=(1,2). 1234 auto activations = builder.ConstantR4FromArray4D<float>( 1235 Array4D<float>(1, 1, 1, 4, /*values=*/{1, 2, 3, 4})); 1236 auto gradients = builder.ConstantR4FromArray4D<float>( 1237 Array4D<float>(1, 1, 1, 3, /*values=*/{100, 10, 1})); 1238 auto forward_conv = builder.ConvGeneralDilated( 1239 activations, gradients, 1240 /*window_strides=*/{1, 1}, 1241 /*padding=*/{{0, 0}, {1, 2}}, 1242 /*lhs_dilation=*/{}, /*rhs_dilation=*/{1, 2}, 1243 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 1244 builder.Transpose(forward_conv, {0, 1, 2, 3}); 1245 1246 ComputeAndCompareR4<float>(&builder, {{{{24, 130, 240}}}}, {}, error_spec_); 1247 } 1248 1249 XLA_TEST_F(ConvolutionVariantsTest, 1250 BackwardFilterLowPaddingGreaterThanHighPadding) { 1251 ComputationBuilder builder(client_, TestName()); 1252 1253 // activations: 1,2,3,4 ---pad--> 0,0,1,2,3,4 1254 // gradients: 100,10,1 -dilate-> 100,0,10,0,1 1255 // weight gradients: 13,24 1256 // 1257 // This pattern will be fused to backward convolution with padding=(2,1). 1258 // Note: both (2,1) and (2,0) are valid padding for the backward convolution 1259 // because the stride is 2. 1260 auto activations = builder.ConstantR4FromArray4D<float>( 1261 Array4D<float>(1, 1, 1, 4, /*values=*/{1, 2, 3, 4})); 1262 auto gradients = builder.ConstantR4FromArray4D<float>( 1263 Array4D<float>(1, 1, 1, 3, /*values=*/{100, 10, 1})); 1264 auto forward_conv = builder.ConvGeneralDilated( 1265 activations, gradients, 1266 /*window_strides=*/{1, 1}, 1267 /*padding=*/{{0, 0}, {2, 0}}, 1268 /*lhs_dilation=*/{}, /*rhs_dilation=*/{1, 2}, 1269 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 1270 builder.Transpose(forward_conv, {0, 1, 2, 3}); 1271 1272 ComputeAndCompareR4<float>(&builder, {{{{13, 24}}}}, {}, error_spec_); 1273 } 1274 1275 XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding) { 1276 ComputationBuilder builder(client_, TestName()); 1277 1278 // activations: 1,2,3,4 ---pad--> 0,0,1,2,3,4,0 1279 // gradients: 100,10,1 -dilate-> 100,0,10,0,1 1280 // weight gradients: 13,24,130 1281 // 1282 // This pattern will be fused to backward convolution with padding=(2,2). 1283 // Note: both (2,1) and (2,2) are valid padding for the backward convolution 1284 // because the stride is 2. ConvolutionFolding prefers (2,2) because cuDNN 1285 // supports even padding only -- using (2,1) would need extra effort of 1286 // canonicalization. 1287 auto activations = builder.ConstantR4FromArray4D<float>( 1288 Array4D<float>(1, 1, 1, 4, /*values=*/{1, 2, 3, 4})); 1289 auto gradients = builder.ConstantR4FromArray4D<float>( 1290 Array4D<float>(1, 1, 1, 3, /*values=*/{100, 10, 1})); 1291 auto forward_conv = builder.ConvGeneralDilated( 1292 activations, gradients, 1293 /*window_strides=*/{1, 1}, 1294 /*padding=*/{{0, 0}, {2, 1}}, 1295 /*lhs_dilation=*/{}, /*rhs_dilation=*/{1, 2}, 1296 ComputationBuilder::CreateDefaultConvDimensionNumbers()); 1297 builder.Transpose(forward_conv, {0, 1, 2, 3}); 1298 1299 ComputeAndCompareR4<float>(&builder, {{{{13, 24, 130}}}}, {}, error_spec_); 1300 } 1301 1302 XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding1D) { 1303 ComputationBuilder builder(client_, TestName()); 1304 1305 auto gradients = builder.ConstantR3FromArray3D<float>( 1306 Array3D<float>(1, 1, 1, /*value=*/1)); 1307 auto weights = 1308 builder.ConstantR3FromArray3D<float>(Array3D<float>({{{1, 10, 100}}})); 1309 auto mirrored_weights = builder.Rev(weights, {2}); 1310 builder.ConvWithGeneralPadding(gradients, mirrored_weights, 1311 /*window_strides=*/{1}, 1312 /*padding=*/{{1, 1}}); 1313 ComputeAndCompareR3<float>(&builder, {{{10}}}, {}, error_spec_); 1314 } 1315 1316 XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding1D) { 1317 ComputationBuilder builder(client_, TestName()); 1318 1319 auto activations = 1320 builder.ConstantR3FromArray3D<float>(Array3D<float>({{{1, 2, 3, 4}}})); 1321 auto gradients = 1322 builder.ConstantR3FromArray3D<float>(Array3D<float>({{{100, 10, 1}}})); 1323 auto forward_conv = builder.ConvGeneralDilated( 1324 activations, gradients, 1325 /*window_strides=*/{1}, 1326 /*padding=*/{{2, 1}}, 1327 /*lhs_dilation=*/{}, /*rhs_dilation=*/{2}, 1328 ComputationBuilder::CreateDefaultConvDimensionNumbers( 1329 /*num_spatial_dims=*/1)); 1330 builder.Transpose(forward_conv, {0, 1, 2}); 1331 1332 ComputeAndCompareR3<float>(&builder, {{{13, 24, 130}}}, {}, error_spec_); 1333 } 1334 1335 XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding3D) { 1336 ComputationBuilder builder(client_, TestName()); 1337 1338 auto gradients_flat = Literal::CreateR1<float>({1}); 1339 auto gradients_literal = 1340 gradients_flat->Reshape({1, 1, 1, 1, 1}).ConsumeValueOrDie(); 1341 auto gradients = builder.ConstantLiteral(*gradients_literal); 1342 1343 auto weights_flat = Literal::CreateR1<float>({1, 10, 100}); 1344 auto weights_literal = 1345 weights_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); 1346 auto weights = builder.ConstantLiteral(*weights_literal); 1347 1348 auto expected_flat = Literal::CreateR1<float>({10}); 1349 auto expected_literal = 1350 expected_flat->Reshape({1, 1, 1, 1, 1}).ConsumeValueOrDie(); 1351 1352 auto mirrored_weights = builder.Rev(weights, {2, 3, 4}); 1353 builder.ConvWithGeneralPadding(gradients, mirrored_weights, 1354 /*window_strides=*/{1, 1, 1}, 1355 /*padding=*/{{0, 0}, {0, 0}, {1, 1}}); 1356 ComputeAndCompareLiteral(&builder, *expected_literal, {}, error_spec_); 1357 } 1358 1359 XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding3D) { 1360 ComputationBuilder builder(client_, TestName()); 1361 1362 auto activations_flat = Literal::CreateR1<float>({1, 2, 3, 4}); 1363 auto activations_literal = 1364 activations_flat->Reshape({1, 1, 1, 1, 4}).ConsumeValueOrDie(); 1365 auto activations = builder.ConstantLiteral(*activations_literal); 1366 1367 auto gradients_flat = Literal::CreateR1<float>({100, 10, 1}); 1368 auto gradients_literal = 1369 gradients_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); 1370 auto gradients = builder.ConstantLiteral(*gradients_literal); 1371 1372 auto expected_flat = Literal::CreateR1<float>({13, 24, 130}); 1373 auto expected_literal = 1374 expected_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); 1375 1376 auto forward_conv = builder.ConvGeneralDilated( 1377 activations, gradients, 1378 /*window_strides=*/{1, 1, 1}, 1379 /*padding=*/{{0, 0}, {0, 0}, {2, 1}}, 1380 /*lhs_dilation=*/{}, /*rhs_dilation=*/{1, 1, 2}, 1381 ComputationBuilder::CreateDefaultConvDimensionNumbers( 1382 /*num_spatial_dims=*/3)); 1383 builder.Transpose(forward_conv, {0, 1, 2, 3, 4}); 1384 ComputeAndCompareLiteral(&builder, *expected_literal, {}, error_spec_); 1385 } 1386 1387 } // namespace 1388 } // namespace xla 1389
