Method Build

tensorflow/compiler/xla/service/call_graph.cc:279–318

static */

/* static */
std::unique_ptr<CallGraph> CallGraph::Build(const HloModule* module) {
  // Constructor for CallGraph is private so absl::make_unique can't be used.
  auto call_graph = absl::WrapUnique<CallGraph>(new CallGraph(module));

  VLOG(3) << "Building call graph for:";
  XLA_VLOG_LINES(3, module->ToString());

  // Construct nodes of the call graph and populate the callsites.
  for (HloComputation* computation : module->computations()) {
    auto it_added = call_graph->node_indices_.insert(
        {computation, call_graph->nodes_.size()});
    // All computations should be unique, so the computation should not already
    // exist in the map.
    CHECK(it_added.second);
    call_graph->nodes_.emplace_back(computation);

    // Add all callsites in this computation.
    for (HloInstruction* instruction : computation->instructions()) {
      call_graph->nodes_.back().AddCallSiteForInstruction(instruction);
    }
  }

  // Add caller callsites to each node.
  for (const HloComputation* computation : module->computations()) {
    for (const CallSite& callsite :
         call_graph->GetNode(computation).callsites()) {
      for (auto* callee : callsite.called_computations()) {
        // Add caller callsites.
        call_graph->GetNode(callee).AddCallerCallSite(callsite);
      }
    }
  }

  call_graph->SetCallContexts();
  call_graph->SetNodeDepths();

  XLA_VLOG_LINES(2, call_graph->ToString());

  return call_graph;
}

Status CallGraph::VisitNodesInternal(
    const VisitorFunction& visitor_func, const CallGraphNode& node,