1 /* Copyright 2018 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 <functional> 17 #include <memory> 18 19 #include "tensorflow/core/common_runtime/device.h" 20 #include "tensorflow/core/common_runtime/device_factory.h" 21 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h" 22 #include "tensorflow/core/framework/allocator.h" 23 #include "tensorflow/core/framework/fake_input.h" 24 #include "tensorflow/core/framework/node_def_builder.h" 25 #include "tensorflow/core/framework/op_kernel.h" 26 #include "tensorflow/core/framework/tensor.h" 27 #include "tensorflow/core/framework/types.h" 28 #include "tensorflow/core/framework/types.pb.h" 29 #include "tensorflow/core/kernels/ops_testutil.h" 30 #include "tensorflow/core/kernels/ops_util.h" 31 #include "tensorflow/core/lib/io/path.h" 32 #include "tensorflow/core/lib/strings/strcat.h" 33 #include "tensorflow/core/platform/test.h" 34 #include "tensorflow/core/platform/test_benchmark.h" 35 36 namespace tensorflow { 37 namespace { 38 39 class RollOpTest : public OpsTestBase { 40 protected: 41 void MakeOp(DataType data_type, DataType index_type) { 42 TF_ASSERT_OK(NodeDefBuilder("myop", "Roll") 43 .Input(FakeInput(data_type)) 44 .Input(FakeInput(index_type)) 45 .Input(FakeInput(index_type)) 46 .Finalize(node_def())); 47 TF_ASSERT_OK(InitOp()); 48 } 49 }; 50 51 TEST_F(RollOpTest, ScalarIndices) { 52 MakeOp(DT_FLOAT, DT_INT32); 53 54 // Feed and run 55 AddInputFromArray<float>(TensorShape({5}), {0, 1, 2, 3, 4}); 56 AddInputFromArray<int32>(TensorShape({}), {3}); 57 AddInputFromArray<int32>(TensorShape({}), {0}); 58 TF_ASSERT_OK(RunOpKernel()); 59 60 // Check the output. 61 Tensor expected(allocator(), DT_FLOAT, TensorShape({5})); 62 test::FillValues<float>(&expected, {2, 3, 4, 0, 1}); 63 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 64 } 65 66 TEST_F(RollOpTest, ScalarIndices_NoMemcpy) { 67 MakeOp(DT_STRING, DT_INT32); 68 69 // Feed and run 70 AddInputFromArray<string>(TensorShape({5}), {"a", "b", "c", "d", "e"}); 71 AddInputFromArray<int32>(TensorShape({}), {3}); 72 AddInputFromArray<int32>(TensorShape({}), {0}); 73 TF_ASSERT_OK(RunOpKernel()); 74 75 // Check the output. 76 Tensor expected(allocator(), DT_STRING, TensorShape({5})); 77 test::FillValues<string>(&expected, {"c", "d", "e", "a", "b"}); 78 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 79 } 80 81 TEST_F(RollOpTest, ScalarIndices_Complex) { 82 MakeOp(DT_COMPLEX64, DT_INT32); 83 84 // Feed and run 85 AddInputFromArray<std::complex<float>>( 86 TensorShape({5}), {std::complex<float>(0, 10), std::complex<float>(1, 11), 87 std::complex<float>(2, 12), std::complex<float>(3, 13), 88 std::complex<float>(4, 14)}); 89 AddInputFromArray<int32>(TensorShape({}), {3}); 90 AddInputFromArray<int32>(TensorShape({}), {0}); 91 TF_ASSERT_OK(RunOpKernel()); 92 93 // Check the output. 94 Tensor expected(allocator(), DT_COMPLEX64, TensorShape({5})); 95 test::FillValues<std::complex<float>>( 96 &expected, {std::complex<float>(2, 12), std::complex<float>(3, 13), 97 std::complex<float>(4, 14), std::complex<float>(0, 10), 98 std::complex<float>(1, 11)}); 99 test::ExpectTensorEqual<std::complex<float>>(expected, *GetOutput(0)); 100 } 101 102 TEST_F(RollOpTest, Simple_TwoD32) { 103 MakeOp(DT_FLOAT, DT_INT32); 104 105 // Feed and run 106 AddInputFromArray<float>(TensorShape({3, 5}), 107 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}); 108 AddInputFromArray<int32>(TensorShape({2}), {2, -1}); 109 AddInputFromArray<int32>(TensorShape({2}), {0, 1}); 110 TF_ASSERT_OK(RunOpKernel()); 111 112 // Check the output. 113 Tensor expected(allocator(), DT_FLOAT, TensorShape({3, 5})); 114 test::FillValues<float>(&expected, 115 {6, 7, 8, 9, 5, 11, 12, 13, 14, 10, 1, 2, 3, 4, 0}); 116 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 117 } 118 119 TEST_F(RollOpTest, Simple_TwoD32_NoMemcpy) { 120 MakeOp(DT_STRING, DT_INT32); 121 122 // Feed and run 123 AddInputFromArray<string>(TensorShape({3, 5}), 124 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", 125 "k", "l", "m", "n", "o"}); 126 AddInputFromArray<int32>(TensorShape({2}), {2, -1}); 127 AddInputFromArray<int32>(TensorShape({2}), {0, 1}); 128 TF_ASSERT_OK(RunOpKernel()); 129 130 // Check the output. 131 Tensor expected(allocator(), DT_STRING, TensorShape({3, 5})); 132 test::FillValues<string>(&expected, {"g", "h", "i", "j", "f", "l", "m", "n", 133 "o", "k", "b", "c", "d", "e", "a"}); 134 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 135 } 136 137 TEST_F(RollOpTest, Simple_ThreeD32) { 138 MakeOp(DT_FLOAT, DT_INT32); 139 140 // Feed and run 141 AddInputFromArray<float>(TensorShape({2, 2, 3}), 142 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); 143 AddInputFromArray<int32>(TensorShape({3}), {1, -1, -1}); 144 AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2}); 145 TF_ASSERT_OK(RunOpKernel()); 146 147 // Check the output. 148 Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 3})); 149 test::FillValues<float>(&expected, {10, 11, 9, 7, 8, 6, 4, 5, 3, 1, 2, 0}); 150 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 151 } 152 153 TEST_F(RollOpTest, Simple_ThreeD32_NoMemcpy) { 154 MakeOp(DT_STRING, DT_INT32); 155 156 // Feed and run 157 AddInputFromArray<string>( 158 TensorShape({2, 2, 3}), 159 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"}); 160 AddInputFromArray<int32>(TensorShape({3}), {1, -1, -1}); 161 AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2}); 162 TF_ASSERT_OK(RunOpKernel()); 163 164 // Check the output. 165 Tensor expected(allocator(), DT_STRING, TensorShape({2, 2, 3})); 166 test::FillValues<string>( 167 &expected, {"k", "l", "j", "h", "i", "g", "e", "f", "d", "b", "c", "a"}); 168 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 169 } 170 171 TEST_F(RollOpTest, Simple_TwoD64) { 172 MakeOp(DT_FLOAT, DT_INT64); 173 174 // Feed and run 175 AddInputFromArray<float>(TensorShape({5, 3}), 176 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}); 177 AddInputFromArray<int64>(TensorShape({2}), {-1, 4}); 178 AddInputFromArray<int64>(TensorShape({2}), {0, 1}); 179 TF_ASSERT_OK(RunOpKernel()); 180 181 // Check the output. 182 Tensor expected(allocator(), DT_FLOAT, TensorShape({5, 3})); 183 test::FillValues<float>(&expected, 184 {5, 3, 4, 8, 6, 7, 11, 9, 10, 14, 12, 13, 2, 0, 1}); 185 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 186 } 187 188 TEST_F(RollOpTest, Simple_TwoD64_NoMemcpy) { 189 MakeOp(DT_STRING, DT_INT64); 190 191 // Feed and run 192 AddInputFromArray<string>(TensorShape({5, 3}), 193 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", 194 "k", "l", "m", "n", "o"}); 195 AddInputFromArray<int64>(TensorShape({2}), {-1, 4}); 196 AddInputFromArray<int64>(TensorShape({2}), {0, 1}); 197 TF_ASSERT_OK(RunOpKernel()); 198 199 // Check the output. 200 Tensor expected(allocator(), DT_STRING, TensorShape({5, 3})); 201 test::FillValues<string>(&expected, {"f", "d", "e", "i", "g", "h", "l", "j", 202 "k", "o", "m", "n", "c", "a", "b"}); 203 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 204 } 205 206 TEST_F(RollOpTest, Simple_ThreeD64) { 207 MakeOp(DT_FLOAT, DT_INT64); 208 209 // Feed and run 210 AddInputFromArray<float>(TensorShape({4, 1, 3}), 211 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); 212 AddInputFromArray<int64>(TensorShape({3}), {4, 3, 2}); 213 AddInputFromArray<int64>(TensorShape({3}), {0, 1, 2}); 214 TF_ASSERT_OK(RunOpKernel()); 215 216 // Check the output. 217 Tensor expected(allocator(), DT_FLOAT, TensorShape({4, 1, 3})); 218 test::FillValues<float>(&expected, {1, 2, 0, 4, 5, 3, 7, 8, 6, 10, 11, 9}); 219 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 220 } 221 222 TEST_F(RollOpTest, Simple_ThreeD64_NoMemcpy) { 223 MakeOp(DT_STRING, DT_INT64); 224 225 // Feed and run 226 AddInputFromArray<string>( 227 TensorShape({4, 1, 3}), 228 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"}); 229 AddInputFromArray<int64>(TensorShape({3}), {4, 3, 2}); 230 AddInputFromArray<int64>(TensorShape({3}), {0, 1, 2}); 231 TF_ASSERT_OK(RunOpKernel()); 232 233 // Check the output. 234 Tensor expected(allocator(), DT_STRING, TensorShape({4, 1, 3})); 235 test::FillValues<string>( 236 &expected, {"b", "c", "a", "e", "f", "d", "h", "i", "g", "k", "l", "j"}); 237 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 238 } 239 240 TEST_F(RollOpTest, ZeroShift_ThreeD32) { 241 MakeOp(DT_FLOAT, DT_INT32); 242 243 // Feed and run 244 AddInputFromArray<float>(TensorShape({2, 2, 3}), 245 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); 246 AddInputFromArray<int32>(TensorShape({3}), {0, 0, 0}); 247 AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2}); 248 TF_ASSERT_OK(RunOpKernel()); 249 250 // Check the output. 251 Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 3})); 252 test::FillValues<float>(&expected, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); 253 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 254 } 255 256 TEST_F(RollOpTest, ZeroShift_ThreeD32_NoMemcpy) { 257 MakeOp(DT_STRING, DT_INT32); 258 259 // Feed and run 260 AddInputFromArray<string>( 261 TensorShape({2, 2, 3}), 262 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"}); 263 AddInputFromArray<int32>(TensorShape({3}), {0, 0, 0}); 264 AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2}); 265 TF_ASSERT_OK(RunOpKernel()); 266 267 // Check the output. 268 Tensor expected(allocator(), DT_STRING, TensorShape({2, 2, 3})); 269 test::FillValues<string>( 270 &expected, {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"}); 271 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 272 } 273 274 TEST_F(RollOpTest, ZeroSize_ThreeD32) { 275 MakeOp(DT_FLOAT, DT_INT32); 276 277 // Feed and run 278 AddInputFromArray<float>(TensorShape({5, 0, 0}), {}); 279 AddInputFromArray<int32>(TensorShape({}), {1}); 280 AddInputFromArray<int32>(TensorShape({}), {0}); 281 TF_ASSERT_OK(RunOpKernel()); 282 283 // Check the output. 284 Tensor expected(allocator(), DT_FLOAT, TensorShape({5, 0, 0})); 285 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 286 } 287 288 TEST_F(RollOpTest, ZeroSize_ThreeD32_NoMemcpy) { 289 MakeOp(DT_STRING, DT_INT32); 290 291 // Feed and run 292 AddInputFromArray<string>(TensorShape({5, 0, 0}), {}); 293 AddInputFromArray<int32>(TensorShape({}), {1}); 294 AddInputFromArray<int32>(TensorShape({}), {0}); 295 TF_ASSERT_OK(RunOpKernel()); 296 297 // Check the output. 298 Tensor expected(allocator(), DT_STRING, TensorShape({5, 0, 0})); 299 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 300 } 301 302 TEST_F(RollOpTest, OneSize_ThreeD32) { 303 MakeOp(DT_FLOAT, DT_INT32); 304 305 // Feed and run 306 AddInputFromArray<float>(TensorShape({1, 1, 1}), {5}); 307 AddInputFromArray<int32>(TensorShape({}), {1}); 308 AddInputFromArray<int32>(TensorShape({}), {0}); 309 TF_ASSERT_OK(RunOpKernel()); 310 311 // Check the output. 312 Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 1, 1})); 313 test::FillValues<float>(&expected, {5}); 314 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 315 } 316 317 TEST_F(RollOpTest, OneSize_ThreeD32_NoMemcpy) { 318 MakeOp(DT_STRING, DT_INT32); 319 320 // Feed and run 321 AddInputFromArray<string>(TensorShape({1, 1, 1}), {"a"}); 322 AddInputFromArray<int32>(TensorShape({}), {1}); 323 AddInputFromArray<int32>(TensorShape({}), {0}); 324 TF_ASSERT_OK(RunOpKernel()); 325 326 // Check the output. 327 Tensor expected(allocator(), DT_STRING, TensorShape({1, 1, 1})); 328 test::FillValues<string>(&expected, {"a"}); 329 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 330 } 331 332 TEST_F(RollOpTest, MultiShifts_TwoD32) { 333 MakeOp(DT_FLOAT, DT_INT32); 334 335 // Feed and run 336 AddInputFromArray<float>(TensorShape({3, 5}), 337 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}); 338 AddInputFromArray<int32>(TensorShape({4}), {-2, 2, -1, 1}); 339 AddInputFromArray<int32>(TensorShape({4}), {1, 0, 0, 1}); 340 TF_ASSERT_OK(RunOpKernel()); 341 342 // Check the output. 343 Tensor expected(allocator(), DT_FLOAT, TensorShape({3, 5})); 344 test::FillValues<float>(&expected, 345 {11, 12, 13, 14, 10, 1, 2, 3, 4, 0, 6, 7, 8, 9, 5}); 346 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 347 } 348 349 TEST_F(RollOpTest, MultiShifts_TwoD32_NoMemcpy) { 350 MakeOp(DT_STRING, DT_INT32); 351 352 // Feed and run 353 AddInputFromArray<string>(TensorShape({3, 5}), 354 {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", 355 "k", "l", "m", "n", "o"}); 356 AddInputFromArray<int32>(TensorShape({4}), {-2, 2, -1, 1}); 357 AddInputFromArray<int32>(TensorShape({4}), {1, 0, 0, 1}); 358 TF_ASSERT_OK(RunOpKernel()); 359 360 // Check the output. 361 Tensor expected(allocator(), DT_STRING, TensorShape({3, 5})); 362 test::FillValues<string>(&expected, {"l", "m", "n", "o", "k", "b", "c", "d", 363 "e", "a", "g", "h", "i", "j", "f"}); 364 test::ExpectTensorEqual<string>(expected, *GetOutput(0)); 365 } 366 367 TEST_F(RollOpTest, Error_InputMustBeVectorOrHigher) { 368 MakeOp(DT_FLOAT, DT_INT32); 369 370 // Feed and run 371 AddInputFromArray<float>(TensorShape({}), {7}); 372 AddInputFromArray<int32>(TensorShape({}), {1}); 373 AddInputFromArray<int32>(TensorShape({}), {0}); 374 Status s = RunOpKernel(); 375 EXPECT_TRUE(StringPiece(s.ToString()).contains("input must be 1-D or higher")) 376 << s; 377 } 378 379 TEST_F(RollOpTest, Error_AxisMustBeScalarOrVector) { 380 MakeOp(DT_FLOAT, DT_INT32); 381 382 // Feed and run 383 AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4}); 384 AddInputFromArray<int32>(TensorShape({}), {1}); 385 AddInputFromArray<int32>(TensorShape({1, 2}), {0, 1}); 386 Status s = RunOpKernel(); 387 EXPECT_TRUE(StringPiece(s.ToString()) 388 .contains("axis must be a scalar or a 1-D vector")) 389 << s; 390 } 391 392 TEST_F(RollOpTest, Error_ShiftMustBeScalarOrVector) { 393 MakeOp(DT_FLOAT, DT_INT32); 394 395 // Feed and run 396 AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4}); 397 AddInputFromArray<int32>(TensorShape({1, 2}), {0, 1}); 398 AddInputFromArray<int32>(TensorShape({}), {1}); 399 Status s = RunOpKernel(); 400 EXPECT_TRUE(StringPiece(s.ToString()) 401 .contains("shift must be a scalar or a 1-D vector")) 402 << s; 403 } 404 405 TEST_F(RollOpTest, Error_ShiftAndAxisMustBeSameSize) { 406 MakeOp(DT_FLOAT, DT_INT32); 407 408 // Feed and run 409 AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4}); 410 AddInputFromArray<int32>(TensorShape({1}), {1}); 411 AddInputFromArray<int32>(TensorShape({2}), {0, 1}); 412 Status s = RunOpKernel(); 413 EXPECT_TRUE(StringPiece(s.ToString()) 414 .contains("shift and axis must have the same size")) 415 << s; 416 } 417 418 TEST_F(RollOpTest, Error_AxisOutOfRange) { 419 MakeOp(DT_FLOAT, DT_INT32); 420 421 // Feed and run 422 AddInputFromArray<float>(TensorShape({4}), {1, 2, 3, 4}); 423 AddInputFromArray<int32>(TensorShape({}), {1}); 424 AddInputFromArray<int32>(TensorShape({}), {1}); 425 Status s = RunOpKernel(); 426 EXPECT_TRUE(StringPiece(s.ToString()).contains("is out of range")) << s; 427 } 428 429 // isd - (inner shift dimension) The inner most dimension to be shifted. 430 // All outer dimensions will also be shifted for testing. 431 static Graph* RollGraph(const TensorShape& shape, int isd) { 432 Graph* g = new Graph(OpRegistry::Global()); 433 Tensor input(DT_FLOAT, shape); 434 input.flat<float>().setRandom(); 435 const int dims = static_cast<int>(input.dims()); 436 Tensor shift(DT_INT32, TensorShape({dims})); 437 for (int i = 0; i < dims; i++) { 438 // shift the inner shift dimension and all outer dimensions 439 shift.flat<int32>()(i) = (i <= isd) ? 2 : 0; 440 } 441 Tensor axis(DT_INT32, TensorShape({dims})); 442 for (int i = 0; i < dims; i++) { 443 axis.flat<int32>()(i) = i; 444 } 445 test::graph::Roll(g, test::graph::Constant(g, input), 446 test::graph::Constant(g, shift), 447 test::graph::Constant(g, axis)); 448 return g; 449 } 450 451 #define BM_ROLL_OUTER(DEVICE) \ 452 static void BM_##DEVICE##_roll_outer(int iters, int rows, int columns) { \ 453 TensorShape shape{rows, columns}; \ 454 const int64 num_items = static_cast<int64>(iters) * shape.num_elements(); \ 455 testing::ItemsProcessed(num_items); \ 456 testing::BytesProcessed(num_items * sizeof(float)); \ 457 testing::UseRealTime(); \ 458 test::Benchmark(#DEVICE, RollGraph(shape, 0)).Run(iters); \ 459 } \ 460 BENCHMARK(BM_##DEVICE##_roll_outer) \ 461 ->ArgPair(256, 256) \ 462 ->ArgPair(512, 512) \ 463 ->ArgPair(1024, 1024) \ 464 ->ArgPair(2048, 2048) 465 466 #define BM_ROLL_ALL(DEVICE) \ 467 static void BM_##DEVICE##_roll_all(int iters, int rows, int columns) { \ 468 TensorShape shape{rows, columns}; \ 469 const int64 num_items = static_cast<int64>(iters) * shape.num_elements(); \ 470 testing::ItemsProcessed(num_items); \ 471 testing::BytesProcessed(num_items * sizeof(float)); \ 472 testing::UseRealTime(); \ 473 test::Benchmark(#DEVICE, RollGraph(shape, 1)).Run(iters); \ 474 } \ 475 BENCHMARK(BM_##DEVICE##_roll_all) \ 476 ->ArgPair(256, 256) \ 477 ->ArgPair(512, 512) \ 478 ->ArgPair(1024, 1024) \ 479 ->ArgPair(2048, 2048) 480 481 BM_ROLL_OUTER(cpu); 482 BM_ROLL_ALL(cpu); 483 } // namespace 484 } // namespace tensorflow 485
