Function GetLoopInvariants

tensorflow/compiler/tf2xla/kernels/while_op.cc:128–140

Populates loop invariant indices to true in `loop_invariants`.

// Populates loop invariant indices to true in `loop_invariants`.
void GetLoopInvariants(XlaOpKernelContext* ctx,
                       const NameAttrList& body_name_attr,
                       std::vector<bool>* const loop_invariants) {
  const FunctionBody* body;
  OP_REQUIRES_OK(ctx, ctx->compiler()->FindFunctionBody(body_name_attr, &body));
  for (int i = 0; i < body->ret_nodes.size(); i++) {
    const Node* arg = body->arg_nodes[i];
    const Node* ret = body->ret_nodes[i];
    const Node* ret_input_0;
    OP_REQUIRES_OK(ctx, ret->input_node(0, &ret_input_0));
    (*loop_invariants)[i] = ret_input_0->id() == arg->id();
  }
}

// Converts entries in `args` which are loop invariants and have compile
// time constant inputs to constants so that they can be propagated in the loop