Method Benchmark

tensorflow/core/common_runtime/kernel_benchmark_testlib.cc:45–105

TODO(hongm): Convert `g` and `init` to using std::unique_ptr.

// TODO(hongm): Convert `g` and `init` to using std::unique_ptr.
Benchmark::Benchmark(const string& device, Graph* g,
                     const SessionOptions* options, Graph* init,
                     Rendezvous* rendez, const char* executor_type) {
  SessionOptions default_options;
  if (!options) {
    options = &default_options;
  }

  testing::StopTiming();
  string t = absl::AsciiStrToUpper(device);
  // Allow NewDevice to allocate a new threadpool with different number of
  // threads for each new benchmark.
  LocalDevice::set_use_global_threadpool(false);
  device_ =
      DeviceFactory::NewDevice(t, *options, "/job:localhost/replica:0/task:0");
  CHECK(device_) << "Could not create a " << device << " device";

  pool_ =
      new thread::ThreadPool(options->env, "blocking", port::MaxParallelism());

  auto runner = [this](std::function<void()> closure) {
    pool_->Schedule(closure);
  };
  auto cost_runner = [this](std::function<void()> closure, int64 cost) {
    pool_->CostSchedule(closure, cost);
  };

  if (rendez == nullptr) {
    rendez_ = NewLocalRendezvous();
  } else {
    rendez_ = rendez;
  }

  const int graph_def_version = g->versions().producer();

  LocalExecutorParams params;
  params.device = device_.get();
  params.function_library = nullptr;
  params.create_kernel = [this, graph_def_version](const NodeDef& ndef,
                                                   OpKernel** kernel) {
    return CreateNonCachedKernel(device_.get(), nullptr, ndef,
                                 graph_def_version, kernel);
  };
  params.delete_kernel = [](OpKernel* kernel) {
    DeleteNonCachedKernel(kernel);
  };

  if (init) {
    std::unique_ptr<Executor> init_exec;
    TF_CHECK_OK(NewExecutor(executor_type, params, std::unique_ptr<Graph>(init),
                            &init_exec));
    Executor::Args args;
    args.rendezvous = rendez_;
    args.runner = runner;
    args.cost_runner = cost_runner;
    TF_CHECK_OK(init_exec->Run(args));
  }