Function BM_ForEach

tensorflow/compiler/xla/shape_tree_test.cc:574–589

}

void BM_ForEach(int iters, int depth, int fan_out) {
  tensorflow::testing::StopTiming();
  Shape shape = ShapeUtil::MakeShape(F32, {32, 64, 128});
  for (int i = 0; i < depth; ++i) {
    std::vector<xla::Shape> shapes(fan_out, shape);
    shape = ShapeUtil::MakeTupleShape(shapes);
  }
  tensorflow::testing::StartTiming();

  ShapeTree<int> shape_tree(shape);
  for (int i = 0; i < iters; ++i) {
    shape_tree.ForEachMutableElement([](const ShapeIndex& index, int* data) {
      tensorflow::testing::DoNotOptimize(index);
    });
  }
}

void BM_Iterate(int iters, int depth, int fan_out) {
  tensorflow::testing::StopTiming();