Function PropagateConstIntoWhileNode

tensorflow/compiler/tf2xla/tf2xla_util.cc:247–301

For a "While" node in graph `g`, if it has Const node inputs, rewrite its cond/body function to replace _Arg nodes with those Const inputs.

// For a "While" node in graph `g`, if it has Const node inputs, rewrite its
// cond/body function to replace _Arg nodes with those Const inputs.
Status PropagateConstIntoWhileNode(Graph* g, Node* while_node,
                                   const FunctionLibraryDefinition* lookup_fld,
                                   FunctionLibraryDefinition* fld) {
  // For "While" node, we should only replace _Arg nodes which are loop
  // invariants. For such _Arg nodes, the return value's input will come
  // directly from the corresponding arg.
  std::unordered_map<int, const Node*> const_input_index_to_node;
  NameAttrList body_attr;
  TF_RETURN_IF_ERROR(GetNodeAttr(while_node->def(), "body", &body_attr));
  const FunctionDef* body_func = lookup_fld->Find(body_attr.name());
  if (!body_func) {
    return errors::Internal("Cannot find body function ", body_attr.name(),
                            " for While node ", while_node->name());
  }
  for (int i = 0; i < while_node->num_inputs(); i++) {
    const Node* input_node;
    TF_RETURN_IF_ERROR(while_node->input_node(i, &input_node));
    if (input_node->type_string() != "Const") {
      continue;
    }

    // Check if i-th retval's input comes from i-th arg directly.
    // For resource variable input of While nodes, TF2XLA convention is to place
    // them at the end of all inputs (after all data inputs), and *not* return
    // them. So number of While node inputs might be larger than number of its
    // outputs.
    if (i >= body_func->signature().output_arg_size()) {
      continue;
    }
    const OpDef_ArgDef& output_arg = body_func->signature().output_arg(i);
    auto output_arg_input = body_func->ret().find(output_arg.name());
    if (output_arg_input == body_func->ret().end()) {
      return errors::Internal("Cannot find input for output arg ",
                              output_arg.name(), " in function ",
                              body_attr.name());
    }
    const OpDef_ArgDef& input_arg = body_func->signature().input_arg(i);
    if (output_arg_input->second != input_arg.name()) {
      continue;
    }

    const_input_index_to_node[i] = input_node;
  }
  if (const_input_index_to_node.empty()) {
    return Status::OK();
  }

  // Rewrite "cond" and "body" function, replace usage of those _Arg nodes with
  // corresponding const node.
  for (const auto& attr_name : std::vector<string>{"cond", "body"}) {
    TF_RETURN_IF_ERROR(PropagateConstIntoFuncAttr(
        while_node, attr_name, const_input_index_to_node, lookup_fld, fld));
  }
  return Status::OK();
}

}  // namespace