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