Function BM_Helper

tensorflow/core/distributed_runtime/rpcbench_test.cc:169–214

TODO: Support sharding and depth.

// TODO: Support sharding and depth.
static void BM_Helper(int iters, int width, int num_stages, int tensor_size,
                      bool use_multiple_devices) {
  testing::StopTiming();
  const Cluster* cluster = GetCluster();

  // Creates a session.
  std::unique_ptr<Session> session(NewSession(cluster->options));
  GraphDef def = CreateGraphDef(num_stages, width, tensor_size,
                                use_multiple_devices, cluster);
  graph::SetDefaultDevice(cluster->devices[0].name(), &def);

  TF_CHECK_OK(session->Create(def));

  // Randomly initialize the input.
  Tensor x(DT_FLOAT, TensorShape({tensor_size, 1}));

  testing::SetLabel(
      strings::StrCat(def.node_size(), " nodes; ",
                      use_multiple_devices ? "Multi device" : "Single device",
                      "; tensor bytes/send: ", tensor_size * sizeof(float)));

  std::vector<Tensor> outputs;

  // Do a few warmup iterations.
  for (int i = 0; i < 3; i++) {
    outputs.clear();
    TF_CHECK_OK(session->Run({{"x", x}}, {"y:0"}, {}, &outputs));
    CHECK_EQ(size_t{1}, outputs.size());

    if (i == 0) {
      // Print out x, and y.
      const Tensor& y = outputs[0];
      VLOG(1) << DebugString(x, y, tensor_size);
    }
  }

  // Iterations.
  testing::StartTiming();
  for (int i = 0; i < iters; i++) {
    outputs.clear();
    TF_CHECK_OK(session->Run({{"x", x}}, {"y:0"}, {}, &outputs));
    CHECK_EQ(size_t{1}, outputs.size());
  }
  testing::StopTiming();
  TF_CHECK_OK(session->Close());
}
static void BM_ShardedProgram(int iters, int width, int num_stages) {
  BM_Helper(iters, width, num_stages, 2 /*tensor_size*/, true /*multi-device*/);
}