Function ReplayComputation

tensorflow/compiler/xla/tools/replay_computation.cc:221–344

Invokes the given computation passing arbitrary data for every (unbound) parameter if use_fake_data, Otherwise use recorded data if available. Similarly, infeeds fake data of shape fake_infeed_shape if it is provided. If generate_fake_infeed is true, the required infeed shape is derived from the computation and then used to provide a fake infeed shape. If neither generate_fake_infeed is true nor

// If neither generate_fake_infeed is true nor a fake_infeed_shape is provided,
// no infeed is performed.
StatusOr<Literal> ReplayComputation(const HloSnapshot& module,
                                    LocalExecutable* executable,
                                    LocalClient* client, const Options& opts) {
  XlaComputation computation(module.hlo().hlo_module());

  // Build the `argument_ptrs` vector, which contains ShapedBuffer*s to our
  // arguments.  This is a bit involved, because we may have to convert from
  // GlobalData to ShapedBuffer*, and we have to manage the lifetime of all our
  // objects.
  std::vector<ScopedShapedBuffer> scoped_shaped_buffer_arguments;
  std::vector<std::unique_ptr<GlobalData>> global_data_arguments;
  std::vector<const ShapedBuffer*> argument_ptrs;
  if (opts.use_fake_data) {
    // Run fake computations with debug options ignoring XLA_FLAGS.  Users very
    // likely want XLA_FLAGS only to apply to the "real" computation being run,
    // not to the fake computations we use for generating arguments.
    auto debug_opts_flags = GetDebugOptionsFromFlags();
    auto debug_opts = DefaultDebugOptionsIgnoringFlags();

    // ptxas can be called during the generation of fake data.
    // As it is cached, we want it to not ignore this flag.
    debug_opts.set_xla_gpu_asm_extra_flags(debug_opts_flags.xla_gpu_asm_extra_flags());

    global_data_arguments =
        MakeFakeArgumentsOrDie(computation, client, &debug_opts);
    for (const auto& data : global_data_arguments) {
      argument_ptrs.push_back(
          client->GlobalDataToShapedBuffer(data->handle(), /*device_ordinal=*/0)
              .ValueOrDie());
    }
  } else {  // use recorded data if available
    for (const auto& proto : module.arguments()) {
      TF_ASSIGN_OR_RETURN(Literal literal, Literal::CreateFromProto(proto));
      TF_ASSIGN_OR_RETURN(
          ScopedShapedBuffer data,
          client->LiteralToShapedBuffer(literal, /*device_ordinal=*/0));
      scoped_shaped_buffer_arguments.push_back(std::move(data));
    }
    for (const auto& argument : scoped_shaped_buffer_arguments) {
      argument_ptrs.push_back(&argument);
    }
  }

  if (absl::optional<Shape> infeed_shape = GetXfeedShape(
          /*is_infeed=*/true, computation.proto(), opts)) {
    auto infeed_data = std::make_shared<Literal>(
        std::move(MakeFakeLiteral(*infeed_shape)).ValueOrDie());
    xla::gpu::GetOrCreateInfeedManager()
        ->RegisterBeforeGetNextDestinationCallback([infeed_data, client] {
          TF_CHECK_OK(client->TransferToInfeed(*infeed_data));
        });
  }

  absl::optional<tensorflow::thread::ThreadPool> outfeed_thread_pool;
  if (absl::optional<Shape> outfeed_shape = GetXfeedShape(
          /*is_infeed=*/false, computation.proto(), opts)) {
    // For each an outfeed that runs, enqueue a task that will consume it.  We
    // need a thread pool because the act of running an outfeed blocks on there