Test TakeSubTree with different depths (depth of ShapeTree) and fan-outs (cardinality of each non-leaf node's children).
| 175 | // Test TakeSubTree with different depths (depth of ShapeTree) and fan-outs |
| 176 | // (cardinality of each non-leaf node's children). |
| 177 | void BM_TakeSubTree(int iters, int depth, int fan_out) { |
| 178 | tensorflow::testing::StopTiming(); |
| 179 | TestAllocator allocator; |
| 180 | xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {32, 64, 128}); |
| 181 | for (int i = 0; i < depth; ++i) { |
| 182 | std::vector<xla::Shape> shapes(fan_out, shape); |
| 183 | shape = xla::ShapeUtil::MakeTupleShape(shapes); |
| 184 | } |
| 185 | xla::ScopedShapedBuffer shaped_buffer(shape, shape, /*allocator=*/&allocator, |
| 186 | /*device_ordinal=*/0); |
| 187 | tensorflow::testing::StartTiming(); |
| 188 | for (int i = 0; i < iters; ++i) { |
| 189 | // Extract a buffer from approximately the middle of the first level of the |
| 190 | // tree. |
| 191 | (void)shaped_buffer.TakeSubTree(/*index=*/{fan_out / 2}).release(); |
| 192 | } |
| 193 | tensorflow::testing::StopTiming(); |
| 194 | } |
| 195 | |
| 196 | BENCHMARK(BM_TakeSubTree) |
| 197 | ->ArgPair(1, 4) |
nothing calls this directly
no test coverage detected