Method Compile

tensorflow/compiler/tf2xla/graph_compiler.cc:116–199

namespace

}
}  // namespace
Status GraphCompiler::Compile() {
  // Check that the graph has no illegal cycles.
  TF_RETURN_IF_ERROR(graph::ValidateGraphHasNoCycle(*graph_));
  // Maintain a mapping from node id to node outputs.
  using NodeOutputs = std::vector<TensorValue>;
  std::vector<NodeOutputs> output_registry(graph_->num_node_ids());
  auto output_registry_cleanup = gtl::MakeCleanup([&output_registry] {
    for (const NodeOutputs& outputs : output_registry) {
      for (const TensorValue& value : outputs) {
        CHECK(!value.is_ref());
        delete value.tensor;
      }
    }
  });

  // XLA requires determinism, generate a stable ordering from DFS.
  std::vector<Node*> topo_sorted_nodes;
  GetReversePostOrder(*graph_, &topo_sorted_nodes,
                      /*stable_comparator=*/NodeComparatorName());

  OpKernelContext::Params params;
  PartiallySetupParams(&params);

  for (Node* n : topo_sorted_nodes) {
    OpKernel* op_kernel_raw = nullptr;
    // The kernel is not actually run for functional ops, we just need it
    // for metadata.
    Status s = flib_->CreateKernel(n->def(), &op_kernel_raw);
    // Transfer ownership of the kernel to a local smart pointer.
    std::unique_ptr<OpKernel> op_kernel(op_kernel_raw);

    if (!s.ok()) {
      s = AttachDef(s, *n);
      LOG(ERROR) << "Executor failed to create kernel. " << s;
      return s;
    }

    TF_RET_CHECK(!n->IsRecv() && !n->IsSend() && !n->IsSwitch())
        << "Not supported node: " << n->DebugString();
    params.op_kernel = op_kernel.get();
    absl::InlinedVector<AllocatorAttributes, 4> output_attr(n->num_outputs());
    params.output_attr_array = output_attr.data();

    // tensor_inputs_ is a buffer reused across graph traversal. We clean up and
    // reinitialize the buffer before we visit a new node.
    tensor_inputs_.clear();
    tensor_inputs_.resize(n->num_inputs());

    // Set up inputs from outputs of previous nodes.
    for (auto* e : n->in_edges()) {
      if (e->IsControlEdge()) continue;
      const Node* src = e->src();
      TF_RET_CHECK(src->id() < output_registry.size());
      const NodeOutputs& src_outputs = output_registry[src->id()];

      tensor_inputs_.at(e->dst_input()) = src_outputs.at(e->src_output());
    }