/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/async_transferer.cc
 * \brief The AsyncTransferer implementation.
 */


#include "async_transferer.h"

#include <dgl/packed_func_ext.h>
#include <dgl/runtime/registry.h>
#include <dgl/runtime/device_api.h>
#include <vector>
#include <utility>

#ifdef DGL_USE_CUDA
#include <cuda_runtime.h>
#include "../runtime/cuda/cuda_common.h"
#endif


namespace dgl {

using namespace runtime;

namespace dataloading {

using TransferId = AsyncTransferer::TransferId;

struct AsyncTransferer::Event {
  #ifdef DGL_USE_CUDA
  cudaEvent_t id;

  ~Event() {
    CUDA_CALL(cudaEventDestroy(id));
  }
  #endif
};

AsyncTransferer::AsyncTransferer(
    DGLContext ctx) :
  ctx_(ctx),
  next_id_(0),
  transfers_(),
  stream_(nullptr) {
  if (ctx_.device_type == kDLGPU) {
    stream_ = DeviceAPI::Get(ctx_)->CreateStream(ctx_);
  }
}


AsyncTransferer::~AsyncTransferer() {
  if (stream_) {
    DeviceAPI::Get(ctx_)->FreeStream(ctx_, stream_);
  }
}

TransferId AsyncTransferer::StartTransfer(
    NDArray src,
    DGLContext dst_ctx) {
  const TransferId id = GenerateId();

  Transfer t;
  t.src = src;

  DLDataType dtype = src->dtype;
  std::vector<int64_t> shape(src->shape, src->shape+src->ndim);
  t.dst = NDArray::Empty(shape, dtype, dst_ctx);

  if (stream_) {
    #ifdef DGL_USE_CUDA
    // get tensor information
    t.event.reset(new Event);
    CUDA_CALL(cudaEventCreate(&t.event->id));
    t.dst.CopyFrom(t.src, stream_);

    CUDA_CALL(cudaEventRecord(t.event->id, static_cast<cudaStream_t>(stream_)));
    #else
    LOG(FATAL) << "GPU support not compiled.";
    #endif
  } else {
    // copy synchronously since we don't have the notion of streams on the CPU
    t.event.reset(nullptr);
    t.dst.CopyFrom(t.src);
  }
  transfers_.emplace(id, std::move(t));

  return id;
}

NDArray AsyncTransferer::Wait(
    const TransferId id) {
  auto iter = transfers_.find(id);
  CHECK(iter != transfers_.end()) << "Unknown transfer: " << id;

  Transfer t = std::move(iter->second);
  transfers_.erase(iter);

  if (t.event) {
    #ifdef DGL_USE_CUDA
    // wait for it
    CUDA_CALL(cudaEventSynchronize(t.event->id));
    #endif
  }

  return t.dst;
}

TransferId AsyncTransferer::GenerateId() {
  return ++next_id_;
}

DGL_REGISTER_GLOBAL("dataloading.async_transferer._CAPI_DGLAsyncTransfererCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    DGLContext ctx = args[0];
    *rv = AsyncTransfererRef(std::make_shared<AsyncTransferer>(ctx));
});

DGL_REGISTER_GLOBAL("dataloading.async_transferer._CAPI_DGLAsyncTransfererStartTransfer")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
  AsyncTransfererRef ref = args[0];
  NDArray array = args[1];
  DGLContext ctx = args[2];
  int id = ref->StartTransfer(array, ctx);
  *rv = id;
});

DGL_REGISTER_GLOBAL("dataloading.async_transferer._CAPI_DGLAsyncTransfererWait")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
  AsyncTransfererRef ref = args[0];
  int id = args[1];
  NDArray arr = ref->Wait(id);
  *rv = arr;
});

}  // namespace dataloading
}  // namespace dgl