#include "caffe2/operators/conv_transpose_op.h"

#include <vector>

#include "caffe2/core/context_gpu.h"
#include "caffe2/core/cudnn_wrappers.h"
#include "caffe2/operators/conv_op_cache_cudnn.h"
#include "caffe2/operators/op_utils_cudnn.h"

namespace caffe2 {

class CudnnConvTransposeOpBase : public ConvTransposeUnpoolBase<CUDAContext> {
 public:
  template <class... Args>
  explicit CudnnConvTransposeOpBase(Args&&... args)
      : ConvTransposeUnpoolBase<CUDAContext>(std::forward<Args>(args)...),
        cudnn_wrapper_(&context_),
        cudnn_ws_nbytes_limit_(OperatorBase::GetSingleArgument<size_t>(
            "ws_nbytes_limit",
            kCONV_CUDNN_WORKSPACE_LIMIT_BYTES)),
        exhaustive_search_(
            OperatorBase::GetSingleArgument<int>("exhaustive_search", 0)),
        deterministic_(
            OperatorBase::GetSingleArgument<int>("deterministic", 0)),
        cudnn_state_(OperatorBase::GetSingleArgument<int>("cudnn_state", 0)),
        force_algo_(OperatorBase::GetRepeatedArgument<int>(
            "force_algo",
            vector<int>{-1, -1, -1})),
        enable_tensor_core_(
            OperatorBase::GetSingleArgument<bool>("enable_tensor_core", 1)) {
    CAFFE_ENFORCE(!deterministic_ || !exhaustive_search_);

    bool individual_force_algo = OperatorBase::HasArgument("force_algo_fwd") ||
        OperatorBase::HasArgument("force_algo_dgrad") ||
        OperatorBase::HasArgument("force_algo_wgrad");
    if (OperatorBase::HasArgument("force_algo")) {
      CAFFE_ENFORCE(
          !individual_force_algo,
          "Cannot specify both force_algo and any of",
          "force_algo_fwd, force_algo_dgrad, force_algo_wgrad");
    } else {
      force_algo_ = std::vector<int>{-1, -1, -1};
      force_algo_[ALGO_FWD] =
          OperatorBase::GetSingleArgument<int>("force_algo_fwd", -1);
      force_algo_[ALGO_DGRAD] =
          OperatorBase::GetSingleArgument<int>("force_algo_dgrad", -1);
      force_algo_[ALGO_WGRAD] =
          OperatorBase::GetSingleArgument<int>("force_algo_wgrad", -1);
    }

    CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&bottom_desc_));
    CUDNN_ENFORCE(cudnnCreateFilterDescriptor(&filter_desc_));
    if (InputSize() == 3) {
      CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&bias_desc_));
      CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&top_desc_for_bias_));
    }
    CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&top_desc_));
    CUDNN_ENFORCE(cudnnCreateConvolutionDescriptor(&conv_desc_));
  }

  ~CudnnConvTransposeOpBase() override {
    CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(bottom_desc_));
    CUDNN_ENFORCE(cudnnDestroyFilterDescriptor(filter_desc_));
    if (InputSize() == 3) {
      CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(bias_desc_));
      CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(top_desc_for_bias_));
    }
    CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(top_desc_));
    CUDNN_ENFORCE(cudnnDestroyConvolutionDescriptor(conv_desc_));
  }

 protected:
  void SetTensor4DDescriptorWithGroup(
      const cudnnDataType_t data_type,
      const int N,
      const int C,
      const int H,
      const int W,
      cudnnTensorDescriptor_t* desc) const {
#if CUDNN_VERSION_MIN(7, 0, 0)
    const int CC = C;
#else
    const int CC = C / group_;
#endif
    switch (order_) {
      case StorageOrder::NCHW: {
        CUDNN_ENFORCE(cudnnSetTensor4dDescriptorEx(
            *desc, data_type, N, CC, H, W, C * H * W, H * W, W, 1));
        break;
      }
      case StorageOrder::NHWC: {
        CUDNN_ENFORCE(cudnnSetTensor4dDescriptorEx(
            *desc, data_type, N, CC, H, W, H * W * C, 1, W * C, C));
        break;
      }
      default: {
        LOG(FATAL) << "Unknown storage order: " << order_;
      }
    }
  }

  std::vector<std::int64_t> cudnn_input_dims_;
  std::vector<std::int64_t> cudnn_filter_dims_;

  CuDNNWrapper cudnn_wrapper_;
  cudnnTensorDescriptor_t bottom_desc_;
  cudnnFilterDescriptor_t filter_desc_;
  cudnnTensorDescriptor_t bias_desc_;
  cudnnTensorDescriptor_t top_desc_;
  cudnnTensorDescriptor_t top_desc_for_bias_;
  cudnnConvolutionDescriptor_t conv_desc_;

  const size_t cudnn_ws_nbytes_limit_;
  size_t cudnn_ws_nbytes_;
  bool exhaustive_search_;
  bool deterministic_;
  size_t cudnn_state_;
  std::vector<int> force_algo_; // stored as FWD, dFILTER, dDATA
  bool enable_tensor_core_;
};

template <typename T>
class CudnnConvTransposeOp final : public CudnnConvTransposeOpBase {
 public:
  template <class... Args>
  explicit CudnnConvTransposeOp(Args&&... args)
      : CudnnConvTransposeOpBase(std::forward<Args>(args)...) {}

  ~CudnnConvTransposeOp() override {}

  bool RunOnDevice() override;

 private:
  AlgorithmsCache<cudnnConvolutionBwdDataAlgo_t> data_algo_cache_;
  cudnnConvolutionBwdDataAlgo_t bwd_data_algo_;
  // Input: X, W, b
  // Output: Y
  INPUT_TAGS(INPUT, FILTER, BIAS);
};

template <typename T>
class CudnnConvTransposeGradientOp final : public CudnnConvTransposeOpBase {
 public:
  template <class... Args>
  explicit CudnnConvTransposeGradientOp(Args&&... args)
      : CudnnConvTransposeOpBase(std::forward<Args>(args)...),
        no_bias_(OperatorBase::GetSingleArgument<bool>("no_bias", false)) {
    CAFFE_ENFORCE(
        !(no_bias_ && OutputSize() == 3),
        "If bias is not present, you should not have 3 grad output.");
  }

  ~CudnnConvTransposeGradientOp() override {}

  bool RunOnDevice() override;

 private:
  cudnnConvolutionFwdAlgo_t algo_;
  cudnnConvolutionBwdFilterAlgo_t bwd_filter_algo_;
  AlgorithmsCache<cudnnConvolutionFwdAlgo_t> forward_algo_cache_;
  AlgorithmsCache<cudnnConvolutionBwdFilterAlgo_t> filter_algo_cache_;
  const bool no_bias_;
  // input: X, W, dY
  // output: dW, optionally db and dX
  INPUT_TAGS(INPUT, FILTER, OUTPUT_GRAD);
  OUTPUT_TAGS(FILTER_GRAD, BIAS_OR_INPUT_GRAD, INPUT_GRAD);
};

////////////////////////////////////////////////////////////////////////////////
// Implementations
////////////////////////////////////////////////////////////////////////////////

template <typename T>
bool CudnnConvTransposeOp<T>::RunOnDevice() {
  auto& X = Input(INPUT);
  auto& filter = Input(FILTER);
  int C = 0;
  switch (order_) {
    case StorageOrder::NHWC:
      C = filter.dim32(3) * group_;
      break;
    case StorageOrder::NCHW:
      C = filter.dim32(1) * group_;
      break;
    default:
      LOG(FATAL) << "Unknown storage order: " << order_;
  }

  auto sizes = ConvTransposeUnpoolBase<CUDAContext>::GetOutputSize(X, C);
  auto* Y = Output(0, sizes, at::dtype<T>());

  if (X.numel() == 0) {
    VLOG(2) << "Number on elements is 0 in CudnnConvTransposeOp";
    return true;
  }

  int N = 0, M = 0, H = 0, W = 0, H_out = 0, W_out = 0;
  switch (order_) {
    case StorageOrder::NHWC:
      N = X.dim32(0);
      H = X.dim32(1);
      W = X.dim32(2);
      M = X.dim32(3);
      H_out = Y->dim32(1);
      W_out = Y->dim32(2);
      CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
      CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_w());
      CAFFE_ENFORCE_EQ(filter.dim32(3), C / group_);
      break;
    case StorageOrder::NCHW:
      N = X.dim32(0);
      M = X.dim32(1);
      H = X.dim32(2);
      W = X.dim32(3);
      H_out = Y->dim32(2);
      W_out = Y->dim32(3);
      CAFFE_ENFORCE_EQ(filter.dim32(1), C / group_);
      CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_h());
      CAFFE_ENFORCE_EQ(filter.dim32(3), kernel_w());
      break;
    default:
      LOG(FATAL) << "Unknown storage order: " << order_;
  }
  CAFFE_ENFORCE_EQ(M % group_, 0);

  if (InputSize() == 3) {
    auto& bias = Input(BIAS);
    CAFFE_ENFORCE_EQ(bias.dim(), 1);
    CAFFE_ENFORCE_EQ(bias.dim32(0), C);
  }

  // Set up the cudnn algorithms & workspace if necessary
  const bool input_changed = (X.sizes() != cudnn_input_dims_);
  const bool filter_changed = (filter.sizes() != cudnn_filter_dims_);

  if (input_changed || filter_changed) {
    VLOG(1) << "Changing the cudnn descriptor configurations.";
    if (input_changed) {
      cudnn_input_dims_ = X.sizes().vec();
      SetTensor4DDescriptorWithGroup(
          cudnnTypeWrapper<T>::type, N, M, H, W, &bottom_desc_);
    }
    if (filter_changed) {
      cudnn_filter_dims_ = filter.sizes().vec();
#if CUDNN_VERSION_MIN(7, 0, 0)
      const int MM = M;
#else
      const int MM = M / group_;
#endif
      CUDNN_ENFORCE(cudnnSetFilter4dDescriptor(
          filter_desc_,
          cudnnTypeWrapper<T>::type,
          GetCudnnTensorFormat(order_),
          MM,
          C / group_,
          kernel_h(),
          kernel_w()));
      if (InputSize() == 3) {
        CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
            bias_desc_,
            GetCudnnTensorFormat(order_),
            cudnnTypeWrapper<T>::type,
            1,
            C,
            1,
            1));
      }
    }
    // Set the output
    SetTensor4DDescriptorWithGroup(
        cudnnTypeWrapper<T>::type, N, C, H_out, W_out, &top_desc_);
    if (InputSize() == 3) {
      CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
          top_desc_for_bias_,
          GetCudnnTensorFormat(order_),
          cudnnTypeWrapper<T>::type,
          N,
          C,
          H_out,
          W_out));
    }

    // Set the convolution descriptor
    CAFFE_ENFORCE_EQ(
        pad_t(),
        pad_b(),
        "The current padding scheme leads to unequal padding on the top and "
        "bottom, which is not supported by cudnn.");
    CAFFE_ENFORCE_EQ(
        pad_l(),
        pad_r(),
        "The current padding scheme leads to unequal padding on the left "
        "and right, which is not supported by cudnn.");
    // Set the convolution descriptor
#if CUDNN_VERSION_MIN(6, 0, 0)
    CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
        conv_desc_,
        pad_t(),
        pad_l(),
        stride_h(),
        stride_w(),
        1,
        1,
        CUDNN_CROSS_CORRELATION,
        cudnnTypeWrapper<T>::type));
#else
    CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
        conv_desc_,
        pad_t(),
        pad_l(),
        stride_h(),
        stride_w(),
        1,
        1,
        CUDNN_CROSS_CORRELATION));
#endif

#if CUDNN_VERSION_MIN(7, 0, 0)
    // enable TensorCore math if desired
    enable_tensor_core_ &= TensorCoreAvailable();
    if (enable_tensor_core_) {
      CUDNN_ENFORCE(
          cudnnSetConvolutionMathType(conv_desc_, CUDNN_TENSOR_OP_MATH));
    }
    // set cuDNN groups if appropriate
    CUDNN_ENFORCE(cudnnSetConvolutionGroupCount(conv_desc_, group_));
#endif

    if (force_algo_[ALGO_DGRAD] >= 0) {
      bwd_data_algo_ = (cudnnConvolutionBwdDataAlgo_t)force_algo_[ALGO_DGRAD];
    } else if (deterministic_) {
      bwd_data_algo_ = CUDNN_CONVOLUTION_BWD_DATA_ALGO_1;
    } else if (exhaustive_search_) {
      bwd_data_algo_ =
          data_algo_cache_.getAlgorithm(X.sizes(), filter.sizes(), 0, [&]() {
            int returned_algo_count;
            std::array<
                cudnnConvolutionBwdDataAlgoPerf_t,
                kNUM_CUDNN_BWD_DATA_ALGS>
                data_perf_stat;
            cudnn_wrapper_.with_cudnn_state(
                cudnn_state_, [&](CuDNNState* state) {
                  state->workspace().reset();
                  CUDNN_ENFORCE(cudnnFindConvolutionBackwardDataAlgorithm(
                      state->cudnn_handle(),
                      filter_desc_,
                      bottom_desc_,
                      conv_desc_,
                      top_desc_,
                      kNUM_CUDNN_BWD_DATA_ALGS,
                      &returned_algo_count,
                      data_perf_stat.data()));
                });

            LogCuDNNPerfStats(data_perf_stat, returned_algo_count);
            return data_perf_stat[0].algo;
          });
    } else {
      constexpr int nalgo = CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT;
      int valid_algos;
      cudnnConvolutionBwdDataAlgoPerf_t algos[nalgo];
      CUDNN_ENFORCE(cudnnGetConvolutionBackwardDataAlgorithm_v7(
          cudnn_wrapper_.inline_cudnn_handle(),
          filter_desc_,
          bottom_desc_,
          conv_desc_,
          top_desc_,
          nalgo,
          &valid_algos,
          algos));
      bool found = false;
      for (int i = 0; i < valid_algos; i++) {
        auto a = algos[i];
        if (a.memory <= cudnn_ws_nbytes_limit_) {
          bwd_data_algo_ = a.algo;
          found = true;
          break;
        }
      }
      CAFFE_ENFORCE(found, "Unable to find algorithms for cuDNN backward data");
    }

    size_t bwd_data_ws_size;
    CUDNN_ENFORCE(cudnnGetConvolutionBackwardDataWorkspaceSize(
        cudnn_wrapper_.inline_cudnn_handle(),
        filter_desc_,
        bottom_desc_,
        conv_desc_,
        top_desc_,
        bwd_data_algo_,
        &bwd_data_ws_size));
    cudnn_ws_nbytes_ = bwd_data_ws_size;
    VLOG(1) << "CuDNN algorithm: " << bwd_data_algo_;
    VLOG(1) << "CuDNN workspace size: " << bwd_data_ws_size;
  }

  const T* X_data = X.template data<T>();
  const T* filter_data = filter.template data<T>();
  T* Y_data = Y->template mutable_data<T>();

  // Now, actually run the computation.
  // Filter
#if CUDNN_VERSION_MIN(7, 0, 0)
  cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
    CUDNN_ENFORCE(cudnnConvolutionBackwardData(
        state->cudnn_handle(),
        cudnnTypeWrapper<T>::kOne(),
        filter_desc_,
        filter_data,
        bottom_desc_,
        X_data,
        conv_desc_,
        bwd_data_algo_,
        state->workspace().get(cudnn_ws_nbytes_),
        cudnn_ws_nbytes_,
        cudnnTypeWrapper<T>::kZero(),
        top_desc_,
        Y_data));
  });
#else
  const int X_HxW = H * W;
  const int Y_HxW = H_out * W_out;
  const int group_offset_X =
      order_ == StorageOrder::NCHW ? M / group_ * X_HxW : M / group_;
  const int group_offset_Y =
      order_ == StorageOrder::NCHW ? C / group_ * Y_HxW : C / group_;
  const int group_offset_filter = filter.numel() / group_;
  for (int i = 0; i < group_; ++i) {
    cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
      CUDNN_ENFORCE(
          cudnnConvolutionBackwardData(state->cudnn_handle(),
                                       cudnnTypeWrapper<T>::kOne(),
                                       filter_desc_,
                                       filter_data + i * group_offset_filter,
                                       bottom_desc_,
                                       X_data + i * group_offset_X;
                                       conv_desc_,
                                       bwd_data_algo_,
                                       state->workspace().get(cudnn_ws_nbytes_),
                                       cudnn_ws_nbytes_,
                                       cudnnTypeWrapper<T_DX>::kZero(),
                                       top_desc_,
                                       Y_data + i * group_offset_Y));
    });
  }
#endif
  // Bias
  if (InputSize() == 3) {
    CUDNN_ENFORCE(cudnnAddTensor(
        cudnn_wrapper_.inline_cudnn_handle(),
        cudnnTypeWrapper<T>::kOne(),
        bias_desc_,
        Input(BIAS).template data<T>(),
        cudnnTypeWrapper<T>::kOne(),
        top_desc_for_bias_,
        Y->template mutable_data<T>()));
  }
  // Done.
  return true;
}

// TODO(Yangqing): a lot of the function contents are very similar. Consider
// consolidating them.
template <typename T>
bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
  const auto& X = Input(INPUT);
  const auto& filter = Input(FILTER);
  const auto& dY = Input(OUTPUT_GRAD);

  CAFFE_ENFORCE_EQ(X.dim(), 4);
  CAFFE_ENFORCE_EQ(filter.dim(), 4);
  int C = 0;
  switch (order_) {
    case StorageOrder::NHWC:
      C = filter.dim32(3) * group_;
      break;
    case StorageOrder::NCHW:
      C = filter.dim32(1) * group_;
      break;
    default:
      LOG(FATAL) << "Unknown storage order: " << order_;
  }

  int N = 0, M = 0, H = 0, W = 0, H_out = 0, W_out = 0;
  switch (order_) {
    case StorageOrder::NHWC:
      N = X.dim32(0);
      H = X.dim32(1);
      W = X.dim32(2);
      M = X.dim32(3);
      H_out = dY.dim32(1);
      W_out = dY.dim32(2);
      CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
      CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
      CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_w());
      CAFFE_ENFORCE_EQ(filter.dim32(3), C / group_);
      break;
    case StorageOrder::NCHW:
      N = X.dim32(0);
      M = X.dim32(1);
      H = X.dim32(2);
      W = X.dim32(3);
      H_out = dY.dim32(2);
      W_out = dY.dim32(3);
      CAFFE_ENFORCE_EQ(filter.dim32(1), C / group_);
      CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_h());
      CAFFE_ENFORCE_EQ(filter.dim32(3), kernel_w());
      break;
    default:
      LOG(FATAL) << "Unknown storage order: " << order_;
  }
  CAFFE_ENFORCE_EQ(M % group_, 0);

  // Since we only handle LegacyPadding::NOTSET, we don't need to
  // compute padding.
  auto* dfilter = Output(FILTER_GRAD, filter.sizes(), at::dtype<T>());

  // Set up the cudnn algorithms & workspace if necessary
  const bool input_changed = (X.sizes() != cudnn_input_dims_);
  const bool filter_changed = (filter.sizes() != cudnn_filter_dims_);
  if (input_changed || filter_changed) {
    VLOG(1) << "Changing the cudnn descriptor configurations.";
    if (input_changed) {
      cudnn_input_dims_ = X.sizes().vec();
      SetTensor4DDescriptorWithGroup(
          cudnnTypeWrapper<T>::type, N, M, H, W, &bottom_desc_);
    }
    if (filter_changed) {
      cudnn_filter_dims_ = filter.sizes().vec();
#if CUDNN_VERSION_MIN(7, 0, 0)
      const int MM = M;
#else
      const int MM = M / group_;
#endif
      CUDNN_ENFORCE(cudnnSetFilter4dDescriptor(
          filter_desc_,
          cudnnTypeWrapper<T>::type,
          GetCudnnTensorFormat(order_),
          MM,
          C / group_,
          kernel_h(),
          kernel_w()));
      if (!no_bias_) {
        CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
            bias_desc_,
            GetCudnnTensorFormat(order_),
            cudnnTypeWrapper<T>::type,
            1,
            C,
            1,
            1));
      }
    }
    // Set the output
    SetTensor4DDescriptorWithGroup(
        cudnnTypeWrapper<T>::type, N, C, H_out, W_out, &top_desc_);
    if (!no_bias_) {
      CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
          top_desc_for_bias_,
          GetCudnnTensorFormat(order_),
          cudnnTypeWrapper<T>::type,
          N,
          C,
          H_out,
          W_out));
    }

    // Set the convolution descriptor
    CAFFE_ENFORCE_EQ(
        pad_t(),
        pad_b(),
        "The current padding scheme leads to unequal padding on the top and "
        "bottom, which is not supported by cudnn.");
    CAFFE_ENFORCE_EQ(
        pad_l(),
        pad_r(),
        "The current padding scheme leads to unequal padding on the left "
        "and right, which is not supported by cudnn.");
#if CUDNN_VERSION_MIN(6, 0, 0)
    CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
        conv_desc_,
        pad_t(),
        pad_l(),
        stride_h(),
        stride_w(),
        1,
        1,
        CUDNN_CROSS_CORRELATION,
        cudnnTypeWrapper<T>::type));
#else
    CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
        conv_desc_,
        pad_t(),
        pad_l(),
        stride_h(),
        stride_w(),
        1,
        1,
        CUDNN_CROSS_CORRELATION));
#endif
#if CUDNN_VERSION_MIN(7, 0, 0)
    // enable TensorCore math if desired
    enable_tensor_core_ &= TensorCoreAvailable();
    if (enable_tensor_core_) {
      CUDNN_ENFORCE(
          cudnnSetConvolutionMathType(conv_desc_, CUDNN_TENSOR_OP_MATH));
    }
    // set cuDNN groups if appropriate
    CUDNN_CHECK(cudnnSetConvolutionGroupCount(conv_desc_, group_));
#endif
    if (force_algo_[ALGO_WGRAD] >= 0) {
      bwd_filter_algo_ =
          (cudnnConvolutionBwdFilterAlgo_t)force_algo_[ALGO_WGRAD];
    } else if (deterministic_) {
      algo_ = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
      bwd_filter_algo_ = CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1;
    } else if (exhaustive_search_) {
      bwd_filter_algo_ =
          filter_algo_cache_.getAlgorithm(X.sizes(), filter.sizes(), 0, [&]() {
            LOG(INFO) << "CUDNN Convolution bwd: doing exhaustive search.";
            // When we do an exhaustive search, we will ignore the workspace
            // size
            // limit and simply go for the fastest algorithm. If you happen to
            // run
            // out of memory later, you will be on your own...
            int returned_algo_count;
            // We clean up the current workspace memory so that the forward
            // algorithm
            // is free to allocate memory.
            // Actually run the search.
            std::array<
                cudnnConvolutionBwdFilterAlgoPerf_t,
                kNUM_CUDNN_BWD_FILTER_ALGS>
                filter_perf_stat;

            cudnn_wrapper_.with_cudnn_state(
                cudnn_state_, [&](CuDNNState* state) {
                  state->workspace().reset();
                  CUDNN_ENFORCE(cudnnFindConvolutionBackwardFilterAlgorithm(
                      state->cudnn_handle(),
                      top_desc_,
                      bottom_desc_,
                      conv_desc_,
                      filter_desc_,
                      kNUM_CUDNN_BWD_FILTER_ALGS,
                      &returned_algo_count,
                      filter_perf_stat.data()));
                });
            LogCuDNNPerfStats(filter_perf_stat, returned_algo_count);
            return filter_perf_stat[0].algo;
          });

      algo_ =
          forward_algo_cache_.getAlgorithm(X.sizes(), filter.sizes(), 0, [&]() {
            int returned_algo_count;
            std::array<cudnnConvolutionFwdAlgoPerf_t, kNUM_CUDNN_FWD_ALGS>
                fwd_perf_stat;
            cudnn_wrapper_.with_cudnn_state(
                cudnn_state_, [&](CuDNNState* state) {
                  state->workspace().reset();
                  CUDNN_ENFORCE(cudnnFindConvolutionForwardAlgorithm(
                      state->cudnn_handle(),
                      top_desc_,
                      filter_desc_,
                      conv_desc_,
                      bottom_desc_,
                      kNUM_CUDNN_BWD_DATA_ALGS,
                      &returned_algo_count,
                      fwd_perf_stat.data()));
                });

            LogCuDNNPerfStats(fwd_perf_stat, returned_algo_count);
            return fwd_perf_stat[0].algo;
          });
    } else {
      // choose backward algorithm for filter
      {
      constexpr int nalgo = CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT;
      int valid_algos;
      cudnnConvolutionBwdFilterAlgoPerf_t algos[nalgo];
      CUDNN_ENFORCE(cudnnGetConvolutionBackwardFilterAlgorithm_v7(
          cudnn_wrapper_.inline_cudnn_handle(),
          top_desc_,
          bottom_desc_,
          conv_desc_,
          filter_desc_,
          nalgo,
          &valid_algos,
          algos));
      bool found = false;
      for (int i = 0; i < valid_algos; i++) {
        auto a = algos[i];
        if (a.memory <= cudnn_ws_nbytes_limit_) {
          bwd_filter_algo_ = a.algo;
          found = true;
          break;
        }
      }
      CAFFE_ENFORCE(found, "Unable to find algorithms for cuDNN backward filter");
      }
      // choose backward algo for data
      {
      constexpr int nalgo = CUDNN_CONVOLUTION_FWD_ALGO_COUNT;
      int valid_algos;
      cudnnConvolutionFwdAlgoPerf_t algos[nalgo];
      CUDNN_ENFORCE(cudnnGetConvolutionForwardAlgorithm_v7(
          cudnn_wrapper_.inline_cudnn_handle(),
          top_desc_,
          filter_desc_,
          conv_desc_,
          bottom_desc_,
          nalgo,
          &valid_algos,
          algos));
      bool found = false;
      for (int i = 0; i < valid_algos; i++) {
        auto a = algos[i];
        if (a.memory <= cudnn_ws_nbytes_limit_) {
          algo_ = a.algo;
          found = true;
          break;
        }
      }
      CAFFE_ENFORCE(found, "Unable to find algorithms for cuDNN forward");
      }
    }
    // get workspace for backwards filter algorithm
    size_t bwd_filter_ws_size, fwd_ws_size;
    CUDNN_ENFORCE(cudnnGetConvolutionBackwardFilterWorkspaceSize(
        cudnn_wrapper_.inline_cudnn_handle(),
        top_desc_,
        bottom_desc_,
        conv_desc_,
        filter_desc_,
        bwd_filter_algo_,
        &bwd_filter_ws_size));
    // get workspace for backwards data algorithm
    CUDNN_ENFORCE(cudnnGetConvolutionForwardWorkspaceSize(
        cudnn_wrapper_.inline_cudnn_handle(),
        top_desc_,
        filter_desc_,
        conv_desc_,
        bottom_desc_,
        algo_,
        &fwd_ws_size));
    cudnn_ws_nbytes_ = std::max(bwd_filter_ws_size, fwd_ws_size);

    VLOG(1) << "CuDNN bwd algorithm: " << bwd_filter_algo_ << ", " << algo_;
    VLOG(1) << "CuDNN workspace size: " << cudnn_ws_nbytes_;
  }

  // Now, actually run the computation.
  if (!no_bias_) {
    auto* dbias = Output(BIAS_OR_INPUT_GRAD, {C}, at::dtype<T>());
    CUDNN_ENFORCE(cudnnConvolutionBackwardBias(
        cudnn_wrapper_.inline_cudnn_handle(),
        cudnnTypeWrapper<T>::kOne(),
        top_desc_for_bias_,
        dY.template data<T>(),
        cudnnTypeWrapper<T>::kZero(),
        bias_desc_,
        dbias->template mutable_data<T>()));
  }

#if CUDNN_VERSION_MIN(7, 0, 0)
  cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
    CUDNN_ENFORCE(cudnnConvolutionBackwardFilter(
        state->cudnn_handle(),
        cudnnTypeWrapper<T>::kOne(),
        top_desc_,
        dY.template data<T>(),
        bottom_desc_,
        X.template data<T>(),
        conv_desc_,
        bwd_filter_algo_,
        state->workspace().get(cudnn_ws_nbytes_),
        cudnn_ws_nbytes_,
        cudnnTypeWrapper<T>::kZero(),
        filter_desc_,
        dfilter->template mutable_data<T>()));

    if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
      // Compute the gradient w.r.t. the input.
      auto* dX = Output(
          no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD,
          X.sizes(),
          at::dtype<T>());
      CUDNN_ENFORCE(cudnnConvolutionForward(
          state->cudnn_handle(),
          cudnnTypeWrapper<T>::kOne(),
          top_desc_,
          dY.template data<T>(),
          filter_desc_,
          filter.template data<T>(),
          conv_desc_,
          algo_,
          state->workspace().get(cudnn_ws_nbytes_),
          cudnn_ws_nbytes_,
          cudnnTypeWrapper<T>::kZero(),
          bottom_desc_,
          dX->template mutable_data<T>()));
    }
  });
#else
  const int X_HxW = H * W;
  const int Y_HxW = H_out * W_out;
  const int group_offset_X =
      order_ == StorageOrder::NCHW ? M / group_ * X_HxW : M / group_;
  const int group_offset_Y =
      order_ == StorageOrder::NCHW ? C / group_ * Y_HxW : C / group_;
  const int group_offset_filter = filter.numel() / group_;
  for (int i = 0; i < group_; ++i) {
    cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
      CUDNN_ENFORCE(cudnnConvolutionBackwardFilter(
          state->cudnn_handle(),
          cudnnTypeWrapper<T>::kOne(),
          top_desc_,
          dY.template data<T>() + i * group_offset_Y,
          bottom_desc_,
          X.template data<T>() + i * group_offset_X,
          conv_desc_,
          bwd_filter_algo_,
          state->workspace().get(cudnn_ws_nbytes_),
          cudnn_ws_nbytes_,
          cudnnTypeWrapper<T>::kZero(),
          filter_desc_,
          dfilter->template mutable_data<T>() + i * group_offset_filter));
      if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
        // Compute the gradient w.r.t. the input.
        auto* dX = Output(
            no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD,
            X.sizes(),
            at::dtype<T>());
        cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
          CUDNN_ENFORCE(cudnnConvolutionForward(
              state->cudnn_handle(),
              cudnnTypeWrapper<T>::kOne(),
              top_desc_,
              dY.template data<T>() + i * group_offset_Y,
              filter_desc_,
              filter.template data<T>() + i * group_offset_filter,
              conv_desc_,
              algo_,
              state->workspace().get(cudnn_ws_nbytes_),
              cudnn_ws_nbytes_,
              cudnnTypeWrapper<T>::kZero(),
              bottom_desc_,
              dX->template mutable_data<T>() + i * group_offset_X));
        });
      }
  }
#endif
  return true;
}

REGISTER_CUDNN_OPERATOR(ConvTranspose, CudnnConvTransposeOp<float>);
REGISTER_CUDNN_OPERATOR(
    ConvTransposeGradient,
    CudnnConvTransposeGradientOp<float>);

} // namespace caffe2