Method operator()

tensorflow/compiler/jit/extract_outside_compilation_pass.cc:2026–2128

}  // namespace

Status RewriteOutsideCompilationSubgraphFn::operator()(
    const std::vector<OutputTensor>& arg_source_tensors,
    std::unique_ptr<Graph>* graph, std::vector<int>* input_permutation,
    std::vector<int>* output_permutation, NodeDef* node_def) {
  string old_name = node_def->op();
  string new_name =
      absl::StrCat(xla_cluster_name_, "_", new_function_name_, "_", old_name);
  node_def->set_op(new_name);
  node_def->set_name(new_name);

  // Later we will run PruneForReverseReachability(), so make sure all original
  // nodes are reachable from sink node and won't be removed.
  FixupSourceAndSinkEdges(graph->get());

  // Step 1: create a key placeholder node.
  TF_ASSIGN_OR_RETURN(
      Node * key_placeholder,
      AddHostComputeKeyPlaceholder(xla_cluster_name_, graph->get()));

  // Step 2: build RecvAtHost node, and replace all _Arg nodes with it.
  std::vector<DataType> recv_at_host_dtypes;
  TF_ASSIGN_OR_RETURN(
      Node * recv_at_host_node,
      ReplaceArgNodesWithRecvAtHostNode(graph->get(), new_name,
                                        &recv_at_host_dtypes, key_placeholder));

  // Step 3: build SendFromHost node, and replace all _Retval nodes with it.
  std::vector<DataType> send_from_host_dtypes;
  TF_ASSIGN_OR_RETURN(
      Node * send_from_host_node,
      ReplaceRetNodesWithSendFromHostNode(
          graph->get(), new_name, &send_from_host_dtypes, key_placeholder));

  // Step 4: add XLA cluster and outside compilation attr.
  for (Node* n : (*graph)->nodes()) {
    if (IsKeyPlaceholderNode(*n)) {
      continue;
    }

    n->AddAttr(xla_cluster_attr_name_, xla_cluster_name_);
    n->AddAttr(outside_compilation_attr_name_, old_name);
  }

  // Check whether we have all input shapes for XlaSendFromHost. If we do, we
  // will set `shapes` attr for the call node; otherwise we will save the
  // shape inference graph and set `shape_inference_graph` for the call node.
  absl::optional<std::vector<PartialTensorShape>> shapes =
      GetInferredInputShapes(send_from_host_dtypes.size(), send_from_host_node);
  for (Node* n : (*graph)->nodes()) {
    n->ClearAttr(kXlaInferredShapesAttrName);
  }

  // Step 5: add control edges for originally XLA <-> outside compilation
  // control edges.
  for (Node* n : (*graph)->nodes()) {
    if (HasNodeAttr(n->def(), kXlaConnectedToXlaComputationAttrName)) {
      (*graph)->AddControlEdge(n, send_from_host_node);
      n->ClearAttr(kXlaConnectedToXlaComputationAttrName);