diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc index 9e59fd6639..24d1f4e3c7 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc +++ b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc @@ -119,8 +119,14 @@ Status CudnnRnnBase::CacheCudnnRnnWeights(const OpKernelInfo& info) { if (get_W && get_R) { CudnnRNN tmp_rnn_desc; - ORT_RETURN_IF_ERROR(tmp_rnn_desc.Set(CudnnHandle(), hidden_size_, RNN_NUM_LAYERS, cudnn_dropout_desc_, - cudnn_direction_mode_, rnn_mode_, CudnnTensor::GetDataType())); + ORT_RETURN_IF_ERROR(tmp_rnn_desc.Set(CudnnHandle(), + hidden_size_, + RNN_NUM_LAYERS, + cudnn_dropout_desc_, + cudnn_direction_mode_, + rnn_mode_, + CudnnTensor::GetDataType(), + GetDeviceProp())); if (get_B) { ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_, tmp_rnn_desc)); } else { @@ -211,8 +217,14 @@ Status CudnnRnnBase::ComputeInternal(OpKernelContext* ctx) const { const int32_t* sequence_lens_data = (sequence_lens == nullptr) ? nullptr : sequence_lens->template Data(); CudnnRNN rnn_desc; - ORT_RETURN_IF_ERROR(rnn_desc.Set(CudnnHandle(), hidden_size_, RNN_NUM_LAYERS, cudnn_dropout_desc_, - cudnn_direction_mode_, rnn_mode_, CudnnTensor::GetDataType())); + ORT_RETURN_IF_ERROR(rnn_desc.Set(CudnnHandle(), + hidden_size_, + RNN_NUM_LAYERS, + cudnn_dropout_desc_, + cudnn_direction_mode_, + rnn_mode_, + CudnnTensor::GetDataType(), + GetDeviceProp())); // Prepare the weight data IAllocatorUniquePtr w_data; diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h index 41a6c94c8c..5281904a2b 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h +++ b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h @@ -38,7 +38,7 @@ class CudnnRNN { Status Set(const cudnnHandle_t& cudnnHandle, int64_t hidden_size, int num_layers, cudnnDropoutDescriptor_t cudnn_dropout_desc, cudnnDirectionMode_t cudnn_direction_model, - cudnnRNNMode_t rnn_mode, cudnnDataType_t dataType) { + cudnnRNNMode_t rnn_mode, cudnnDataType_t dataType, const cudaDeviceProp& prop) { if (!cudnn_rnn_desc_) CUDNN_RETURN_IF_ERROR(cudnnCreateRNNDescriptor(&cudnn_rnn_desc_)); @@ -53,6 +53,10 @@ class CudnnRNN { CUDNN_RNN_ALGO_STANDARD, //CUDNN_RNN_ALGO_PERSIST_STATIC, CUDNN_RNN_ALGO_PERSIST_DYNAMIC dataType)); + if (prop.major >= 7 && dataType == CUDNN_DATA_HALF) { + cudnnSetRNNMatrixMathType(cudnn_rnn_desc_, CUDNN_TENSOR_OP_MATH); + } + return Status::OK(); } @@ -91,7 +95,7 @@ class CudnnRnnBase : public CudaKernel { rnn_mode_ = CUDNN_LSTM; weight_cached_ = false; w_data_cache_ = nullptr; - + size_t state_size; cudnn_dropout_desc_.CreateDescriptorIfNeeded(); cudnn_dropout_desc_.GetCudnnDropoutStatesSize(CudnnHandle(), state_size);