/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "./cudnn_convolution.h"
#ifdef USE_CUDNN
#include <cudnn.h>
#include <chrono>
#include "./cudnn_utils.h"
#include "singa/utils/logging.h"

namespace singa {
RegisterLayerClass(cudnn_convolution, CudnnConvolution);
CudnnConvolution::~CudnnConvolution() {
  if (bias_desc_ != nullptr)
    CUDNN_CHECK(cudnnDestroyTensorDescriptor(bias_desc_));
  if (filter_desc_ != nullptr)
    CUDNN_CHECK(cudnnDestroyFilterDescriptor(filter_desc_));
  if (conv_desc_ != nullptr)
    CUDNN_CHECK(cudnnDestroyConvolutionDescriptor(conv_desc_));
  if (x_desc_ != nullptr) CUDNN_CHECK(cudnnDestroyTensorDescriptor(x_desc_));
  if (y_desc_ != nullptr) CUDNN_CHECK(cudnnDestroyTensorDescriptor(y_desc_));
}

void CudnnConvolution::Setup(const Shape& in_sample, const LayerConf &conf) {
  Convolution::Setup(in_sample, conf);
  ConvolutionConf conv_conf = conf.convolution_conf();
  // convert MB to bytes
  workspace_byte_limit_ = conv_conf.workspace_byte_limit() << 20;
  prefer_ = ToLowerCase(conv_conf.prefer());
  CHECK(prefer_ == "fastest" || prefer_ == "limited_workspace" ||
        prefer_ == "no_workspace" || prefer_ == "autotune")
      << "CudnnConvolution only supports four algorithm preferences: fastest, "
      "limited_workspace, no_workspace and autotune";
}

void CudnnConvolution::ToDevice(std::shared_ptr<Device> device) {
  Convolution::ToDevice(device);
  workspace_.ToDevice(device);
}

void CudnnConvolution::InitCudnn(const Tensor &input) {
  DataType dtype = input.data_type();
  auto dev = input.device();
  Context *ctx = dev->context(0);
  size_t batchsize = input.shape(0);
  if (!has_init_cudnn_) {
    CUDNN_CHECK(cudnnCreateTensorDescriptor(&x_desc_));
    CUDNN_CHECK(cudnnCreateTensorDescriptor(&y_desc_));
    if (bias_term_)
      CUDNN_CHECK(cudnnCreateTensorDescriptor(&bias_desc_));
    CUDNN_CHECK(cudnnCreateFilterDescriptor(&filter_desc_));
    CUDNN_CHECK(cudnnCreateConvolutionDescriptor(&conv_desc_));
  }

  CUDNN_CHECK(cudnnSetTensor4dDescriptor(x_desc_, CUDNN_TENSOR_NCHW,
                                         GetCudnnDataType(dtype), batchsize,
                                         channels_, height_, width_));
  CUDNN_CHECK(cudnnSetTensor4dDescriptor(
                y_desc_, CUDNN_TENSOR_NCHW, GetCudnnDataType(dtype), batchsize,
                num_filters_, conv_height_, conv_width_));
  if (bias_term_)
    CUDNN_CHECK(cudnnSetTensor4dDescriptor(bias_desc_, CUDNN_TENSOR_NCHW,
                                           GetCudnnDataType(dtype), 1,
                                           num_filters_, 1, 1));
  CUDNN_CHECK(cudnnSetConvolution2dDescriptor(conv_desc_, pad_h_, pad_w_,
              stride_h_, stride_w_, 1, 1,  // dilation x and y
              CUDNN_CROSS_CORRELATION
#if CUDNN_MAJOR >= 7
              , GetCudnnDataType(dtype)
#endif  // CUDNN_MAJOR
                                             ));
  CUDNN_CHECK(cudnnSetFilter4dDescriptor(filter_desc_, GetCudnnDataType(dtype),
                                         CUDNN_TENSOR_NCHW, num_filters_,
                                         channels_, kernel_h_, kernel_w_));
  if (prefer_ == "fastest" || prefer_ == "limited_workspace" ||
      prefer_ == "no_workspace") {
    cudnnConvolutionFwdPreference_t fwd_pref;
    cudnnConvolutionBwdFilterPreference_t bwd_filt_pref;
    cudnnConvolutionBwdDataPreference_t bwd_data_pref;
    if (prefer_ == "fastest") {
      fwd_pref = CUDNN_CONVOLUTION_FWD_PREFER_FASTEST;
      bwd_filt_pref = CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST;
      bwd_data_pref = CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST;
    } else if (prefer_ == "limited_workspace") {
      fwd_pref = CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT;
      bwd_filt_pref = CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT;
      bwd_data_pref = CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT;
    } else {
      fwd_pref = CUDNN_CONVOLUTION_FWD_NO_WORKSPACE;
      bwd_filt_pref = CUDNN_CONVOLUTION_BWD_FILTER_NO_WORKSPACE;
      bwd_data_pref = CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT;
    }
    CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(
                  ctx->cudnn_handle, x_desc_, filter_desc_, conv_desc_, y_desc_, fwd_pref,
                  workspace_byte_limit_, &fp_alg_));
    CUDNN_CHECK(cudnnGetConvolutionBackwardFilterAlgorithm(
                  ctx->cudnn_handle, x_desc_, y_desc_, conv_desc_, filter_desc_,
                  bwd_filt_pref, workspace_byte_limit_, &bp_filter_alg_));
    // deprecated in cudnn v7
    CUDNN_CHECK(cudnnGetConvolutionBackwardDataAlgorithm(
                  ctx->cudnn_handle, filter_desc_, y_desc_, conv_desc_, x_desc_,
                  bwd_data_pref, workspace_byte_limit_, &bp_data_alg_));
  } else if (prefer_ == "autotune") {
    const int topk = 1;
    int num_fp_alg, num_bp_filt_alg, num_bp_data_alg;
    cudnnConvolutionFwdAlgoPerf_t fp_alg_perf[topk];
    cudnnConvolutionBwdFilterAlgoPerf_t bp_filt_perf[topk];
    cudnnConvolutionBwdDataAlgoPerf_t bp_data_perf[topk];
    CUDNN_CHECK(cudnnFindConvolutionForwardAlgorithm(
                  ctx->cudnn_handle, x_desc_, filter_desc_, conv_desc_, y_desc_, topk,
                  &num_fp_alg, fp_alg_perf));
    fp_alg_ = fp_alg_perf[0].algo;
    CUDNN_CHECK(cudnnFindConvolutionBackwardFilterAlgorithm(
                  ctx->cudnn_handle, x_desc_, y_desc_, conv_desc_, filter_desc_, topk,
                  &num_bp_filt_alg, bp_filt_perf));
    bp_filter_alg_ = bp_filt_perf[0].algo;
    CUDNN_CHECK(cudnnFindConvolutionBackwardDataAlgorithm(
                  ctx->cudnn_handle, filter_desc_, y_desc_, conv_desc_, x_desc_, topk,
                  &num_bp_data_alg, bp_data_perf));
    bp_data_alg_ = bp_data_perf[0].algo;
  } else {
    LOG(FATAL) << "Preferred algorithm is not available!";
  }

  size_t fp_byte, bp_data_byte, bp_filter_byte;
  CUDNN_CHECK(cudnnGetConvolutionForwardWorkspaceSize(
                ctx->cudnn_handle, x_desc_, filter_desc_, conv_desc_, y_desc_, fp_alg_,
                &fp_byte));
  CUDNN_CHECK(cudnnGetConvolutionBackwardDataWorkspaceSize(
                ctx->cudnn_handle, filter_desc_, y_desc_, conv_desc_, x_desc_,
                bp_data_alg_, &bp_data_byte));
  CUDNN_CHECK(cudnnGetConvolutionBackwardFilterWorkspaceSize(
                ctx->cudnn_handle, x_desc_, y_desc_, conv_desc_, filter_desc_,
                bp_filter_alg_, &bp_filter_byte));
  workspace_count_ = std::max(std::max(fp_byte, bp_data_byte), bp_filter_byte) /
                     sizeof(float) +
                     1;
  if (workspace_count_ * sizeof(float) > workspace_byte_limit_)
    LOG(WARNING) << "The required memory for workspace ("
                 << workspace_count_ * sizeof(float)
                 << ") is larger than the expected Bytes ("
                 << workspace_byte_limit_ << ")";
  workspace_ = Tensor(Shape{workspace_count_}, dev, dtype);
  has_init_cudnn_ = true;
}

const Tensor CudnnConvolution::Forward(int flag, const Tensor &input) {
  CHECK(buf_.empty());
  CHECK_EQ(input.device()->lang(), kCuda);
  CHECK_EQ(input.nDim(), 4u);
  if (flag & kTrain) buf_.push(input);  // buffer the input for backward
  size_t batchsize = input.shape()[0];
  DataType dtype = input.data_type();
  auto dev = input.device();

  if (!has_init_cudnn_) {
    InitCudnn(input);
  } else {
    int n, c, h, w, s;
    cudnnDataType_t type;
    CUDNN_CHECK(cudnnGetTensor4dDescriptor(x_desc_, &type, &n, &c, &h, &w,
                                           &s, &s, &s, &s));
    if (batchsize != static_cast<size_t>(n))
      InitCudnn(input);
    CHECK(input.shape(1) == static_cast<size_t>(c)
          && input.shape(2) == static_cast<size_t>(h)
          && input.shape(3) == static_cast<size_t>(w))
        << "input sample shape should not change"
        << "previous shape " << c << ", " << h << ", " << w
        << "current shape " << input.shape(1) << ", " << input.shape(2) << ", "
        << input.shape(3);
  }

  Shape shape{batchsize, num_filters_, conv_height_, conv_width_};
  Tensor output(shape, dev, dtype);
  output.device()->Exec([input, output, this](Context * ctx) {
    Block *inblock = input.block(), *outblock = output.block(),
           *wblock = this->weight_.block();
    float alpha = 1.f, beta = 0.f;
    cudnnConvolutionForward(ctx->cudnn_handle, &alpha, this->x_desc_,
                            inblock->data(), this->filter_desc_, wblock->data(),
                            this->conv_desc_, this->fp_alg_,
                            this->workspace_.block()->mutable_data(),
                            this->workspace_count_ * sizeof(float), &beta,
                            this->y_desc_, outblock->mutable_data());
  }, {input.block(), weight_.block()}, {output.block(), workspace_.block()});

  if (bias_term_) {
    output.device()->Exec([output, this](Context * ctx) {
      float beta = 1.f, alpha = 1.0f;
      Block *outblock = output.block(), *bblock = this->bias_.block();
      cudnnAddTensor(ctx->cudnn_handle, &alpha, this->bias_desc_,
                     bblock->data(), &beta, this->y_desc_,
                     outblock->mutable_data());
    }, {output.block(), bias_.block()}, {output.block()});
  }
  return output;
}

const std::pair<Tensor, vector<Tensor>> CudnnConvolution::Backward(
int flag, const Tensor &grad) {
  CHECK(has_init_cudnn_);
  CHECK_EQ(grad.device()->lang(), kCuda);
  CHECK_EQ(grad.nDim(), 4u);
  CHECK(!buf_.empty());
  Tensor src_data = buf_.top();
  buf_.pop();
  vector<Tensor> param_grad;
  Tensor dx;
  dx.ResetLike(src_data);
  Tensor db, dw;
  dw.ResetLike(weight_);

  // LOG(ERROR) << "backward bias";
  if (bias_term_) {
    db.ResetLike(bias_);
    dx.device()->Exec([grad, db, this](Context * ctx) {
      Block *dyblock = grad.block(), *dbblock = db.block();
      float alpha = 1.f, beta = 0.f;
      cudnnConvolutionBackwardBias(ctx->cudnn_handle, &alpha, this->y_desc_,
                                   dyblock->data(), &beta, this->bias_desc_,
                                   dbblock->mutable_data());
    }, {grad.block()}, {db.block()});
  }
  // LOG(ERROR) << "backward w";
  dx.device()->Exec([grad, dw, src_data, this](Context * ctx) {
    Block *inblock = src_data.block(), *dyblock = grad.block(),
           *dwblock = dw.block();
    float alpha = 1.f, beta = 0.f;
    cudnnConvolutionBackwardFilter(
      ctx->cudnn_handle, &alpha, this->x_desc_, inblock->data(),
      this->y_desc_, dyblock->data(), this->conv_desc_, this->bp_filter_alg_,
      this->workspace_.block()->mutable_data(),
      this->workspace_count_ * sizeof(float), &beta, this->filter_desc_,
      dwblock->mutable_data());
  }, {grad.block(), src_data.block()}, {dw.block(), workspace_.block()});

  // LOG(ERROR) << "backward src";
  dx.device()->Exec([dx, grad, this](Context * ctx) {
    Block *wblock = this->weight_.block(), *dyblock = grad.block(),
           *dxblock = dx.block();
    float alpha = 1.f, beta = 0.f;
    cudnnConvolutionBackwardData(ctx->cudnn_handle, &alpha, this->filter_desc_,
                                 wblock->data(), this->y_desc_, dyblock->data(),
                                 this->conv_desc_, this->bp_data_alg_,
                                 this->workspace_.block()->mutable_data(),
                                 this->workspace_count_ * sizeof(float), &beta,
                                 this->x_desc_, dxblock->mutable_data());
  }, {grad.block(), weight_.block()}, {dx.block(), workspace_.block()});
  param_grad.push_back(dw);
  if (bias_term_)
    param_grad.push_back(db);
  return std::make_pair(dx, param_grad);
}

}  // namespace singa
#endif  // USE_CUDNN