Function BM_SubgraphHelper

tensorflow/core/graph/subgraph_test.cc:345–384

REGISTER_OP("Op").Input("i: float").Output("o: float");

static void BM_SubgraphHelper(int iters, int num_nodes,
                              bool use_function_convention) {
  DeviceAttributes device_info;
  device_info.set_name("/job:a/replica:0/task:0/cpu:0");
  device_info.set_device_type(DeviceType(DEVICE_CPU).type());
  device_info.set_incarnation(0);

  testing::StopTiming();
  Graph g(OpRegistry::Global());
  {  // Scope for temporary variables used to construct g.
    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
    Node* last_node = nullptr;
    for (int i = 0; i < num_nodes; i++) {
      string name = strings::StrCat("N", i);
      if (i > 0) {
        last_node = ops::UnaryOp("Op", last_node, b.opts().WithName(name));
      } else {
        last_node = ops::SourceOp("In", b.opts().WithName(name));
      }
    }
    TF_CHECK_OK(GraphDefBuilderToGraph(b, &g));
  }

  std::vector<string> fed;
  if (num_nodes > 1000) {
    fed.push_back(strings::StrCat("N", num_nodes - 1000));
  }
  std::vector<string> fetch;
  std::vector<string> targets = {strings::StrCat("N", num_nodes - 1)};
  testing::StartTiming();
  while (--iters > 0) {
    Graph* subgraph = new Graph(OpRegistry::Global());
    CopyGraph(g, subgraph);
    subgraph::RewriteGraphMetadata metadata;
    TF_CHECK_OK(subgraph::RewriteGraphForExecution(
        subgraph, fed, fetch, targets, device_info, use_function_convention,
        &metadata));
    delete subgraph;
  }
}

static void BM_Subgraph(int iters, int num_nodes) {
  BM_SubgraphHelper(iters, num_nodes, false /* use_function_convention */);