Function TEST_F

tensorflow/cc/profiler/profiler_test.cc:95–174

}

TEST_F(ProfilerTest, Basics) {
  SessionOptions options;
  options.config.set_allow_soft_placement(true);
  std::unique_ptr<Session> session(NewSession(options));
  GraphDef def = CreateGraphDef();
  if (options.target.empty()) {
    graph::SetDefaultDevice("/gpu:0", &def);
  }

  TF_CHECK_OK(session->Create(def));

  Tensor x(DT_FLOAT, TensorShape({2, 1}));
  auto x_flat = x.flat<float>();
  x_flat.setRandom();
  Eigen::Tensor<float, 0, Eigen::RowMajor> inv_norm =
      x_flat.square().sum().sqrt().inverse();
  x_flat = x_flat * inv_norm();

  std::vector<Tensor> outputs;
  RunOptions run_options;
  run_options.set_trace_level(RunOptions::FULL_TRACE);
  RunMetadata run_metadata;
  outputs.clear();

  Profiler profiler(def);
  for (int i = 0; i < 2; ++i) {
    TF_CHECK_OK(session->Run(run_options, {{"x", x}}, {"y:0", "y_normalized:0"},
                             {}, &outputs, &run_metadata));
    profiler.AddStep(i, run_metadata);
    CHECK_EQ(size_t{2}, outputs.size());
  }

  std::vector<DeviceAttributes> resp;
  TF_CHECK_OK(session->ListDevices(&resp));
  bool has_gpu = false;
  for (const auto& dev : resp) {
    if (dev.device_type() == "GPU") {
      has_gpu = true;
    }
  }

  GraphNodeProto ret = profiler.ProfileNameScope(Default());
  const GraphNodeProto* matmul = ExtractNode(ret, "y");
  EXPECT_TRUE(matmul);
  EXPECT_GT(matmul->exec_micros(), 0);
  if (has_gpu) {
    EXPECT_GT(matmul->accelerator_exec_micros(), 0);
  } else {
    EXPECT_EQ(matmul->accelerator_exec_micros(), 0);
  }
  const GraphNodeProto* square = ExtractNode(ret, "Square");
  EXPECT_TRUE(square);
  EXPECT_GT(square->exec_micros(), 0);
  if (has_gpu) {
    EXPECT_GT(square->accelerator_exec_micros(), 0);
  } else {
    EXPECT_EQ(square->accelerator_exec_micros(), 0);
  }