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

Cherry pick for 1.17.2 source code release (#19679)

Browse source 
### Description
<!-- Describe your changes. -->

As title.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Co-authored-by: Markus Tavenrath <mtavenrath@users.noreply.github.com>
This commit is contained in:
Rachel Guo 2024-02-28 10:20:01 -08:00 committed by GitHub
parent 8f5c79cb63
commit 3a4599792d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 239 additions and 301 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -24,12 +24,12 @@ class CudnnTensor final {
operator cudnnTensorDescriptor_t() const { return tensor_; }
Status CreateTensorIfNeeded();
template <typename T>
static cudnnDataType_t GetDataType();
private:
Status CreateTensorIfNeeded();
cudnnTensorDescriptor_t tensor_;
};

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

@ -9,40 +9,49 @@ namespace onnxruntime {
namespace cuda {
template <typename T>
void CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
const cudnnRNNDescriptor_t rnn_desc,
const int pseudo_layer,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
const cudnnFilterDescriptor_t filter_desc,
const void* reorganized_w_data,
const int lin_layer_id,
const T* pos,
int& offset,
bool is_matrix,
cudaStream_t cuda_stream) const {
Status CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
const cudnnRNNDescriptor_t rnn_desc,
const int pseudo_layer,
size_t reorganized_w_data_size,
const void* reorganized_w_data,
const int lin_layer_id,
const T* pos,
int& offset,
bool is_matrix,
cudaStream_t cuda_stream) const {
int numDims;
std::vector<int> matDims(3);
std::array<int, 3> matDims;
std::array<int, 3> strideA;
cudnnDataType_t dt;
cudnnTensorFormat_t tf;
T* mem_offset;
if (is_matrix) {
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, reorganized_w_data, lin_layer_id, filter_desc, (void**)&mem_offset);
}
CudnnTensor tensor_desc_matrix, tensor_desc_bias;
ORT_RETURN_IF_ERROR(tensor_desc_bias.CreateTensorIfNeeded());
ORT_RETURN_IF_ERROR(tensor_desc_matrix.CreateTensorIfNeeded());
cudnnGetFilterNdDescriptor(filter_desc, 3, &dt, &tf, &numDims, matDims.data());
T *mem_offset_matrix, *mem_offset_bias;
CUDNN_RETURN_IF_ERROR(cudnnGetRNNWeightParams(
handle, rnn_desc, pseudo_layer, reorganized_w_data_size, reorganized_w_data,
lin_layer_id, tensor_desc_matrix, (void**)&mem_offset_matrix, tensor_desc_bias, (void**)&mem_offset_bias));
CUDNN_RETURN_IF_ERROR(cudnnGetTensorNdDescriptor(
is_matrix ? tensor_desc_matrix : tensor_desc_bias, 3, &dt, &numDims, matDims.data(), strideA.data()));
mem_offset = is_matrix ? mem_offset_matrix : mem_offset_bias;
int count = matDims[0] * matDims[1] * matDims[2];
if (strideA[0] != count) {
return ORT_MAKE_STATUS(ONNXRUNTIME, StatusCode::INVALID_ARGUMENT, "Stride is not packed");
}
CUDA_CALL_THROW(cudaMemcpyAsync(mem_offset, pos + offset, count * sizeof(T), cudaMemcpyDeviceToDevice, cuda_stream));
offset += count;
return Status::OK();
}
template <typename T>
Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
const cudnnRNNDescriptor_t rnn_desc,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
size_t reorganized_w_data_size,
void* reorganized_w_data,
const T* W_data,
const T* R_data,
@ -51,18 +60,22 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
int w_offset = 0;
int r_offset = 0;
int bias_offset = 0;
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, reorganized_w_data, W_lin_layer_id_[idx], W_data, w_offset, true, cuda_stream);
ORT_RETURN_IF_ERROR(SetWeightBias(
cudnn_handle, rnn_desc, layer, reorganized_w_data_size, reorganized_w_data,
W_lin_layer_id_[idx], W_data, w_offset, true, cuda_stream));
if (B_data != nullptr) {
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, cuda_stream);
ORT_RETURN_IF_ERROR(SetWeightBias(cudnn_handle, rnn_desc, layer, reorganized_w_data_size, reorganized_w_data,
W_lin_layer_id_[idx], B_data, bias_offset, false, cuda_stream));
}
}
for (size_t idx = 0; idx < R_lin_layer_id_.size(); ++idx) {
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, cuda_stream);
ORT_RETURN_IF_ERROR(SetWeightBias(cudnn_handle, rnn_desc, layer, reorganized_w_data_size, reorganized_w_data,
R_lin_layer_id_[idx], R_data, r_offset, true, cuda_stream));
if (B_data != nullptr) {
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, cuda_stream);
ORT_RETURN_IF_ERROR(SetWeightBias(cudnn_handle, rnn_desc, layer, reorganized_w_data_size, reorganized_w_data,
R_lin_layer_id_[idx], B_data, bias_offset, false, cuda_stream));
}
}
}
@ -72,6 +85,7 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
template <typename T>
Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, const Tensor* B,
size_t& reorganized_w_data_size_in_bytes,
IAllocatorUniquePtr<void>& reorganized_w_data,
CudnnFilterDescriptor& target_w_desc,
CudnnRNN& rnn_desc, onnxruntime::Stream* ort_stream) const {
@ -91,19 +105,16 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
TensorShapeVector dims_w({w_size, 1, 1});
ORT_RETURN_IF_ERROR(target_w_desc.Set(dims_w, CudnnTensor::GetDataType<CudaT>()));
TensorShapeVector fake_dims_x({1, input_size, 1});
CudnnTensor fake_x_desc;
ORT_RETURN_IF_ERROR(fake_x_desc.Set(fake_dims_x, CudnnTensor::GetDataType<CudaT>()));
// Prepare the weight data
reorganized_w_data = GetScratchBuffer<void>(w_size * sizeof(T), ort_stream);
reorganized_w_data_size_in_bytes = w_size * sizeof(T);
reorganized_w_data = GetScratchBuffer<void>(reorganized_w_data_size_in_bytes, ort_stream);
// In many cases, this allocation is bigger than needed, leaving part of
// the buffer unintialized. non-zero garbage data leads to wrong result
// the buffer uninitialized. non-zero garbage data leads to wrong result
// in call to cudnnRNNForwardInference()
// TODO! refine allocation size for each case.
cudaStream_t cuda_stream = ort_stream ? static_cast<cudaStream_t>(ort_stream->GetHandle()) : nullptr;
cudaMemsetAsync(reorganized_w_data.get(), 0, w_size * sizeof(T), cuda_stream);
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(reorganized_w_data.get(), 0, reorganized_w_data_size_in_bytes, cuda_stream));
const T* W_data = W->Data<T>();
const T* R_data = R->Data<T>();
@ -111,8 +122,9 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
auto* ort_cuda_stream = dynamic_cast<CudaStream*>(ort_stream);
cudnnHandle_t cudnn_handle = ort_cuda_stream ? ort_cuda_stream->cudnn_handle_ : DefaultCudnnHandle();
ORT_RETURN_IF_ERROR(SetCudnnRnnWeightBias(cudnn_handle, rnn_desc, fake_x_desc, target_w_desc,
reorganized_w_data.get(), W_data, R_data, B_data, cuda_stream));
ORT_RETURN_IF_ERROR(SetCudnnRnnWeightBias(cudnn_handle, rnn_desc,
reorganized_w_data_size_in_bytes, reorganized_w_data.get(),
W_data, R_data, B_data, cuda_stream));
return Status::OK();
}
@ -128,22 +140,31 @@ Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
bool get_R = info.TryGetConstantInput(RNN_Input_Index::R, &R);
bool get_B = info.TryGetConstantInput(RNN_Input_Index::B, &B);
bool has_bias = B != nullptr;
if (get_W && get_R) {
CudnnRNN tmp_rnn_desc;
ORT_RETURN_IF_ERROR(tmp_rnn_desc.Set(DefaultCudnnHandle(),
auto proj_size = hidden_size_;
ORT_RETURN_IF_ERROR(tmp_rnn_desc.Set(W->Shape()[2], // input_size
hidden_size_,
proj_size,
RNN_NUM_LAYERS,
cudnn_dropout_desc_,
cudnn_direction_mode_,
rnn_mode_,
CudnnTensor::GetDataType<CudaT>(),
GetDeviceProp()));
has_bias,
CudnnTensor::GetDataType<CudaT>()));
if (get_B) {
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_, tmp_rnn_desc, nullptr));
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B,
w_data_cache_size_in_bytes_, w_data_cache_, w_desc_cache_,
tmp_rnn_desc, nullptr));
} else {
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, nullptr, w_data_cache_, w_desc_cache_, tmp_rnn_desc, nullptr));
ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, nullptr,
w_data_cache_size_in_bytes_, w_data_cache_, w_desc_cache_,
tmp_rnn_desc, nullptr));
}
cudaStreamSynchronize(nullptr);
weight_cached_ = true;
}
@ -158,17 +179,72 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
ORT_ENFORCE(nullptr != X);
// optional inputs
const Tensor* sequence_lens = ctx->Input<Tensor>(RNN_Input_Index::sequence_lens); // [batch_size]
const Tensor* initial_h = ctx->Input<Tensor>(RNN_Input_Index::initial_h); // initial hidden. [num_directions_, batch_size, hidden_size_]
// [batch_size]
const Tensor* sequence_lens = ctx->Input<Tensor>(RNN_Input_Index::sequence_lens);
// initial hidden. [num_directions_, batch_size, hidden_size_]
const Tensor* initial_h = ctx->Input<Tensor>(RNN_Input_Index::initial_h);
const Tensor* initial_c(nullptr);
if (rnn_mode_ == CUDNN_LSTM) {
initial_c = ctx->Input<Tensor>(RNN_Input_Index::initial_c); // initial cell. [num_directions_, batch_size, hidden_size_]
// initial cell. [num_directions_, batch_size, hidden_size_]
initial_c = ctx->Input<Tensor>(RNN_Input_Index::initial_c);
}
size_t proj_size = hidden_size_;
int64_t seq_length = X->Shape()[0];
int64_t batch_size = X->Shape()[1];
int64_t input_size = X->Shape()[2];
// we thread a single input as sequence_lens of length 1, require to expand to [batch_size]?
std::vector<int32_t> sequence_lengths_temp;
if (!sequence_lens) {
sequence_lengths_temp.resize(batch_size, gsl::narrow_cast<int32_t>(seq_length));
}
const int32_t* sequence_lens_data = (sequence_lens == nullptr)
? sequence_lengths_temp.data()
: sequence_lens->Data<int32_t>();
// cuDNN doesn't support 0 sequence inside the batch, find the 0 sequence and set it to 1
// there's a ZeroMask kernel to reset the result to 0 for the 0 sequence
int64_t zero_seq_count = 0;
std::vector<int32_t> zero_seq_index_cache(batch_size, 0);
CudaAsyncBuffer<int32_t> sequence_lens_buffer(this, batch_size);
int32_t* seq_len_array = sequence_lens_buffer.CpuPtr();
// 0-len sequences are not supported by cuDNN.
// Replace them by sequences of len 1 and mask them out with SetZeroSequences
for (int i = 0; i < batch_size; ++i) {
if (0 == sequence_lens_data[i]) {
seq_len_array[i] = 1;
zero_seq_index_cache[zero_seq_count] = i;
++zero_seq_count;
} else {
seq_len_array[i] = sequence_lens_data[i];
}
}
// Calculate the zero position cache for reverse direction if it's bidirectional
// The cache is for Y_h or Y_c, and the 1st sequence for Y, no need to do it for other sequence in Y since
// we hacked the 0 sequence to 1
if (zero_seq_count && num_directions_ > 1) {
zero_seq_index_cache.resize(zero_seq_count * num_directions_);
for (int64_t i = 0; i < zero_seq_count; ++i) {
zero_seq_index_cache[static_cast<size_t>(zero_seq_count) + i] =
static_cast<int32_t>(batch_size + zero_seq_index_cache[i]);
}
zero_seq_count *= num_directions_;
}
// Prior to cuDNN 8.9.1 the sequence lens buffer must be passed to cudnnRNNForward and thus is must
// be copied to the GPU always.
ORT_RETURN_IF_ERROR(sequence_lens_buffer.CopyToGpu(ctx->GetComputeStream()));
// Starting with cuDNN 8.9.1 the sequence lens buffer is ignored by cudnnRNNForward and thus it must
// be copied to the GPU only for the ReverseBySequence kernels.
// if (reverse_) {
// ORT_RETURN_IF_ERROR(sequence_lens_buffer.CopyToGpu(ctx->GetComputeStream()));
// }
// optional outputs
TensorShapeVector dims_Y({seq_length, num_directions_, batch_size, hidden_size_});
TensorShapeVector dims_hxy({RNN_NUM_LAYERS * num_directions_, batch_size, hidden_size_});
@ -177,25 +253,6 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
Tensor* Y_h = ctx->Output(Output_Index::Y_h, dims_hxy);
Tensor* Y_c = ctx->Output(Output_Index::Y_c, dims_yc);
std::vector<int64_t> dims_x({batch_size, input_size, 1});
std::vector<int64_t> dims_y({batch_size, hidden_size_ * num_directions_, 1});
CudnnTensor x_desc_temp;
ORT_RETURN_IF_ERROR(x_desc_temp.Set(dims_x, CudnnTensor::GetDataType<CudaT>()));
CudnnTensor y_desc_temp;
ORT_RETURN_IF_ERROR(y_desc_temp.Set(dims_y, CudnnTensor::GetDataType<CudaT>()));
std::vector<cudnnTensorDescriptor_t> x_desc(seq_length, x_desc_temp);
std::vector<cudnnTensorDescriptor_t> y_desc(seq_length, y_desc_temp);
CudnnTensor hx_desc;
CudnnTensor cx_desc;
CudnnTensor y_h_desc;
CudnnTensor y_c_desc;
ORT_RETURN_IF_ERROR(hx_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(cx_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
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>()));
IAllocatorUniquePtr<T> x_reversed_data;
const T* x_data = X->Data<T>();
if (reverse_) {
@ -203,6 +260,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
x_reversed_data = GetScratchBuffer<T>(seq_length * batch_size * input_size, ctx->GetComputeStream());
ReverseBySequence(Stream(ctx),
gsl::narrow_cast<int32_t>(seq_length),
sequence_lens_buffer.GpuPtr(),
gsl::narrow_cast<int32_t>(batch_size),
gsl::narrow_cast<int32_t>(input_size),
reinterpret_cast<const CudaT*>(x_data),
@ -226,115 +284,82 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
y_data = y_alloc_data.get();
}
const int32_t* sequence_lens_data = (sequence_lens == nullptr) ? nullptr : sequence_lens->Data<int32_t>();
const Tensor* B = ctx->Input<Tensor>(RNN_Input_Index::B);
bool has_bias = B != nullptr;
CudnnRNN rnn_desc;
ORT_RETURN_IF_ERROR(rnn_desc.Set(GetCudnnHandle(ctx),
ORT_RETURN_IF_ERROR(rnn_desc.Set(input_size,
hidden_size_,
proj_size,
RNN_NUM_LAYERS,
cudnn_dropout_desc_,
cudnn_direction_mode_,
rnn_mode_,
CudnnTensor::GetDataType<CudaT>(),
GetDeviceProp()));
has_bias,
CudnnTensor::GetDataType<CudaT>()));
// Prepare the weight data
size_t w_data_size_in_bytes = 0;
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);
ORT_RETURN_IF_ERROR(ReorganizeWeights(&W, &R, B, w_data, w_desc, rnn_desc, ctx->GetComputeStream()));
ORT_RETURN_IF_ERROR(ReorganizeWeights(&W, &R, B, w_data_size_in_bytes, w_data, w_desc,
rnn_desc, ctx->GetComputeStream()));
}
// 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));
CudnnDataTensor x_desc1;
ORT_RETURN_IF_ERROR(x_desc1.Set(CudnnTensor::GetDataType<CudaT>(), seq_length, batch_size,
input_size, seq_len_array));
CudnnDataTensor y_desc1;
ORT_RETURN_IF_ERROR(y_desc1.Set(CudnnTensor::GetDataType<CudaT>(), seq_length, batch_size,
((rnn_mode_ == CUDNN_LSTM) ? proj_size : hidden_size_) * num_directions_,
seq_len_array));
size_t workspace_bytes;
CUDNN_RETURN_IF_ERROR(cudnnGetRNNWorkspaceSize(GetCudnnHandle(ctx), rnn_desc, gsl::narrow_cast<int>(seq_length), x_desc.data(), &workspace_bytes));
CudnnTensor cx_desc;
ORT_RETURN_IF_ERROR(cx_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
CudnnTensor hx_desc;
ORT_RETURN_IF_ERROR(hx_desc.Set(dims_hxy, CudnnTensor::GetDataType<CudaT>()));
// reserveSpaceSize is not required cudnnRNNForward, but returned by cudnnGetRNNTempSpaceSizes
size_t workspace_bytes, reservespace_bytes;
CUDNN_RETURN_IF_ERROR(cudnnGetRNNTempSpaceSizes(GetCudnnHandle(ctx), rnn_desc, CUDNN_FWD_MODE_INFERENCE,
x_desc1, &workspace_bytes, &reservespace_bytes));
auto workspace_cuda = GetScratchBuffer<void>(workspace_bytes, ctx->GetComputeStream());
int64_t zero_seq_count = 0;
std::vector<int32_t> zero_seq_index_cache(batch_size, 0);
int64_t zero_seq_index_cache_size = 0;
auto reservespace_cuda = GetScratchBuffer<void>(reservespace_bytes, ctx->GetComputeStream());
if (CUDNN_RNN_RELU == rnn_mode_ || CUDNN_RNN_TANH == rnn_mode_ || nullptr == sequence_lens_data) {
CUDNN_RETURN_IF_ERROR(cudnnRNNForwardInference(GetCudnnHandle(ctx),
rnn_desc,
gsl::narrow_cast<int>(seq_length),
x_desc.data(),
x_data_input,
hx_desc,
hx_data,
cx_desc,
cx_data,
weight_cached_ ? w_desc_cache_ : w_desc,
weight_cached_ ? w_data_cache_.get() : w_data.get(),
y_desc.data(),
y_data,
y_h_desc,
y_h_data,
y_c_desc,
y_c_data,
workspace_cuda.get(),
workspace_bytes));
} else {
// cudnn doesn't support 0 sequence inside the batch, find the 0 sequence and set it to 1
// there's a ZeroMask kernel to reset the result to 0 for the 0 sequence
std::vector<int32_t> seq_len_array(sequence_lens_data, sequence_lens_data + batch_size);
for (int i = 0; i < batch_size; ++i) {
if (0 == seq_len_array[i]) {
seq_len_array[i] = 1;
zero_seq_index_cache[zero_seq_count] = i;
++zero_seq_count;
}
}
CUDNN_RETURN_IF_ERROR(cudnnRNNForward(GetCudnnHandle(ctx),
rnn_desc,
CUDNN_FWD_MODE_INFERENCE,
sequence_lens_buffer.GpuPtr(), // should be zero starting with cudnn 8.9.1
x_desc1,
x_data_input,
y_desc1,
y_data, // output
hx_desc,
hx_data, // input
y_h_data, // output
cx_desc, cx_data, y_c_data,
weight_cached_ ? w_data_cache_size_in_bytes_ : w_data_size_in_bytes,
weight_cached_ ? w_data_cache_.get() : w_data.get(),
workspace_bytes,
workspace_cuda.get(),
reservespace_bytes,
reservespace_cuda.get()));
// Calculate the zero position cache for reverse direction if it's bidirectional
// The cache is for Y_h or Y_c, and the 1st sequence for Y, no need to do it for other sequence in Y since
// we hacked the 0 sequence to 1
if (zero_seq_count && num_directions_ > 1) {
zero_seq_index_cache_size = zero_seq_count * num_directions_;
zero_seq_index_cache.resize(zero_seq_index_cache_size);
for (int64_t i = 0; i < zero_seq_count; ++i) {
zero_seq_index_cache[static_cast<size_t>(zero_seq_count) + i] = static_cast<int32_t>(batch_size + zero_seq_index_cache[i]);
}
}
CudnnDataTensor x_desc1;
ORT_RETURN_IF_ERROR(x_desc1.Set(CudnnTensor::GetDataType<CudaT>(), seq_length, batch_size, input_size, seq_len_array.data()));
CudnnDataTensor y_desc1;
ORT_RETURN_IF_ERROR(y_desc1.Set(CudnnTensor::GetDataType<CudaT>(), seq_length, batch_size, hidden_size_ * num_directions_, seq_len_array.data()));
CUDNN_RETURN_IF_ERROR(cudnnRNNForwardInferenceEx(GetCudnnHandle(ctx),
rnn_desc,
x_desc1,
x_data_input,
hx_desc,
hx_data,
cx_desc,
cx_data,
weight_cached_ ? w_desc_cache_ : w_desc,
weight_cached_ ? w_data_cache_.get() : w_data.get(),
y_desc1,
y_data,
y_h_desc,
y_h_data,
y_c_desc,
y_c_data,
nullptr, nullptr, nullptr, nullptr,
nullptr, nullptr, nullptr, nullptr,
workspace_cuda.get(),
workspace_bytes));
// Early terminate for this case since Y data is not required, and Y_h is obtained correctly, no need the following code to retrive Y_h from Y data.
if (nullptr == Y) {
// Early terminate for this case since Y data is not required, and Y_h is obtained correctly,
// no need the following code to retrieve Y_h from Y data.
if (nullptr == Y) {
// Mask on output for 0 sequence batches
if (zero_seq_count > 0) {
// Mask on output for 0 sequence batches
if (zero_seq_count > 0) {
SetZeroSequences(zero_seq_index_cache_size, zero_seq_index_cache, y_data, y_h_data, y_c_data, ctx->GetComputeStream());
}
return Status::OK();
SetZeroSequences(zero_seq_count, zero_seq_index_cache, y_data, y_h_data, y_c_data, ctx->GetComputeStream());
}
return Status::OK();
}
IAllocatorUniquePtr<T> y_reorganized_data;
@ -345,6 +370,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
// reverse output data
ReverseBySequence(Stream(ctx),
gsl::narrow_cast<int32_t>(seq_length),
sequence_lens_buffer.GpuPtr(),
gsl::narrow_cast<int32_t>(batch_size),
gsl::narrow_cast<int32_t>(hidden_size_),
reinterpret_cast<CudaT*>(y_data),
@ -361,8 +387,9 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
}
if (Y != nullptr) {
// User specified this optional output, so need to copy the reversed data to orignial place
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(y_data, y_reorganized_data.get(), output_size * sizeof(T), cudaMemcpyDeviceToDevice, Stream(ctx)));
// User specified this optional output, so need to copy the reversed data to original place
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(y_data, y_reorganized_data.get(), output_size * sizeof(T),
cudaMemcpyDeviceToDevice, Stream(ctx)));
} else {
y_data = y_reorganized_data.get();
}
@ -370,23 +397,9 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
// Mask on output for 0 sequence batches
if (zero_seq_count > 0) {
SetZeroSequences(zero_seq_index_cache_size, zero_seq_index_cache, y_data, y_h_data, y_c_data, ctx->GetComputeStream());
SetZeroSequences(zero_seq_count, zero_seq_index_cache, y_data, y_h_data, y_c_data, ctx->GetComputeStream());
}
if ((CUDNN_RNN_RELU == rnn_mode_ || CUDNN_RNN_TANH == rnn_mode_) && sequence_lens_data != nullptr && y_h_data != nullptr && y_data != nullptr) {
CudaAsyncBuffer<int32_t> sequence_lens_buffer(this, batch_size);
memcpy(sequence_lens_buffer.CpuPtr(), sequence_lens_data, batch_size * sizeof(int32_t));
ORT_RETURN_IF_ERROR(sequence_lens_buffer.CopyToGpu(ctx->GetComputeStream()));
RnnMaskImpl(Stream(ctx),
gsl::narrow_cast<int32_t>(num_directions_),
gsl::narrow_cast<int32_t>(seq_length),
gsl::narrow_cast<int32_t>(batch_size),
gsl::narrow_cast<int32_t>(hidden_size_),
sequence_lens_buffer.GpuPtr(),
reinterpret_cast<CudaT*>(y_data),
reinterpret_cast<CudaT*>(y_h_data),
output_size);
}
return Status::OK();
}
@ -399,7 +412,8 @@ void CudnnRnnBase<T>::SetZeroSequences(const int64_t zero_seq_index_cache_size,
onnxruntime::Stream* ort_stream) const {
typedef typename ToCudaType<T>::MappedType CudaT;
CudaAsyncBuffer<int32_t> zero_seq_index_cache_async_buffer(this, zero_seq_index_cache_size);
memcpy(zero_seq_index_cache_async_buffer.CpuPtr(), zero_seq_index_cache.data(), zero_seq_index_cache_size * sizeof(int32_t));
memcpy(zero_seq_index_cache_async_buffer.CpuPtr(), zero_seq_index_cache.data(),
zero_seq_index_cache_size * sizeof(int32_t));
ORT_THROW_IF_ERROR(zero_seq_index_cache_async_buffer.CopyToGpu(ort_stream));
cudaStream_t cuda_stream = ort_stream ? static_cast<cudaStream_t>(ort_stream->GetHandle()) : nullptr;
MaskZeroSequences(cuda_stream,

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

@ -38,26 +38,28 @@ class CudnnRNN {
}
}
Status Set(const cudnnHandle_t& cudnnHandle, int64_t hidden_size, int num_layers,
Status Set(int64_t input_size, int64_t hidden_size, int64_t proj_size, int num_layers,
cudnnDropoutDescriptor_t cudnn_dropout_desc, cudnnDirectionMode_t cudnn_direction_model,
cudnnRNNMode_t rnn_mode, cudnnDataType_t dataType, const cudaDeviceProp& prop) {
cudnnRNNMode_t rnn_mode, bool has_bias, cudnnDataType_t dataType) {
if (!cudnn_rnn_desc_)
CUDNN_RETURN_IF_ERROR(cudnnCreateRNNDescriptor(&cudnn_rnn_desc_));
CUDNN_RETURN_IF_ERROR(cudnnSetRNNDescriptor_v6(cudnnHandle,
cudnn_rnn_desc_,
CUDNN_RETURN_IF_ERROR(cudnnSetRNNDescriptor_v8(cudnn_rnn_desc_,
CUDNN_RNN_ALGO_STANDARD, // CUDNN_RNN_ALGO_PERSIST_STATIC, CUDNN_RNN_ALGO_PERSIST_DYNAMIC
rnn_mode,
has_bias ? CUDNN_RNN_DOUBLE_BIAS : CUDNN_RNN_NO_BIAS,
cudnn_direction_model,
CUDNN_LINEAR_INPUT,
dataType,
dataType,
dataType == CUDNN_DATA_HALF ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH,
gsl::narrow_cast<int>(input_size),
gsl::narrow_cast<int>(hidden_size),
gsl::narrow_cast<int>(proj_size), // projected size
num_layers,
cudnn_dropout_desc,
CUDNN_LINEAR_INPUT, // We can also skip the input matrix transformation
cudnn_direction_model,
rnn_mode,
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);
}
// 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_RNN_PADDED_IO_ENABLED));
return Status::OK();
}
@ -119,8 +121,7 @@ class CudnnRnnBase : public CudaKernel {
private:
Status SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
const cudnnRNNDescriptor_t rnn_desc,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
size_t w_data_size,
void* w_data,
const T* W_data,
const T* R_data,
@ -128,23 +129,22 @@ class CudnnRnnBase : public CudaKernel {
cudaStream_t cuda_stream) const;
Status ReorganizeWeights(const Tensor* W, const Tensor* R, const Tensor* B,
size_t& target_w_data_size_in_bytes,
IAllocatorUniquePtr<void>& target_w_data,
CudnnFilterDescriptor& target_w_desc,
CudnnRNN& rnn_desc,
onnxruntime::Stream* ort_stream) const;
void SetWeightBias(const cudnnHandle_t handle,
const cudnnRNNDescriptor_t rnn_desc,
const int pseudo_layer,
const cudnnTensorDescriptor_t x_desc,
const cudnnFilterDescriptor_t w_desc,
const cudnnFilterDescriptor_t filter_desc,
const void* w_data,
const int lin_layer_id,
const T* pos,
int& offset,
bool is_matrix,
cudaStream_t cuda_stream) const;
Status SetWeightBias(const cudnnHandle_t handle,
const cudnnRNNDescriptor_t rnn_desc,
const int pseudo_layer,
size_t w_data_size,
const void* w_data,
const int lin_layer_id,
const T* pos,
int& offset,
bool is_matrix,
cudaStream_t cuda_stream) const;
void SetZeroSequences(const int64_t zero_seq_index_cache_size,
const std::vector<int32_t> zero_seq_index_cache,
@ -167,6 +167,7 @@ class CudnnRnnBase : public CudaKernel {
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_;
size_t w_data_cache_size_in_bytes_;
IAllocatorUniquePtr<void> w_data_cache_;
bool weight_cached_;
int64_t layout_;

3
onnxruntime/core/providers/cuda/rnn/rnn.cc
View file

@ -1,8 +1,9 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/shared_library/provider_api.h"
#include "rnn.h"
#include "core/providers/shared_library/provider_api.h"
#include "rnn_impl.h"
#include "core/providers/cuda/cudnn_common.h"

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

@ -4,6 +4,7 @@
#pragma once
#include "cudnn_rnn_base.h"
#include "core/providers/cuda/cuda_common.h"
#include <cudnn.h>

91
onnxruntime/core/providers/cuda/rnn/rnn_impl.cu
View file

@ -8,22 +8,32 @@ namespace onnxruntime {
namespace cuda {
template <typename T>
__global__ void _ReverseBySequenceKernel(const int32_t seq_length,
__global__ void _ReverseBySequenceKernel(const int32_t max_seq_length,
const int32_t* seq_lengths,
const int32_t block_size,
const fast_divmod div_batch_block,
const fast_divmod div_input_or_hidden_size,
const T* data,
T* reversed_data,
const CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
int seq_id, offset;
div_batch_block.divmod(id, seq_id, offset);
int org_id = (seq_length - seq_id - 1) * block_size + offset;
reversed_data[id] = data[org_id];
int batch, batch_offset;
div_input_or_hidden_size.divmod(offset, batch, batch_offset);
int seq_id_org = seq_lengths[batch] - seq_id - 1;
if (seq_id_org >= 0) {
int org_id = seq_id_org * block_size + offset;
reversed_data[id] = data[org_id];
} else {
reversed_data[id] = T{};
}
}
template <typename T>
void ReverseBySequence(cudaStream_t stream,
const int32_t seq_length,
const int32_t max_seq_length,
const int32_t *seq_lengths,
const int32_t batch_size,
const int32_t input_or_hidden_size,
const T* data,
@ -32,9 +42,10 @@ void ReverseBySequence(cudaStream_t stream,
// kerneral
int32_t block_size = batch_size * input_or_hidden_size;
fast_divmod div_batch_block(block_size);
fast_divmod div_input_or_hidden_size(input_or_hidden_size);
int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
_ReverseBySequenceKernel<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
seq_length, block_size, div_batch_block, data, reversed_data, (CUDA_LONG)N);
max_seq_length, seq_lengths, block_size, div_batch_block, div_input_or_hidden_size, data, reversed_data, (CUDA_LONG)N);
}
template <typename T>
@ -82,60 +93,6 @@ void ReorderBidirectionalDataInSequence(cudaStream_t stream,
data, reordered_data, (CUDA_LONG)N);
}
template <typename T>
__global__ void _RnnMaskKernel(const int32_t seq_length,
const int32_t batch_size,
const int32_t hidden_size,
const int32_t* sequence_lens,
const fast_divmod div_seq_block,
const fast_divmod div_dir_block,
const fast_divmod div_batch_block,
T* y_output_data,
T* y_h_output_data,
const CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
int seq_id, direction_id, batch_id, offset;
div_seq_block.divmod(id, seq_id, offset);
div_dir_block.divmod(offset, direction_id, offset);
div_batch_block.divmod(offset, batch_id, offset);
int32_t batch_seq_length = sequence_lens[batch_id];
if (batch_id >= batch_size || batch_seq_length == seq_length) {
return;
}
if (seq_id >= batch_seq_length) {
y_output_data[id] = 0;
return;
}
if ((y_h_output_data != nullptr) &&
((direction_id == 0 && (seq_id + 1) == batch_seq_length) || (direction_id == 1 && seq_id == 0))) {
int hy_idx = direction_id * batch_size * hidden_size + batch_id * hidden_size + offset;
y_h_output_data[hy_idx] = y_output_data[id];
}
}
template <typename T>
void RnnMaskImpl(cudaStream_t stream,
const int32_t num_directions,
const int32_t seq_length,
const int32_t batch_size,
const int32_t hidden_size,
const int32_t* sequence_lens,
T* y_output_data,
T* y_h_output_data,
const size_t N) {
fast_divmod div_seq_block(batch_size * hidden_size * num_directions);
fast_divmod div_dir_block(batch_size * hidden_size);
fast_divmod div_batch_block(hidden_size);
int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
_RnnMaskKernel<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
seq_length, batch_size, hidden_size, sequence_lens, div_seq_block,
div_dir_block, div_batch_block, y_output_data, y_h_output_data, (CUDA_LONG)N);
}
template <typename T>
__global__ void _MaskZeroSequences(const int32_t hidden_size,
T* y_output_data,
@ -180,17 +137,9 @@ void MaskZeroSequences(cudaStream_t stream,
}
#define SPECIALIZED_RNN_IMPL(T) \
template void RnnMaskImpl<T>(cudaStream_t stream, \
const int32_t num_directions, \
const int32_t seq_length, \
const int32_t batch_size, \
const int32_t hidden_size, \
const int32_t* sequence_lens, \
T* y_output_data, \
T* y_h_output_data, \
const size_t N); \
template void ReverseBySequence<T>(cudaStream_t stream, \
const int32_t seq_length, \
template void ReverseBySequence<T>(cudaStream_t stream, \
const int32_t max_seq_length, \
const int32_t* seq_lengths, \
const int32_t batch_size, \
const int32_t hidden_size, \
const T* data, \
@ -203,7 +152,7 @@ void MaskZeroSequences(cudaStream_t stream,
const T* data, \
T* reordered_data, \
const size_t N); \
template void MaskZeroSequences<T>(cudaStream_t stream, \
template void MaskZeroSequences<T>(cudaStream_t stream, \
const int32_t hidden_size, \
T* y_output_data, \
T* y_h_output_data, \

14
onnxruntime/core/providers/cuda/rnn/rnn_impl.h
View file

@ -10,7 +10,8 @@ namespace cuda {
template <typename T>
void ReverseBySequence(cudaStream_t stream,
const int32_t seq_length,
const int32_t max_seq_length,
const int32_t* seq_lengths,
const int32_t batch_size,
const int32_t input_or_hidden_size,
const T* data,
@ -26,17 +27,6 @@ void ReorderBidirectionalDataInSequence(cudaStream_t stream,
T* reordered_data,
const size_t N);
template <typename T>
void RnnMaskImpl(cudaStream_t stream,
const int32_t num_directions,
const int32_t seq_length,
const int32_t batch_size,
const int32_t hidden_size,
const int32_t* sequence_lens,
T* y_output_data,
T* y_h_output_data,
const size_t N);
template <typename T>
void MaskZeroSequences(cudaStream_t stream,
const int32_t hidden_size,

24
onnxruntime/test/providers/cpu/rnn/rnn_op_test.cc
View file

@ -120,15 +120,11 @@ TEST(RNNTest, RNN_bidirectional_bias_initial_zigged_batch) {
test.AddOutput<float>("Y_h", Y_h_dims, Y_h_data);
// TensorRT failed on RNN tests
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kTensorrtExecutionProvider});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
// Doesn't work with CUDA 11.4 on Windows. Need investigation.
#if defined(USE_CUDA) && defined(_WIN32)
TEST(RNNTest, DISABLED_RNN_bidirectional_zigged_batch) {
#else
TEST(RNNTest, RNN_bidirectional_zigged_batch) {
#endif
OpTester test("RNN");
int64_t num_directions = 2, input_size = 2, hidden_size = 3, seq_length = 5;
@ -275,15 +271,11 @@ TEST(RNNTest, RNN_reverse_direction_zigged_batch) {
std::vector<float> Y_h_data({0.87014002F, 0.09402763F, -0.54269236F, 0.64809889F, -0.19472955F, -0.24271242F});
test.AddOutput<float>("Y_h", Y_h_dims, Y_h_data);
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kTensorrtExecutionProvider});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
// Doesn't work with CUDA 11.4 on Windows. Need investigation.
#if defined(USE_CUDA) && defined(_WIN32)
TEST(RNNTest, DISABLED_RNN_forward_direction_zigged_batch) {
#else
TEST(RNNTest, RNN_forward_direction_zigged_batch) {
#endif
OpTester test("RNN");
int64_t num_directions = 1, input_size = 2, hidden_size = 3, seq_length = 5;
@ -357,12 +349,7 @@ TEST(RNNTest, RNN_forward_direction_zigged_batch) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
// Doesn't work with CUDA 11.4 on Windows. Need investigation.
#if defined(USE_CUDA) && defined(_WIN32)
TEST(RNNTest, DISABLED_RNN_bidirectional_0) {
#else
TEST(RNNTest, RNN_bidirectional_0) {
#endif
OpTester test("RNN");
int64_t num_directions = 2, input_size = 2, hidden_size = 3, batch_size = 1, seq_length = 5;
@ -424,12 +411,7 @@ TEST(RNNTest, RNN_bidirectional_0) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
// Doesn't work with CUDA 11.4 on Windows. Need investigation.
#if defined(USE_CUDA) && defined(_WIN32)
TEST(RNNTest, DISABLED_RNN_bidirectional_1) {
#else
TEST(RNNTest, RNN_bidirectional_1) {
#endif
OpTester test("RNN");
int64_t num_directions = 2, input_size = 2, hidden_size = 2, batch_size = 1, seq_length = 1;
@ -597,7 +579,7 @@ TEST(RNNTest, DISABLED_RNN_default_attributes_and_forward_direction) {
}
}
TEST(RNNTest, DISABLED_RNN_reverse_direction) {
TEST(RNNTest, RNN_reverse_direction) {
int64_t num_directions = 1, input_size = 2, hidden_size = 3, batch_size = 1, seq_length = 5;
// In case of useDefault, attributes, inputs or outputs are not set.