1 /* Copyright 2015 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/kernel_benchmark_testlib.h" 20 #include "tensorflow/core/framework/allocator.h" 21 #include "tensorflow/core/framework/fake_input.h" 22 #include "tensorflow/core/framework/node_def_builder.h" 23 #include "tensorflow/core/framework/op_kernel.h" 24 #include "tensorflow/core/framework/tensor.h" 25 #include "tensorflow/core/framework/types.h" 26 #include "tensorflow/core/framework/types.pb.h" 27 #include "tensorflow/core/graph/node_builder.h" 28 #include "tensorflow/core/graph/testlib.h" 29 #include "tensorflow/core/kernels/ops_testutil.h" 30 #include "tensorflow/core/kernels/ops_util.h" 31 #include "tensorflow/core/lib/core/status_test_util.h" 32 #include "tensorflow/core/lib/random/simple_philox.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 DynamicPartitionOpTest : public OpsTestBase { 40 protected: 41 void MakeOp() { 42 TF_ASSERT_OK(NodeDefBuilder("myop", "DynamicPartition") 43 .Input(FakeInput(DT_FLOAT)) 44 .Input(FakeInput(DT_INT32)) 45 .Attr("num_partitions", 4) 46 .Finalize(node_def())); 47 TF_ASSERT_OK(InitOp()); 48 } 49 }; 50 51 TEST_F(DynamicPartitionOpTest, Simple_OneD) { 52 MakeOp(); 53 54 // Similar to how we would use this to split embedding ids to be looked up 55 56 // Feed and run 57 AddInputFromArray<float>(TensorShape({6}), {0, 13, 2, 39, 4, 17}); 58 AddInputFromArray<int32>(TensorShape({6}), {0, 0, 2, 3, 2, 1}); 59 TF_ASSERT_OK(RunOpKernel()); 60 61 // Check the output sizes 62 { // Output 0 63 Tensor expected(allocator(), DT_FLOAT, TensorShape({2})); 64 test::FillValues<float>(&expected, {0, 13}); 65 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 66 } 67 { // Output 1 68 Tensor expected(allocator(), DT_FLOAT, TensorShape({1})); 69 test::FillValues<float>(&expected, {17}); 70 test::ExpectTensorEqual<float>(expected, *GetOutput(1)); 71 } 72 { // Output 2 73 Tensor expected(allocator(), DT_FLOAT, TensorShape({2})); 74 test::FillValues<float>(&expected, {2, 4}); 75 test::ExpectTensorEqual<float>(expected, *GetOutput(2)); 76 } 77 { // Output 3 78 Tensor expected(allocator(), DT_FLOAT, TensorShape({1})); 79 test::FillValues<float>(&expected, {39}); 80 test::ExpectTensorEqual<float>(expected, *GetOutput(3)); 81 } 82 } 83 84 TEST_F(DynamicPartitionOpTest, Simple_TwoD) { 85 MakeOp(); 86 87 // Feed and run 88 AddInputFromArray<float>( 89 TensorShape({6, 3}), 90 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17}); 91 AddInputFromArray<int32>(TensorShape({6}), {0, 0, 2, 3, 2, 1}); 92 TF_ASSERT_OK(RunOpKernel()); 93 94 // Check the output sizes 95 { // Output 0 96 Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 3})); 97 test::FillValues<float>(&expected, {0, 1, 2, 3, 4, 5}); 98 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 99 } 100 { // Output 1 101 Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 3})); 102 test::FillValues<float>(&expected, {15, 16, 17}); 103 test::ExpectTensorEqual<float>(expected, *GetOutput(1)); 104 } 105 { // Output 2 106 Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 3})); 107 test::FillValues<float>(&expected, {6, 7, 8, 12, 13, 14}); 108 test::ExpectTensorEqual<float>(expected, *GetOutput(2)); 109 } 110 { // Output 3 111 Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 3})); 112 test::FillValues<float>(&expected, {9, 10, 11}); 113 test::ExpectTensorEqual<float>(expected, *GetOutput(3)); 114 } 115 } 116 117 TEST_F(DynamicPartitionOpTest, SomeOutputsEmpty) { 118 MakeOp(); 119 120 // Feed and run 121 AddInputFromArray<float>(TensorShape({6}), {0, 13, 2, 39, 4, 17}); 122 AddInputFromArray<int32>(TensorShape({6}), {0, 0, 2, 2, 0, 2}); 123 TF_ASSERT_OK(RunOpKernel()); 124 125 TensorShape empty_one_dim; 126 empty_one_dim.AddDim(0); 127 Tensor expected_empty(allocator(), DT_FLOAT, empty_one_dim); 128 129 // Check the output sizes 130 { // Output 0 131 Tensor expected(allocator(), DT_FLOAT, TensorShape({3})); 132 test::FillValues<float>(&expected, {0, 13, 4}); 133 test::ExpectTensorEqual<float>(expected, *GetOutput(0)); 134 } 135 { // Output 1 136 test::ExpectTensorEqual<float>(expected_empty, *GetOutput(1)); 137 } 138 { // Output 2 139 Tensor expected(allocator(), DT_FLOAT, TensorShape({3})); 140 test::FillValues<float>(&expected, {2, 39, 17}); 141 test::ExpectTensorEqual<float>(expected, *GetOutput(2)); 142 } 143 { // Output 3 144 test::ExpectTensorEqual<float>(expected_empty, *GetOutput(3)); 145 } 146 } 147 148 TEST_F(DynamicPartitionOpTest, Error_IndexOutOfRange) { 149 MakeOp(); 150 151 // Feed and run 152 AddInputFromArray<float>(TensorShape({5, 3}), 153 {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}); 154 AddInputFromArray<int32>(TensorShape({5}), {0, 2, 99, 2, 2}); 155 Status s = RunOpKernel(); 156 EXPECT_TRUE( 157 StringPiece(s.ToString()).contains("partitions[2] = 99 is not in [0, 4)")) 158 << s; 159 } 160 161 Node* DynamicPartitionNode(Graph* g, Node* in0, Node* in1, int num_partitions) { 162 Node* ret; 163 TF_CHECK_OK(NodeBuilder(g->NewName("n"), "DynamicPartition") 164 .Input(in0) 165 .Input(in1) 166 .Attr("num_partitions", num_partitions) 167 .Finalize(g, &ret)); 168 return ret; 169 } 170 171 template <typename T> 172 static Graph* DynamicPartition(int num_partitions, int dim) { 173 Graph* g = new Graph(OpRegistry::Global()); 174 // Always use a 128MB buffer. 175 const int kRows = ((128 << 20) / sizeof(T)) / dim; 176 Tensor data(DataTypeToEnum<T>::value, TensorShape({kRows, dim})); 177 data.flat<T>().setRandom(); 178 179 random::PhiloxRandom philox(301, 17); 180 random::SimplePhilox rnd(&philox); 181 Tensor partitions(DT_INT32, TensorShape({kRows})); 182 for (int i = 0; i < kRows; i++) { 183 partitions.flat<int32>()(i) = rnd.Uniform(num_partitions); 184 } 185 DynamicPartitionNode(g, test::graph::Constant(g, data), 186 test::graph::Constant(g, partitions), num_partitions); 187 return g; 188 } 189 190 #define BM_DYNAMIC_PARTITION(DEVICE, T, num) \ 191 static void BM_##DEVICE##_dynpart_##T##_##num(int iters, int dim) { \ 192 const int64 items = ((128 << 20) / sizeof(T)); \ 193 const int64 tot = static_cast<int64>(iters) * items; \ 194 testing::ItemsProcessed(tot); \ 195 testing::UseRealTime(); \ 196 test::Benchmark(#DEVICE, DynamicPartition<T>(num, dim)).Run(iters); \ 197 } \ 198 BENCHMARK(BM_##DEVICE##_dynpart_##T##_##num)->Arg(1)->Arg(256) 199 200 BM_DYNAMIC_PARTITION(cpu, float, 2); 201 BM_DYNAMIC_PARTITION(cpu, float, 100); 202 BM_DYNAMIC_PARTITION(cpu, double, 2); 203 BM_DYNAMIC_PARTITION(cpu, double, 100); 204 BM_DYNAMIC_PARTITION(cpu, complex64, 2); 205 BM_DYNAMIC_PARTITION(cpu, complex64, 100); 206 207 BM_DYNAMIC_PARTITION(gpu, float, 2); 208 BM_DYNAMIC_PARTITION(gpu, float, 100); 209 BM_DYNAMIC_PARTITION(gpu, double, 2); 210 BM_DYNAMIC_PARTITION(gpu, double, 100); 211 BM_DYNAMIC_PARTITION(gpu, complex64, 2); 212 BM_DYNAMIC_PARTITION(gpu, complex64, 100); 213 214 } // namespace 215 } // namespace tensorflow 216
