Method AddNode

tensorflow/core/profiler/internal/tfprof_code.cc:393–434

}  // namespace

void TFCode::AddNode(TFGraphNode* node) {
  if (!node->call_stack() || node->call_stack()->traces().empty()) {
    return;
  }
  // We infer the forward operation name from gradient op name. So, we can
  // map gradient op traces to forward op traces.
  // E.g. gradient node of 'inp_1/Conv2D' would be 'gradients/inp_1/Conv2D_grad.
  string forward_name;
  if (IsGradNode(node->name(), &forward_name)) {
    auto grad_nodes_it = grad_nodes_.find(forward_name);
    if (grad_nodes_it != grad_nodes_.end()) {
      grad_nodes_it->second.push_back(node);
    } else {
      grad_nodes_.insert(
          std::pair<string, std::vector<TFGraphNode*>>(forward_name, {node}));
    }
    return;
  } else {
    forward_nodes_[node->name()] = node;
  }

  if (!root_) {
    graph_root_.reset(new TFMultiGraphNode(kTFProfRoot));
    root_.reset(new CodeNode(graph_root_.get(), nullptr, ""));
  }

  CodeNode* pre_code_node = root_.get();
  // TODO(xpan): Consider to release CodeDef after TFCode is built. It
  // takes a lot of memory.
  std::set<string> traces;
  for (int i = 0; i < node->call_stack()->traces().size(); ++i) {
    // Unlike op name, which is globally unique, trace name is only unique
    // w.r.t. it's parent.
    const string& trace = GetTraceString(node->call_stack()->traces().at(i));
    traces.insert(trace);
    pre_code_node = pre_code_node->AddChildren(
        trace, &node->call_stack()->traces().at(i), "");
    if (i == node->call_stack()->traces().size() - 1) {
      pre_code_node->node->AddGraphNode(node);
    }
  }
}

void TFCode::Build() {
  int64 unaccounted_nodes = 0;