Function BuildLoopBody

tensorflow/compiler/tf2xla/functionalize_while.cc:157–211

Builds a graph for the loop body.

// Builds a graph for the loop body.
Status BuildLoopBody(const Graph& graph, WhileLoopFrame* frame,
                     DataTypeVector* arg_types,
                     std::unique_ptr<Graph>* body_output) {
  VLOG(2) << "Building loop body for " << frame->name;
  *body_output = absl::make_unique<Graph>(graph.op_registry());
  Graph* output = body_output->get();

  // Map from nodes in the original graph to the condition graph.
  std::vector<Node*> node_map(graph.num_node_ids(), nullptr);
  std::vector<bool> squash_src_outputs(graph.num_node_ids(), false);

  // Build one _Arg node for each Enter node.
  std::vector<Node*> next_iterations;
  next_iterations.reserve(frame->args.size());
  arg_types->reserve(frame->args.size());
  for (int i = 0; i < frame->args.size(); ++i) {
    const WhileLoopArg& arg = frame->args[i];

    DataType dtype = arg.enter->input_type(0);
    arg_types->push_back(dtype);

    TF_ASSIGN_OR_RETURN(Node * arg_node, BuildArgNode(output, dtype, i));
    TF_ASSIGN_OR_RETURN(Node * retval_node, BuildRetvalNode(output, dtype, i));
    if (arg.is_loop_invariant) {
      // Argument is loop-invariant. Forward it from the Arg to the Retval.
      node_map[arg.enter->id()] = arg_node;
      output->AddEdge(arg_node, 0, retval_node, 0);
    } else {
      // Argument is loop-varying.
      if (dtype == DT_RESOURCE) {
        // DT_RESOURCE arguments should always be loop-invariant in the graphs
        // generated from TF.
        return errors::Unimplemented("Loop-varying DT_RESOURCE Enter node ",
                                     arg.enter->name(), " is currently not",
                                     " supported.");
      }
      node_map[arg.switch_node->id()] = arg_node;
      // The Switch node has two outputs, but _Arg only has one. This tells
      // the CopySubgraph function to rewrite the output number of edges from
      // the _Arg node to be 0 rather than copying the output number from the
      // Switch node.
      squash_src_outputs[arg.switch_node->id()] = true;
      node_map[arg.next_iteration->id()] = retval_node;
      next_iterations.push_back(arg.next_iteration);
    }
  }

  // Performs a reverse DFS, copying nodes and edges to the output graph.
  // The _Arg and _Retval nodes were added unconditionally above, so we are
  // guaranteed to get the correct function signature.
  TF_RETURN_IF_ERROR(CopySubgraph(graph, frame, std::move(next_iterations),
                                  squash_src_outputs, &node_map, output));

  return Status::OK();
}

// Copy the FunctionDef of given function from lookup_library to library, if
// it can be found in lookup_library but is missing from library.