#ifndef CAFFE2_CORE_COMMON_CUDNN_H_
|
#define CAFFE2_CORE_COMMON_CUDNN_H_
|
|
#include <array>
|
#include <mutex>
|
|
#include "caffe2/core/common.h"
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/logging.h"
|
#include "caffe2/core/types.h"
|
|
#ifndef CAFFE2_USE_CUDNN
|
#error("This Caffe2 install is not built with cudnn, so you should not include this file.");
|
#endif
|
|
#include <cudnn.h>
|
|
static_assert(
|
CUDNN_VERSION >= 5000,
|
"Caffe2 requires cudnn version 5.0 or above.");
|
|
#if CUDNN_VERSION < 6000
|
#pragma message "CUDNN version under 6.0 is supported at best effort."
|
#pragma message "We strongly encourage you to move to 6.0 and above."
|
#pragma message "This message is intended to annoy you enough to update."
|
#endif // CUDNN_VERSION < 6000
|
|
#define CUDNN_VERSION_MIN(major, minor, patch) \
|
(CUDNN_VERSION >= ((major) * 1000 + (minor) * 100 + (patch)))
|
|
namespace caffe2 {
|
|
namespace internal {
|
/**
|
* A helper function to obtain cudnn error strings.
|
*/
|
inline const char* cudnnGetErrorString(cudnnStatus_t status) {
|
switch (status) {
|
case CUDNN_STATUS_SUCCESS:
|
return "CUDNN_STATUS_SUCCESS";
|
case CUDNN_STATUS_NOT_INITIALIZED:
|
return "CUDNN_STATUS_NOT_INITIALIZED";
|
case CUDNN_STATUS_ALLOC_FAILED:
|
return "CUDNN_STATUS_ALLOC_FAILED";
|
case CUDNN_STATUS_BAD_PARAM:
|
return "CUDNN_STATUS_BAD_PARAM";
|
case CUDNN_STATUS_INTERNAL_ERROR:
|
return "CUDNN_STATUS_INTERNAL_ERROR";
|
case CUDNN_STATUS_INVALID_VALUE:
|
return "CUDNN_STATUS_INVALID_VALUE";
|
case CUDNN_STATUS_ARCH_MISMATCH:
|
return "CUDNN_STATUS_ARCH_MISMATCH";
|
case CUDNN_STATUS_MAPPING_ERROR:
|
return "CUDNN_STATUS_MAPPING_ERROR";
|
case CUDNN_STATUS_EXECUTION_FAILED:
|
return "CUDNN_STATUS_EXECUTION_FAILED";
|
case CUDNN_STATUS_NOT_SUPPORTED:
|
return "CUDNN_STATUS_NOT_SUPPORTED";
|
case CUDNN_STATUS_LICENSE_ERROR:
|
return "CUDNN_STATUS_LICENSE_ERROR";
|
default:
|
return "Unknown cudnn error number";
|
}
|
}
|
} // namespace internal
|
|
// A macro that wraps around a cudnn statement so we can check if the cudnn
|
// execution finishes or not.
|
#define CUDNN_ENFORCE(condition) \
|
do { \
|
cudnnStatus_t status = condition; \
|
CAFFE_ENFORCE_EQ( \
|
status, \
|
CUDNN_STATUS_SUCCESS, \
|
", Error at: ", \
|
__FILE__, \
|
":", \
|
__LINE__, \
|
": ", \
|
::caffe2::internal::cudnnGetErrorString(status)); \
|
} while (0)
|
#define CUDNN_CHECK(condition) \
|
do { \
|
cudnnStatus_t status = condition; \
|
CHECK(status == CUDNN_STATUS_SUCCESS) \
|
<< ::caffe2::internal::cudnnGetErrorString(status); \
|
} while (0)
|
|
// report the version of cuDNN Caffe2 was compiled with
|
inline size_t cudnnCompiledVersion() {
|
return CUDNN_VERSION;
|
}
|
// report the runtime version of cuDNN
|
inline size_t cudnnRuntimeVersion() {
|
return cudnnGetVersion();
|
}
|
|
// Check compatibility of compiled and runtime cuDNN versions
|
inline void CheckCuDNNVersions() {
|
// Version format is major*1000 + minor*100 + patch
|
// If compiled with version < 7, major, minor and patch must all match
|
// If compiled with version >= 7, then either
|
// runtime_version > compiled_version
|
// major and minor match
|
bool version_match = cudnnCompiledVersion() == cudnnRuntimeVersion();
|
bool compiled_with_7 = cudnnCompiledVersion() >= 7000;
|
bool backwards_compatible_7 = compiled_with_7 && cudnnRuntimeVersion() >= cudnnCompiledVersion();
|
bool patch_compatible = compiled_with_7 && (cudnnRuntimeVersion() / 100) == (cudnnCompiledVersion() / 100);
|
CAFFE_ENFORCE(version_match || backwards_compatible_7 || patch_compatible,
|
"cuDNN compiled (", cudnnCompiledVersion(), ") and "
|
"runtime (", cudnnRuntimeVersion(), ") versions mismatch");
|
}
|
|
/**
|
* cudnnTypeWrapper is a wrapper class that allows us to refer to the cudnn type
|
* in a template function. The class is specialized explicitly for different
|
* data types below.
|
*/
|
template <typename T>
|
class cudnnTypeWrapper;
|
|
template <>
|
class cudnnTypeWrapper<float> {
|
public:
|
static const cudnnDataType_t type = CUDNN_DATA_FLOAT;
|
typedef const float ScalingParamType;
|
typedef float BNParamType;
|
static ScalingParamType* kOne() {
|
static ScalingParamType v = 1.0;
|
return &v;
|
}
|
static const ScalingParamType* kZero() {
|
static ScalingParamType v = 0.0;
|
return &v;
|
}
|
};
|
|
#if CUDNN_VERSION_MIN(6, 0, 0)
|
template <>
|
class cudnnTypeWrapper<int> {
|
public:
|
static const cudnnDataType_t type = CUDNN_DATA_INT32;
|
typedef const int ScalingParamType;
|
typedef int BNParamType;
|
static ScalingParamType* kOne() {
|
static ScalingParamType v = 1;
|
return &v;
|
}
|
static const ScalingParamType* kZero() {
|
static ScalingParamType v = 0;
|
return &v;
|
}
|
};
|
#endif // CUDNN_VERSION_MIN(6, 0, 0)
|
|
template <>
|
class cudnnTypeWrapper<double> {
|
public:
|
static const cudnnDataType_t type = CUDNN_DATA_DOUBLE;
|
typedef const double ScalingParamType;
|
typedef double BNParamType;
|
static ScalingParamType* kOne() {
|
static ScalingParamType v = 1.0;
|
return &v;
|
}
|
static ScalingParamType* kZero() {
|
static ScalingParamType v = 0.0;
|
return &v;
|
}
|
};
|
|
template <>
|
class cudnnTypeWrapper<at::Half> {
|
public:
|
static const cudnnDataType_t type = CUDNN_DATA_HALF;
|
typedef const float ScalingParamType;
|
typedef float BNParamType;
|
static ScalingParamType* kOne() {
|
static ScalingParamType v = 1.0;
|
return &v;
|
}
|
static ScalingParamType* kZero() {
|
static ScalingParamType v = 0.0;
|
return &v;
|
}
|
};
|
|
/**
|
* A wrapper function to convert the Caffe storage order to cudnn storage order
|
* enum values.
|
*/
|
inline cudnnTensorFormat_t GetCudnnTensorFormat(const StorageOrder& order) {
|
switch (order) {
|
case StorageOrder::NHWC:
|
return CUDNN_TENSOR_NHWC;
|
case StorageOrder::NCHW:
|
return CUDNN_TENSOR_NCHW;
|
default:
|
LOG(FATAL) << "Unknown cudnn equivalent for order: " << order;
|
}
|
// Just to suppress compiler warnings
|
return CUDNN_TENSOR_NCHW;
|
}
|
|
/**
|
* cudnnTensorDescWrapper is the placeholder that wraps around a
|
* cudnnTensorDescriptor_t, allowing us to do descriptor change as-needed during
|
* runtime.
|
*/
|
class cudnnTensorDescWrapper {
|
public:
|
cudnnTensorDescWrapper() {
|
CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&desc_));
|
}
|
~cudnnTensorDescWrapper() noexcept {
|
CUDNN_CHECK(cudnnDestroyTensorDescriptor(desc_));
|
}
|
|
inline cudnnTensorDescriptor_t Descriptor(
|
const cudnnTensorFormat_t format,
|
const cudnnDataType_t type,
|
const vector<int>& dims,
|
bool* changed) {
|
if (type_ == type && format_ == format && dims_ == dims) {
|
// if not changed, simply return the current descriptor.
|
if (changed)
|
*changed = false;
|
return desc_;
|
}
|
CAFFE_ENFORCE_EQ(
|
dims.size(), 4, "Currently only 4-dimensional descriptor supported.");
|
format_ = format;
|
type_ = type;
|
dims_ = dims;
|
CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
|
desc_,
|
format,
|
type,
|
dims_[0],
|
(format == CUDNN_TENSOR_NCHW ? dims_[1] : dims_[3]),
|
(format == CUDNN_TENSOR_NCHW ? dims_[2] : dims_[1]),
|
(format == CUDNN_TENSOR_NCHW ? dims_[3] : dims_[2])));
|
if (changed)
|
*changed = true;
|
return desc_;
|
}
|
|
template <typename T>
|
inline cudnnTensorDescriptor_t Descriptor(
|
const StorageOrder& order,
|
const vector<int>& dims) {
|
return Descriptor(
|
GetCudnnTensorFormat(order), cudnnTypeWrapper<T>::type, dims, nullptr);
|
}
|
|
private:
|
cudnnTensorDescriptor_t desc_;
|
cudnnTensorFormat_t format_;
|
cudnnDataType_t type_;
|
vector<int> dims_;
|
C10_DISABLE_COPY_AND_ASSIGN(cudnnTensorDescWrapper);
|
};
|
|
class cudnnFilterDescWrapper {
|
public:
|
cudnnFilterDescWrapper() {
|
CUDNN_ENFORCE(cudnnCreateFilterDescriptor(&desc_));
|
}
|
~cudnnFilterDescWrapper() noexcept {
|
CUDNN_CHECK(cudnnDestroyFilterDescriptor(desc_));
|
}
|
|
inline cudnnFilterDescriptor_t Descriptor(
|
const StorageOrder& order,
|
const cudnnDataType_t type,
|
const vector<int>& dims,
|
bool* changed) {
|
if (type_ == type && order_ == order && dims_ == dims) {
|
// if not changed, simply return the current descriptor.
|
if (changed)
|
*changed = false;
|
return desc_;
|
}
|
CAFFE_ENFORCE_EQ(
|
dims.size(), 4, "Currently only 4-dimensional descriptor supported.");
|
order_ = order;
|
type_ = type;
|
dims_ = dims;
|
CUDNN_ENFORCE(cudnnSetFilter4dDescriptor(
|
desc_,
|
type,
|
GetCudnnTensorFormat(order),
|
dims_[0],
|
// TODO - confirm that this is correct for NHWC
|
(order == StorageOrder::NCHW ? dims_[1] : dims_[3]),
|
(order == StorageOrder::NCHW ? dims_[2] : dims_[1]),
|
(order == StorageOrder::NCHW ? dims_[3] : dims_[2])));
|
if (changed)
|
*changed = true;
|
return desc_;
|
}
|
|
template <typename T>
|
inline cudnnFilterDescriptor_t Descriptor(
|
const StorageOrder& order,
|
const vector<int>& dims) {
|
return Descriptor(order, cudnnTypeWrapper<T>::type, dims, nullptr);
|
}
|
|
private:
|
cudnnFilterDescriptor_t desc_;
|
StorageOrder order_;
|
cudnnDataType_t type_;
|
vector<int> dims_;
|
C10_DISABLE_COPY_AND_ASSIGN(cudnnFilterDescWrapper);
|
};
|
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_CORE_COMMON_CUDNN_H_
|
