Method Compile

tensorflow/compiler/tf2xla/kernels/while_op.cc:294–603

}

void XlaWhileOp::Compile(XlaOpKernelContext* ctx) {
  VLOG(1) << "WhileOp::Compile";

  // Input resources need to be grouped in the end of the body function
  // according to the convention of the XLA bridge.
  OP_REQUIRES_OK(ctx, VerifyResourceArgsGroupedAtEnd(ctx, body_name_attr_));

  std::vector<XlaCompiler::Argument> arguments;
  bool has_uninitialized_vars;
  bool has_tensor_arrays;
  bool has_uninitialized_tensor_lists;
  OP_REQUIRES_OK(ctx, MakeXlaCompilerArgumentsFromInputs(
                          ctx, &arguments, &has_uninitialized_vars,
                          &has_tensor_arrays, &has_uninitialized_tensor_lists));

  xla::XlaBuilder* builder = ctx->builder();
  XlaCompiler* compiler = ctx->compiler();

  // Indices of loop vars which satisfy the following conditions:
  // 1. They are loop invariants.
  // 2. The op inputs at these indices are compile time constants.
  //
  // These compile time consts do not appear as _Args in the cond/body functions
  // and are replaced by kConstant nodes instead. As as result, the compiled
  // body function does not have matching input and output shape. We fix this
  // by rewriting the body computation (see body_wrapper below) to output
  // just the non compile-time-const values and later pad up the while output
  // with the const args.
  std::vector<bool> compile_time_const_arg_indices(ctx->num_inputs());
  int num_compile_time_const_args = 0;
  if (propagate_compile_time_consts_) {
    OP_REQUIRES_OK(ctx, ConvertLoopInvariantsToConst(
                            ctx, body_name_attr_, &arguments,
                            &compile_time_const_arg_indices,
                            &num_compile_time_const_args, compiler->client()));
  }

  VLOG(1) << "Compiling body";

  // All resource that are inputs to the loop's body must also be
  // present as loop body outputs; the signature of the loop's input and
  // output must match. We ensure this by asking the compiler to include the
  // current values of all resources, even if they haven't been updated by the
  // computation. We must also ask the compiler to keep compile-time constant
  // outputs as part of the generated computation, for the same reason.
  // TODO(phawkins): consider adding loop-invariant inputs to XLA's While()
  // operator.
  XlaCompiler::CompileOptions body_options;
  body_options.use_tuple_arg = true;
  body_options.return_updated_values_for_all_resources = true;
  body_options.resolve_compile_time_constants = false;
  body_options.is_entry_computation = false;
  body_options.add_token_input_output = has_token_input_output_;
  XlaCompiler::CompilationResult body;
  OP_REQUIRES_OK(ctx, compiler->CompileFunction(body_options, body_name_attr_,
                                                arguments, &body));

  // We must use a static shape for parameters to an XLA compilation. However,