saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

enable tensor core for rnn (#3937)

Browse source
This commit is contained in:
Yufeng Li 2020-05-18 16:24:18 -07:00 committed by GitHub
parent c6a94f95cf
commit 29c39f867c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc
View file

@ -119,8 +119,14 @@ Status CudnnRnnBase<T>::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<CudaT>()));
ORT_RETURN_IF_ERROR(tmp_rnn_desc.Set(CudnnHandle(),
hidden_size_,
RNN_NUM_LAYERS,
cudnn_dropout_desc_,
cudnn_direction_mode_,
rnn_mode_,
CudnnTensor::GetDataType<CudaT>(),
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<T>::ComputeInternal(OpKernelContext* ctx) const {
const int32_t* sequence_lens_data = (sequence_lens == nullptr) ? nullptr : sequence_lens->template Data<int32_t>();
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<CudaT>()));
ORT_RETURN_IF_ERROR(rnn_desc.Set(CudnnHandle(),
hidden_size_,
RNN_NUM_LAYERS,
cudnn_dropout_desc_,
cudnn_direction_mode_,
rnn_mode_,
CudnnTensor::GetDataType<CudaT>(),
GetDeviceProp()));
// Prepare the weight data
IAllocatorUniquePtr<void> w_data;

8
onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h
View file

@ -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);