| 311 | |
| 312 | template <typename T> |
| 313 | static Graph* DynamicPartition(int num_partitions, int dim) { |
| 314 | Graph* g = new Graph(OpRegistry::Global()); |
| 315 | // Always use a 128MB buffer. |
| 316 | const int kRows = ((128 << 20) / sizeof(T)) / dim; |
| 317 | Tensor data(DataTypeToEnum<T>::value, TensorShape({kRows, dim})); |
| 318 | data.flat<T>().setRandom(); |
| 319 | |
| 320 | random::PhiloxRandom philox(301, 17); |
| 321 | random::SimplePhilox rnd(&philox); |
| 322 | Tensor partitions(DT_INT32, TensorShape({kRows})); |
| 323 | for (int i = 0; i < kRows; i++) { |
| 324 | partitions.flat<int32>()(i) = rnd.Uniform(num_partitions); |
| 325 | } |
| 326 | DynamicPartitionNode(g, test::graph::Constant(g, data), |
| 327 | test::graph::Constant(g, partitions), num_partitions); |
| 328 | return g; |
| 329 | } |
| 330 | |
| 331 | #define BM_DYNAMIC_PARTITION(DEVICE, T, num) \ |
| 332 | static void BM_##DEVICE##_dynpart_##T##_##num(int iters, int dim) { \ |
no test coverage detected