

#include "tensorflow/compiler/xla/service/gpu/gpu_executable.h"

#include <set>
#include <utility>
#include <vector>

#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/ptr_util.h"
#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/logical_buffer.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/service/transfer_manager.h"
#include "tensorflow/compiler/xla/shape_tree.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"

namespace se = ::perftools::gputools;

namespace xla {
namespace gpu {
namespace {

// A helper class for profiling HLO in the course of GPU program execution.
// All of the profiling is guarded internally, to avoid the caller needing to
// have lots of conditionals sprinkled around.
class HloExecutionProfiler {
 public:
  // If profiling is enabled, start an execution timer running.
  explicit HloExecutionProfiler(bool do_profile, HloExecutionProfile* profile,
                                se::Stream* stream,
                                const HloComputation* computation)
      : do_profile_(do_profile),
        profile_(profile),
        stream_(stream),
        computation_(computation) {
    if (do_profile_) {
      clock_rate_ghz_ =
          stream->parent()->GetDeviceDescription().clock_rate_ghz();
      execution_timer_.reset(new se::Timer(stream->parent()));
      per_op_timer_.reset(new se::Timer(stream->parent()));
      stream->InitTimer(execution_timer_.get())
          .ThenStartTimer(execution_timer_.get());
      stream->InitTimer(per_op_timer_.get());
    }
  }

  // If profiling is enabled, sets the total cycle count on the profile from the
  // execution timer.
  ~HloExecutionProfiler() {
    if (do_profile_) {
      stream_->ThenStopTimer(execution_timer_.get());
      stream_->BlockHostUntilDone();
      profile_->set_total_cycles_executed(
          *computation_, execution_timer_->Nanoseconds() * clock_rate_ghz_);
    }
  }

  // If profiling is enabled, starts the per-operation timer.
  void StartOperation() {
    if (do_profile_) {
      stream_->ThenStartTimer(per_op_timer_.get());
    }
  }

  // If profiling is enabled, stops the per-operation timer and records the time
  // that the hlo_instruction took to execute in the profile.
  void FinishOperation(const HloInstruction* hlo_instruction) {
    if (do_profile_) {
      stream_->ThenStopTimer(per_op_timer_.get());
      stream_->BlockHostUntilDone();
      profile_->SetCyclesTakenBy(
          hlo_instruction, per_op_timer_->Nanoseconds() * clock_rate_ghz_);
    }
  }

 private:
  const bool do_profile_;
  double clock_rate_ghz_;
  HloExecutionProfile* profile_;
  se::Stream* stream_;
  const HloComputation* computation_;
  std::unique_ptr<se::Timer> execution_timer_;
  std::unique_ptr<se::Timer> per_op_timer_;
};

}  // namespace

// Implementation note: HLO profiling is always enabled for GPU executables,
// since we can use timers around thunks.
GpuExecutable::GpuExecutable(
    const string& ptx, const std::vector<uint8>& cubin,
    std::pair<int, int> compute_capability,
    std::unique_ptr<const ThunkSchedule> thunk_schedule,
    std::unique_ptr<const HloModule> hlo_module,
    std::unique_ptr<const BufferAssignment> assignment,
    HloCostAnalysis::ShapeSizeFunction shape_size_function)
    : Executable(std::move(hlo_module)),
      ptx_(ptx),
      cubin_(cubin),
      compute_capability_(compute_capability),
      thunk_schedule_(std::move(thunk_schedule)),
      assignment_(std::move(assignment)),
      shape_size_function_(std::move(shape_size_function)) {}

Status GpuExecutable::ExecuteThunks(
    const ServiceExecutableRunOptions* run_options,
    const BufferAllocations& buffer_allocations, bool block_host_until_done,
    HloExecutionProfile* hlo_execution_profile) {
  se::Stream* main_stream = run_options->stream();

  std::pair<int, int> stream_compute_compatibility;
  main_stream->parent()->GetDeviceDescription().cuda_compute_capability(
      &stream_compute_compatibility.first,
      &stream_compute_compatibility.second);
  TF_RET_CHECK(stream_compute_compatibility == compute_capability_)
      << "Compute capability mismatch; expected {" << compute_capability_.first
      << ", " << compute_capability_.second << "}, but was {"
      << stream_compute_compatibility.first << ", "
      << stream_compute_compatibility.second << "}";

  bool do_profile = hlo_execution_profile != nullptr;
  if (do_profile) {
    LOG(WARNING) << "PROFILING: profiling is enabled";
  }
  HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream,
                                hlo_module_->entry_computation());

  // Stream 0 indicates `main_stream` and substreams start from stream 1.
  std::vector<Pool<se::Stream>::SmartPtr> sub_streams;
  while (sub_streams.size() + 1 < thunk_schedule_->StreamCount()) {
    sub_streams.emplace_back();
    TF_ASSIGN_OR_RETURN(
        sub_streams.back(),
        run_options->BorrowStream(main_stream->parent()->device_ordinal()));
  }

  std::map<const Thunk*, std::unique_ptr<se::Event>> thunk_to_finish_event;
  for (Thunk* thunk : thunk_schedule_->TotalOrder()) {
    TF_RETURN_IF_ERROR(thunk->Initialize(*this));
    int32 stream_no =
        thunk_schedule_->StreamNumberForHlo(*thunk->hlo_instruction());
    se::Stream* stream =
        (stream_no == 0 ? main_stream : sub_streams[stream_no - 1].get());

    for (const Thunk* dependency : thunk_schedule_->DependsOn(thunk)) {
      stream->ThenWaitFor(FindOrDie(thunk_to_finish_event, dependency).get());
    }

    profiler.StartOperation();
    VLOG(2) << "Executing the thunk for "
            << thunk->hlo_instruction()->ToString();
    TF_RETURN_IF_ERROR(thunk->ExecuteOnStream(buffer_allocations, stream));
    if (thunk_schedule_->Depended(thunk)) {
      auto finish_event = MakeUnique<se::Event>(main_stream->parent());
      finish_event->Init();
      stream->ThenRecordEvent(finish_event.get());
      thunk_to_finish_event[thunk] = std::move(finish_event);
    }
    profiler.FinishOperation(thunk->hlo_instruction());
  }

  main_stream->ThenWaitFor(&sub_streams);
  // Make sure kernels are completed before deallocating temporary buffers.
  // TODO(b/30100571): we could potentially postpone deallocating the temp
  // buffers until a different computation is executed.
  if (block_host_until_done && !main_stream->BlockHostUntilDone()) {
    return InternalError("Failed to complete all kernels launched on stream %p",
                         main_stream);
  }

  return Status::OK();
}

StatusOr<se::DeviceMemoryBase> GpuExecutable::ExecuteOnStream(
    const ServiceExecutableRunOptions* run_options,
    tensorflow::gtl::ArraySlice<se::DeviceMemoryBase> arguments,
    HloExecutionProfile* hlo_execution_profile) {
  se::Stream* stream = run_options->stream();
  DeviceMemoryAllocator* memory_allocator = run_options->allocator();

  BufferAllocations::Builder buffer_allocations_builder;
  for (BufferAllocation::Index i = 0; i < assignment_->Allocations().size();
       ++i) {
    const BufferAllocation& allocation = assignment_->GetAllocation(i);
    if (allocation.is_entry_computation_parameter()) {
      buffer_allocations_builder.RegisterBuffer(
          i, arguments[allocation.parameter_number()]);
    }
  }
  se::StreamExecutor* executor = stream->parent();
  TF_ASSIGN_OR_RETURN(
      auto buffer_allocations,
      buffer_allocations_builder.Build(*assignment_, executor->device_ordinal(),
                                       memory_allocator));

  bool block_host_until_done =
      !memory_allocator->AllowsAsynchronousDeallocation();
  TF_RETURN_IF_ERROR(ExecuteThunks(run_options, *buffer_allocations,
                                   block_host_until_done,
                                   hlo_execution_profile));

  HloInstruction* root = hlo_module_->entry_computation()->root_instruction();
  TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice output_slice,
                      assignment_->GetUniqueTopLevelOutputSlice());
  se::DeviceMemoryBase output_buffer_address =
      buffer_allocations->GetDeviceAddress(output_slice.index());

  if (ShapeUtil::IsTuple(root->shape())) {
    std::set<se::DeviceMemoryBase> referred_by_output;
    if (GetRootPointsToSet().IsAmbiguous()) {
      // The points-to set of the root is ambiguous so we need to examine the
      // result data to determine which buffers are contained in the result.
      TF_ASSIGN_OR_RETURN(
          TransferManager * transfer_manager,
          TransferManager::GetForPlatform(executor->platform()));
      TF_ASSIGN_OR_RETURN(referred_by_output,
                          transfer_manager->GatherBufferPointersFromTuple(
                              executor, output_buffer_address, root->shape()));
    } else {
      // The points-to set of the root is unambiguous so it's known statically
      // which buffers are in the result. Gather these buffers using the root's
      // points-to set.
      TF_RETURN_IF_ERROR(GetRootPointsToSet().ForEachElementWithStatus(
          [&referred_by_output, &buffer_allocations, this](
              const ShapeIndex& /*index*/,
              const PointsToSet::BufferList& buffers) {
            // The points to set is unambiguous so the set should be a
            // singleton. That is, we know exactly which instruction produced
            // the array at this element.
            CHECK_EQ(1, buffers.size());
            HloInstruction* hlo = buffers[0]->instruction();
            TF_ASSIGN_OR_RETURN(
                const BufferAllocation::Slice slice,
                this->assignment_->GetUniqueSlice(hlo, buffers[0]->index()));
            CHECK(!slice.allocation()->is_entry_computation_parameter());
            referred_by_output.insert(
                buffer_allocations->GetDeviceAddress(slice.index()));
            return Status::OK();
          }));
    }
    TF_RETURN_IF_ERROR(
        buffer_allocations->TearDown(referred_by_output, *assignment_));
  } else {
    // If the computation result is not a tuple, we can delete all temporary
    // buffers that are not the output.
    TF_RETURN_IF_ERROR(
        buffer_allocations->TearDown({output_buffer_address}, *assignment_));
  }
  return output_buffer_address;
}

StatusOr<std::unique_ptr<ShapedBuffer>> GpuExecutable::ExecuteOnStream(
    const ServiceExecutableRunOptions* run_options,
    tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
    HloExecutionProfile* hlo_execution_profile) {
  DeviceMemoryAllocator* memory_allocator = run_options->allocator();

  if (GetRootPointsToSet().IsAmbiguous()) {
    return Unimplemented("Points-to set of root instruction is ambiguous");
  }

  BufferAllocations::Builder buffer_allocations_builder;
  for (BufferAllocation::Index i = 0; i < assignment_->Allocations().size();
       ++i) {
    const BufferAllocation& allocation = assignment_->GetAllocation(i);
    if (allocation.is_entry_computation_parameter()) {
      auto param_no = allocation.parameter_number();
      buffer_allocations_builder.RegisterBuffer(
          i, arguments[param_no]->buffer(/*index=*/{}));
    }
  }
  se::StreamExecutor* executor = run_options->stream()->parent();
  TF_ASSIGN_OR_RETURN(
      auto buffer_allocations,
      buffer_allocations_builder.Build(*assignment_, executor->device_ordinal(),
                                       memory_allocator));

  bool block_host_until_done =
      !memory_allocator->AllowsAsynchronousDeallocation();
  TF_RETURN_IF_ERROR(ExecuteThunks(run_options, *buffer_allocations,
                                   block_host_until_done,
                                   hlo_execution_profile));

  HloInstruction* root = hlo_module_->entry_computation()->root_instruction();
  auto device_ordinal = executor->device_ordinal();
  auto shaped_buffer = MakeUnique<ShapedBuffer>(
      root->shape(), executor->platform(), device_ordinal);

  // Copy DeviceMemoryBase values which contain the array(s) of the result into
  // the respective location in ShapedBuffer.
  std::set<se::DeviceMemoryBase> buffers_in_result;
  TF_RETURN_IF_ERROR(
      shaped_buffer->mutable_shape_index_to_buffer_entry()
          ->ForEachMutableElementWithStatus(
              [&buffer_allocations, &buffers_in_result, &shaped_buffer, this](
                  const ShapeIndex& index, size_t* buffer_entry) {
                const auto& sources = this->GetRootPointsToSet().element(index);
                // The points-to set is unambiguous so the set should be a
                // singleton. That is, we know exactly which instruction
                // produced the array at this element.
                CHECK_EQ(1, sources.size());
                auto src_hlo = sources[0]->instruction();

                VLOG(4) << "Looking at: " << sources[0];

                // The source instruction should have a non-parameter buffer
                // assigned.
                TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice,
                                    this->assignment_->GetUniqueSlice(
                                        src_hlo, sources[0]->index()));
                CHECK(!slice.allocation()->is_entry_computation_parameter());

                perftools::gputools::DeviceMemoryBase src_base =
                    buffer_allocations->GetDeviceAddress(slice.index());
                CHECK(!src_base.is_null() || src_base.size() == 0);
                shaped_buffer->mutable_buffers()->push_back(src_base);
                *buffer_entry = shaped_buffer->mutable_buffers()->size() - 1;

                buffers_in_result.insert(src_base);
                return Status::OK();
              }));
  TF_RETURN_IF_ERROR(
      buffer_allocations->TearDown(buffers_in_result, *assignment_));

  return std::move(shaped_buffer);
}

StatusOr<se::DeviceMemoryBase> GpuExecutable::ExecuteAsyncOnStream(
    const ServiceExecutableRunOptions* run_options,
    tensorflow::gtl::ArraySlice<se::DeviceMemoryBase> arguments) {
  // TODO(b/30671675): Implement asynchronous execution mode.
  return Unimplemented(
      "Asynchronous execution on stream is not yet supported on GPU.");
}

const PointsToSet& GpuExecutable::GetRootPointsToSet() const {
  return assignment_->points_to_analysis().GetPointsToSet(
      module().entry_computation()->root_instruction());
}

std::unique_ptr<HloCostAnalysis> GpuExecutable::CreateCostAnalysis() const {
  return MakeUnique<HloCostAnalysis>(shape_size_function_);
}

}  // namespace gpu
}  // namespace xla
