| 27 | namespace tensorflow { |
| 28 | |
| 29 | static Graph* MakeGraph(int split_dim, int num_split, |
| 30 | std::initializer_list<int64> chunk_size) { |
| 31 | Graph* g = new Graph(OpRegistry::Global()); |
| 32 | TensorShape in_shape(chunk_size); |
| 33 | in_shape.set_dim(split_dim, in_shape.dim_size(split_dim) * num_split); |
| 34 | Tensor in(DataTypeToEnum<float>::value, in_shape); |
| 35 | in.flat<float>().setRandom(); |
| 36 | Tensor split_dim_tensor = test::AsScalar<int32>(split_dim); |
| 37 | Node* split; |
| 38 | TF_CHECK_OK(NodeBuilder(g->NewName("split"), "Split") |
| 39 | .Input(test::graph::Constant(g, split_dim_tensor)) |
| 40 | .Input(test::graph::Constant(g, in)) |
| 41 | .Attr("num_split", num_split) |
| 42 | .Finalize(g, &split)); |
| 43 | return g; |
| 44 | } |
| 45 | |
| 46 | #define BM_SPLIT_1D(num_split, chunk_size) \ |
| 47 | static void BM_Split_1d_##num_split##_##chunk_size(int iters) { \ |