#ifndef CAFFE2_CUDA_RTC_COMMON_RTC_H_
#define CAFFE2_CUDA_RTC_COMMON_RTC_H_

#include <sstream>
#include <string>

#include <cuda.h>
#include <nvrtc.h>

#define NVRTC_CHECK(condition)                                                 \
  do {                                                                         \
    nvrtcResult result = condition;                                            \
    if (result != NVRTC_SUCCESS) {                                             \
      LOG(FATAL) << "Error at: " << __FILE__ << ":" << __LINE__ << ": "   \
                      << nvrtcGetErrorString(result);                          \
    }                                                                          \
  } while(0)

namespace caffe2 {

template <typename Derived>
class CudaRTCFunction {
 public:
  CudaRTCFunction() : module_loaded_(false) {}
  ~CudaRTCFunction() {
    if (module_loaded_) {
      CUDA_DRIVERAPI_ENFORCE(cuModuleUnload(module_));
    }
  }

  // TODO: this function is nontrivial and since CudaRTCFunction uses CRTP, it
  // may potentially increase the binary size. In that case, move common parts
  // into a separate function.
  template <typename... Args>
  void Compile(Args... args) {
    string src = static_cast<Derived*>(this)->GetSource(args...);
    string name = static_cast<Derived*>(this)->KernelName(args...);
    VLOG(1) << "function name: " << name;
    VLOG(1) << "function src:\n" << src;
    // Actually do the compiling.
    nvrtcProgram prog;
    NVRTC_CHECK(nvrtcCreateProgram(
        &prog, src.c_str(), nullptr, 0, nullptr, nullptr));
    // Compile the program.
    // TODO(Yangqing): how to find the current gpu architecture instead of hard
    // coding it?
    const char *nvrtc_opts[] = {"--gpu-architecture=compute_35",
                                "--use_fast_math"};
    nvrtcResult compile_result = nvrtcCompileProgram(
        prog, 2, nvrtc_opts);
    if (compile_result != NVRTC_SUCCESS) {
      size_t log_size;
      NVRTC_CHECK(nvrtcGetProgramLogSize(prog, &log_size));
      vector<char> nvrtc_log(log_size);
      NVRTC_CHECK(nvrtcGetProgramLog(prog, nvrtc_log.data()));
      LOG(FATAL) << "Compilation failure for nvrtc("
                 << nvrtcGetErrorString(compile_result) << "): \n"
                 << nvrtc_log.data();
    }
    size_t ptx_size;
    NVRTC_CHECK(nvrtcGetPTXSize(prog, &ptx_size));
    vector<char> nvrtc_ptx(ptx_size);
    NVRTC_CHECK(nvrtcGetPTX(prog, nvrtc_ptx.data()));
    NVRTC_CHECK(nvrtcDestroyProgram(&prog));
    // After compilation, load the module.
    if (module_loaded_) {
      CUDA_DRIVERAPI_ENFORCE(cuModuleUnload(module_));
    }
    CUDA_DRIVERAPI_ENFORCE(
        cuModuleLoadDataEx(&module_, nvrtc_ptx.data(), 0, 0, 0));
    module_loaded_ = true;
    CUDA_DRIVERAPI_ENFORCE(
        cuModuleGetFunction(&kernel_, module_, name.c_str()));
  }

  template <typename... Args>
  void Launch(unsigned int gx, unsigned int gy, unsigned int gz,
              unsigned int bx, unsigned int by, unsigned int bz,
              unsigned int shared_mem, cudaStream_t stream,
              Args... args) {
    CAFFE_ENFORCE(
        module_loaded_, "Cannot call Launch before a module is loaded.");
    void * args_voidp[] = {&args...};
    CUDA_DRIVERAPI_ENFORCE(cuLaunchKernel(
        kernel_, gx, gy, gz, bx, by, bz, shared_mem, stream, args_voidp, 0));
  }

  void LaunchEx(unsigned int gx, unsigned int gy, unsigned int gz,
                unsigned int bx, unsigned int by, unsigned int bz,
                unsigned int shared_mem, cudaStream_t stream,
                void** extra) {
    CAFFE_ENFORCE(
        module_loaded_, "Cannot call Launch before a module is loaded.");
    CUDA_DRIVERAPI_ENFORCE(cuLaunchKernel(
        kernel_, gx, gy, gz, bx, by, bz, shared_mem, stream, nullptr, extra));
  }

 private:
  bool module_loaded_;
  CUmodule module_;
  CUfunction kernel_;
};

// TODO: this is in no way unique and is just a hack right now.
inline std::string GetUniqueName() {
  static constexpr int len = 20;
  static const char alpha[] =
      "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

  std::stringstream ss;
  ss << "_cuda_kernel_";
  for (int i = 0; i < len; ++i) {
    ss << alpha[rand() % (sizeof(alpha) - 1)];
  }
  return ss.str();
}

}  // namepsace caffe2

#endif  // CAFFE2_CUDA_RTC_COMMON_RTC_H_