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

Fix race condition issue in RNN/LSTM/GRU (#1544)

Browse source 
Fix race condition issue in RNN/LSTM/GRU.

Description:
The filter_desc and rnn_desc could also be changed in compute which could be in multi-thread. It will cause race condition issue.

Fix:
create temperate cudnn descriptors
cache cudnn_dropout_desc_ which won't change
This commit is contained in:
Hector Li 2019-08-08 14:18:41 -07:00 committed by GitHub
parent 6e430c0526
commit 38d78542c3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 94 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/cuda/cudnn_common.h
View file

@ -97,13 +97,13 @@ class CudnnDropout final {
return dropout_desc_;
}
private:
Status CreateDescriptorIfNeeded() {
if (!dropout_desc_)
CUDNN_RETURN_IF_ERROR(cudnnCreateDropoutDescriptor(&dropout_desc_));
return Status::OK();
}
private:
cudnnDropoutDescriptor_t dropout_desc_;
};

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

@ -15,7 +15,7 @@ void CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
const cudnnFilterDescriptor_t filter_desc,
const void* w_data,
const void* reorganized_w_data,
const int lin_layer_id,
const T* pos,
int& offset,
@ -27,9 +27,9 @@ void CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
T* mem_offset;
if (is_matrix) {
cudnnGetRNNLinLayerMatrixParams(handle, rnn_desc, pseudo_layer, x_desc, w_desc, w_data, lin_layer_id, filter_desc, (void**)&mem_offset);
cudnnGetRNNLinLayerMatrixParams(handle, rnn_desc, pseudo_layer, x_desc, w_desc, reorganized_w_data, lin_layer_id, filter_desc, (void**)&mem_offset);
} else {
cudnnGetRNNLinLayerBiasParams(handle, rnn_desc, pseudo_layer, x_desc, w_desc, w_data, lin_layer_id, filter_desc, (void**)&mem_offset);
cudnnGetRNNLinLayerBiasParams(handle, rnn_desc, pseudo_layer, x_desc, w_desc, reorganized_w_data, lin_layer_id, filter_desc, (void**)&mem_offset);
}
cudnnGetFilterNdDescriptor(filter_desc, 3, &dt, &tf, &numDims, matDims.data());
@ -42,25 +42,25 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
const cudnnRNNDescriptor_t rnn_desc,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
void* w_data,
void* reorganized_w_data,
const T* W_data,
const T* R_data,
const T* B_data) const {
//Onnx only support 1 layer
int w_offset = 0;
int r_offset = 0;
int bias_offset = 0;
for (int layer = 0; layer < num_layers_ * num_directions_; ++layer) {
CudnnFilterDescriptor filter_desc;
for (int layer = 0; layer < RNN_NUM_LAYERS * num_directions_; ++layer) {
for (size_t idx = 0; idx < W_lin_layer_id_.size(); ++idx) {
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc_, w_data, W_lin_layer_id_[idx], W_data, w_offset, true);
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc, reorganized_w_data, W_lin_layer_id_[idx], W_data, w_offset, true);
if (B_data != nullptr) {
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc_, w_data, W_lin_layer_id_[idx], B_data, bias_offset, false);
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc, reorganized_w_data, W_lin_layer_id_[idx], B_data, bias_offset, false);
}
}
for (size_t idx = 0; idx < R_lin_layer_id_.size(); ++idx) {
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc_, w_data, R_lin_layer_id_[idx], R_data, r_offset, true);
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc, reorganized_w_data, R_lin_layer_id_[idx], R_data, r_offset, true);
if (B_data != nullptr) {
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc_, w_data, R_lin_layer_id_[idx], B_data, bias_offset, false);
SetWeightBias(cudnn_handle, rnn_desc, layer, x_desc, w_desc, filter_desc, reorganized_w_data, R_lin_layer_id_[idx], B_data, bias_offset, false);
}
}
}
@ -68,34 +68,11 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
return Status::OK();
}
template <typename T>
Status CudnnRnnBase<T>::SetCudnnRnnDesc() {
typedef typename ToCudaType<T>::MappedType CudaT;
cudnnDirectionMode_t cudnn_direction = CUDNN_UNIDIRECTIONAL;
if (direction_ == "bidirectional") {
cudnn_direction = CUDNN_BIDIRECTIONAL;
} else if (direction_ == "forward") {
cudnn_direction = CUDNN_UNIDIRECTIONAL;
} else if (direction_ == "reverse") {
cudnn_direction = CUDNN_UNIDIRECTIONAL;
// need to reverse data
reverse_ = true;
}
cudnn_dropout_desc_.GetCudnnDropoutStatesSize(CudnnHandle(), state_size_);
state_buffer_ = GetScratchBuffer<void>(state_size_);
cudnn_dropout_desc_.Set(CudnnHandle(), state_buffer_.get(), state_size_);
ORT_RETURN_IF_ERROR(rnn_desc_.Set(CudnnHandle(), hidden_size_, num_layers_, cudnn_dropout_desc_,
cudnn_direction, rnn_mode_, CudnnTensor::GetDataType<CudaT>()));
return Status::OK();
}
template <typename T>
Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, const Tensor* B,
IAllocatorUniquePtr<void>& target_w_data,
CudnnFilterDescriptor& target_w_desc) const {
IAllocatorUniquePtr<void>& reorganized_w_data,
CudnnFilterDescriptor& target_w_desc,
CudnnRNN& rnn_desc) const {
typedef typename ToCudaType<T>::MappedType CudaT;
int64_t input_size = W->Shape()[2];
// RNN W[num_directions_, hidden_size_, input_size]
@ -117,20 +94,21 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
fake_x_desc.Set(fake_dims_x, CudnnTensor::GetDataType<CudaT>());
// Prepare the weight data
target_w_data = GetScratchBuffer<void>(w_size * sizeof(T));
reorganized_w_data = GetScratchBuffer<void>(w_size * sizeof(T));
const T* W_data = W->template Data<T>();
const T* R_data = R->template Data<T>();
const T* B_data = B == nullptr ? nullptr : B->template Data<T>();
ORT_RETURN_IF_ERROR(SetCudnnRnnWeightBias(CudnnHandle(), rnn_desc_, fake_x_desc, target_w_desc,
target_w_data.get(), W_data, R_data, B_data));
ORT_RETURN_IF_ERROR(SetCudnnRnnWeightBias(CudnnHandle(), rnn_desc, fake_x_desc, target_w_desc,
reorganized_w_data.get(), W_data, R_data, B_data));
return Status::OK();
}
template <typename T>
Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
typedef typename ToCudaType<T>::MappedType CudaT;
// Cache the weight
const Tensor* W;
const Tensor* R;
@ -140,10 +118,13 @@ Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
bool get_B = info.TryGetConstantInput(RNN_Input_Index::B, &B);
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>()));
if (get_B) {
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_));
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_, tmp_rnn_desc));
} else {
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, nullptr, w_data_cache_, w_desc_cache_));
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, nullptr, w_data_cache_, w_desc_cache_, tmp_rnn_desc));
}
weight_cached_ = true;
}
@ -173,7 +154,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
// optional outputs
std::vector<int64_t> dims_Y({seq_length, num_directions_, batch_size, hidden_size_});
std::vector<int64_t> dims_hxy({num_layers_ * num_directions_, batch_size, hidden_size_});
std::vector<int64_t> dims_hxy({RNN_NUM_LAYERS * num_directions_, batch_size, hidden_size_});
std::vector<int64_t> dims_yc{num_directions_, batch_size, hidden_size_};
Tensor* Y = ctx->Output(Output_Index::Y, dims_Y);
Tensor* Y_h = ctx->Output(Output_Index::Y_h, dims_hxy);
@ -198,16 +179,6 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
ORT_RETURN_IF_ERROR(y_h_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(y_c_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
// Prepare the weight data
IAllocatorUniquePtr<void> w_data;
CudnnFilterDescriptor w_desc;
if (!weight_cached_) {
const Tensor& W = *ctx->Input<Tensor>(RNN_Input_Index::W);
const Tensor& R = *ctx->Input<Tensor>(RNN_Input_Index::R);
const Tensor* B = ctx->Input<Tensor>(RNN_Input_Index::B);
ReorganizeWeights(&W, &R, B, w_data, w_desc);
}
IAllocatorUniquePtr<T> x_reversed_data;
const T* x_data = X->template Data<T>();
if (reverse_) {
@ -239,16 +210,30 @@ 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>()));
// Prepare the weight data
IAllocatorUniquePtr<void> w_data;
CudnnFilterDescriptor w_desc;
if (!weight_cached_) {
const Tensor& W = *ctx->Input<Tensor>(RNN_Input_Index::W);
const Tensor& R = *ctx->Input<Tensor>(RNN_Input_Index::R);
const Tensor* B = ctx->Input<Tensor>(RNN_Input_Index::B);
ReorganizeWeights(&W, &R, B, w_data, w_desc, rnn_desc);
}
// CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED works with CUDNN_RNN_PADDED_IO_ENABLED, so that it will auto fill 0 for the shorter sequences
CUDNN_RETURN_IF_ERROR(cudnnSetRNNPaddingMode(rnn_desc_, CUDNN_RNN_PADDED_IO_ENABLED));
CUDNN_RETURN_IF_ERROR(cudnnSetRNNPaddingMode(rnn_desc, CUDNN_RNN_PADDED_IO_ENABLED));
size_t workspace_bytes;
CUDNN_RETURN_IF_ERROR(cudnnGetRNNWorkspaceSize(CudnnHandle(), rnn_desc_, gsl::narrow_cast<int>(seq_length), x_desc.data(), &workspace_bytes));
CUDNN_RETURN_IF_ERROR(cudnnGetRNNWorkspaceSize(CudnnHandle(), rnn_desc, gsl::narrow_cast<int>(seq_length), x_desc.data(), &workspace_bytes));
auto workspace_cuda = GetScratchBuffer<void>(workspace_bytes);
if (CUDNN_RNN_RELU == rnn_mode_ || CUDNN_RNN_TANH == rnn_mode_ || nullptr == sequence_lens_data) {
CUDNN_RETURN_IF_ERROR(cudnnRNNForwardInference(CudnnHandle(),
rnn_desc_,
rnn_desc,
gsl::narrow_cast<int>(seq_length),
x_desc.data(),
x_data_input,
@ -273,7 +258,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
y_desc.Set(CudnnTensor::GetDataType<CudaT>(), seq_length, batch_size, hidden_size_ * num_directions_, sequence_lens_data);
CUDNN_RETURN_IF_ERROR(cudnnRNNForwardInferenceEx(CudnnHandle(),
rnn_desc_,
rnn_desc,
x_desc,
x_data_input,
hx_desc,

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

@ -21,26 +21,29 @@ enum RNN_Input_Index {
initial_c = 6
};
// Onnx RNN/GRU/LSTM only support 1 layer
const int RNN_NUM_LAYERS = 1;
class CudnnRNN {
public:
CudnnRNN() : rnn_desc_(nullptr) {
CudnnRNN() : cudnn_rnn_desc_(nullptr) {
}
~CudnnRNN() {
if (rnn_desc_ != nullptr) {
cudnnDestroyRNNDescriptor(rnn_desc_);
rnn_desc_ = nullptr;
if (cudnn_rnn_desc_ != nullptr) {
cudnnDestroyRNNDescriptor(cudnn_rnn_desc_);
cudnn_rnn_desc_ = nullptr;
}
}
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) {
if (!rnn_desc_)
CUDNN_RETURN_IF_ERROR(cudnnCreateRNNDescriptor(&rnn_desc_));
if (!cudnn_rnn_desc_)
CUDNN_RETURN_IF_ERROR(cudnnCreateRNNDescriptor(&cudnn_rnn_desc_));
CUDNN_RETURN_IF_ERROR(cudnnSetRNNDescriptor(cudnnHandle,
rnn_desc_,
cudnn_rnn_desc_,
gsl::narrow_cast<int>(hidden_size),
num_layers,
cudnn_dropout_desc,
@ -54,11 +57,11 @@ class CudnnRNN {
}
operator cudnnRNNDescriptor_t() const {
return rnn_desc_;
return cudnn_rnn_desc_;
}
private:
cudnnRNNDescriptor_t rnn_desc_;
cudnnRNNDescriptor_t cudnn_rnn_desc_;
};
template <typename T>
@ -68,23 +71,40 @@ class CudnnRnnBase : public CudaKernel {
public:
CudnnRnnBase(const OpKernelInfo& info) : CudaKernel{info} {
reverse_ = false;
ORT_ENFORCE(info.GetAttr("direction", &direction_).IsOK());
num_directions_ = direction_ == "bidirectional" ? 2 : 1;
ORT_ENFORCE(allowed_directions.find(direction_) != allowed_directions.end());
std::string direction = "forward";
direction = info.GetAttrOrDefault<std::string>("direction", "forward");
cudnn_direction_mode_ = CUDNN_UNIDIRECTIONAL;
if (direction == "bidirectional") {
cudnn_direction_mode_ = CUDNN_BIDIRECTIONAL;
} else if (direction == "forward") {
cudnn_direction_mode_ = CUDNN_UNIDIRECTIONAL;
} else if (direction == "reverse") {
cudnn_direction_mode_ = CUDNN_UNIDIRECTIONAL;
// need to reverse data
reverse_ = true;
}
num_directions_ = cudnn_direction_mode_ == CUDNN_BIDIRECTIONAL ? 2 : 1;
ORT_ENFORCE(allowed_directions.find(direction) != allowed_directions.end());
ORT_ENFORCE(info.GetAttr("hidden_size", &hidden_size_).IsOK() && hidden_size_ > 0);
rnn_mode_ = CUDNN_LSTM;
num_layers_ = 1;
weight_cached_ = false;
w_data_cache_ = nullptr;
size_t state_size;
cudnn_dropout_desc_.CreateDescriptorIfNeeded();
cudnn_dropout_desc_.GetCudnnDropoutStatesSize(CudnnHandle(), state_size);
state_buffer_ = GetScratchBuffer<void>(state_size);
cudnn_dropout_desc_.Set(CudnnHandle(), state_buffer_.get(), state_size);
}
Status SetCudnnRnnDesc();
Status CacheCudnnRnnWeights(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* ctx) const override;
void SetRNNMode(cudnnRNNMode_t rnn_mode) { rnn_mode_ = rnn_mode; }
private:
Status SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
const cudnnRNNDescriptor_t rnn_desc,
@ -97,7 +117,8 @@ class CudnnRnnBase : public CudaKernel {
Status ReorganizeWeights(const Tensor* W, const Tensor* R, const Tensor* B,
IAllocatorUniquePtr<void>& target_w_data,
CudnnFilterDescriptor& target_w_desc) const;
CudnnFilterDescriptor& target_w_desc,
CudnnRNN& rnn_desc) const;
void SetWeightBias(const cudnnHandle_t handle,
const cudnnRNNDescriptor_t rnn_desc,
@ -112,26 +133,25 @@ class CudnnRnnBase : public CudaKernel {
bool is_matrix) const;
protected:
int64_t num_directions_;
// required
int64_t hidden_size_;
cudnnRNNMode_t rnn_mode_;
// W_lin_layer_id_ & R_lin_layer_id_ are set in Constructor
std::vector<int> W_lin_layer_id_;
std::vector<int> R_lin_layer_id_;
CudnnRNN rnn_desc_;
bool reverse_;
int num_layers_;
private:
// optional
std::string direction_;
cudnnDirectionMode_t cudnn_direction_mode_;
bool reverse_;
int64_t num_directions_;
// hidden_size_ from attribute
int64_t hidden_size_;
cudnnRNNMode_t rnn_mode_;
// w_desc_cache_ & w_data_cache_ are changed in Constructor if we can get the weights as constant input
CudnnFilterDescriptor w_desc_cache_;
CudnnDropout cudnn_dropout_desc_;
CudnnFilterDescriptor filter_desc_;
IAllocatorUniquePtr<void> w_data_cache_;
bool weight_cached_;
// cudnn_dropout_desc_ is a cache, never to be changed
IAllocatorUniquePtr<void> state_buffer_;
size_t state_size_;
CudnnDropout cudnn_dropout_desc_;
enum Output_Index {
Y = 0,

3
onnxruntime/core/providers/cuda/rnn/gru.h
View file

@ -15,8 +15,7 @@ template <typename T>
class GRU final : public CudnnRnnBase<T> {
public:
GRU(const OpKernelInfo& info) : CudnnRnnBase<T>(info) {
CudnnRnnBase<T>::rnn_mode_ = CUDNN_GRU;
CudnnRnnBase<T>::SetCudnnRnnDesc();
CudnnRnnBase<T>::SetRNNMode(CUDNN_GRU);
// ONNX W layout is Wzrh, WBzrh, mapping to RNNLinLayerMatrixParams the linLayerID is 1, 0, 2
CudnnRnnBase<T>::W_lin_layer_id_.assign({1, 0, 2});

3
onnxruntime/core/providers/cuda/rnn/lstm.h
View file

@ -13,8 +13,7 @@ class LSTM final : public CudnnRnnBase<T> {
public:
LSTM(const OpKernelInfo& info) : CudnnRnnBase<T>(info) {
CudnnRnnBase<T>::rnn_mode_ = CUDNN_LSTM;
CudnnRnnBase<T>::SetCudnnRnnDesc();
CudnnRnnBase<T>::SetRNNMode(CUDNN_LSTM);
// ONNX W layout is W[iofc], WB[iofc], mapping to RNNLinLayerMatrixParams the linLayerID is 0, 3, 1, 2
CudnnRnnBase<T>::W_lin_layer_id_.assign({0, 3, 1, 2});

6
onnxruntime/core/providers/cuda/rnn/rnn.h
View file

@ -20,11 +20,9 @@ class RNN final : public CudnnRnnBase<T> {
std::vector<std::string> activations_;
ORT_ENFORCE(info.GetAttrs("activations", activations_).IsOK());
if (activations_[0] == "Relu")
CudnnRnnBase<T>::rnn_mode_ = CUDNN_RNN_RELU;
CudnnRnnBase<T>::SetRNNMode(CUDNN_RNN_RELU);
else if (activations_[0] == "Tanh")
CudnnRnnBase<T>::rnn_mode_ = CUDNN_RNN_TANH;
CudnnRnnBase<T>::SetCudnnRnnDesc();
CudnnRnnBase<T>::SetRNNMode(CUDNN_RNN_TANH);
// ONNX W mapping to RNNLinLayerMatrixParams the linLayerID is 0
CudnnRnnBase<T>::W_lin_layer_id_.assign({0});