Method ResolveAndValidateArguments

tensorflow/compiler/xla/service/service.cc:230–263

}

StatusOr<std::vector<std::vector<const ShapedBuffer*>>>
Service::ResolveAndValidateArguments(
    absl::Span<const GlobalDataHandle* const> arguments,
    absl::Span<se::StreamExecutor* const> stream_executors) const {
  CHECK_EQ(options_.number_of_replicas(), stream_executors.size());
  std::vector<std::vector<const ShapedBuffer*>> replicated_arguments;
  replicated_arguments.resize(options_.number_of_replicas());
  for (size_t i = 0; i < arguments.size(); ++i) {
    auto buffer_status = allocation_tracker_.Resolve(*arguments[i]);
    if (!buffer_status.ok()) {
      return Status(buffer_status.status().code(),
                    StrCat(buffer_status.status().error_message(), ", ",
                           "failed to resolve allocation for parameter ", i));
    }
    auto replicated_buffers = buffer_status.ValueOrDie();
    CHECK_EQ(options_.number_of_replicas(), replicated_buffers.size());
    for (int replica = 0; replica < options_.number_of_replicas(); ++replica) {
      const ShapedBuffer* shaped_buffer = replicated_buffers[replica];
      int replica_device_ordinal = stream_executors[replica]->device_ordinal();
      // Verify allocation is same platform and device as the execution.
      if (shaped_buffer->platform() != execute_backend_->platform() ||
          shaped_buffer->device_ordinal() != replica_device_ordinal) {
        return InvalidArgument(
            "argument %lu is on device %s:%d but computation will be executed "
            "on device %s",
            i, shaped_buffer->platform()->Name(),
            shaped_buffer->device_ordinal(),
            execute_backend_->device_name(replica_device_ordinal));
      }
      replicated_arguments[replica].push_back(shaped_buffer);
    }
  }
  return replicated_arguments;
}

StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig(
    const ProgramShape& program_shape,