#ifdef USE_CUDNN

#include <vector>

#include "thrust/device_vector.h"

#include "caffe/layers/cudnn_softmax_layer.hpp"

namespace caffe {

template <typename Dtype>

void CuDNNSoftmaxLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

SoftmaxLayer<Dtype>::LayerSetUp(bottom, top);

// Initialize CUDNN.

CUDNN_CHECK(cudnnCreate(&handle_));

cudnn::createTensor4dDesc<Dtype>(&bottom_desc_);

cudnn::createTensor4dDesc<Dtype>(&top_desc_);

handles_setup_ = true;

}

template <typename Dtype>

void CuDNNSoftmaxLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

SoftmaxLayer<Dtype>::Reshape(bottom, top);

int N = this->outer_num_;

int K = bottom[0]->shape(this->softmax_axis_);

int H = this->inner_num_;

int W = 1;

cudnn::setTensor4dDesc<Dtype>(&bottom_desc_, N, K, H, W);

cudnn::setTensor4dDesc<Dtype>(&top_desc_, N, K, H, W);

}

template <typename Dtype>

CuDNNSoftmaxLayer<Dtype>::~CuDNNSoftmaxLayer() {

// Check that handles have been setup before destroying.

if (!handles_setup_) { return; }

cudnnDestroyTensorDescriptor(bottom_desc_);

cudnnDestroyTensorDescriptor(top_desc_);

cudnnDestroy(handle_);

}

INSTANTIATE_CLASS(CuDNNSoftmaxLayer);

} // namespace caffe

#endif