44#include " core/providers/shared_library/provider_api.h"
55#include " cudnn_rnn_base.h"
66#include " rnn_impl.h"
7+ #include " core/common/safeint.h"
78
89namespace onnxruntime {
910namespace cuda {
@@ -16,7 +17,7 @@ Status CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
1617 const void * reorganized_w_data,
1718 const int lin_layer_id,
1819 const T* pos,
19- int & offset,
20+ size_t & offset,
2021 bool is_matrix,
2122 cudaStream_t cuda_stream) const {
2223 int numDims;
@@ -37,12 +38,12 @@ Status CudnnRnnBase<T>::SetWeightBias(const cudnnHandle_t handle,
3738 is_matrix ? tensor_desc_matrix : tensor_desc_bias, 3 , &dt, &numDims, matDims.data (), strideA.data ()));
3839
3940 mem_offset = is_matrix ? mem_offset_matrix : mem_offset_bias;
40- int count = matDims[0 ] * matDims[1 ] * matDims[2 ];
41+ size_t count = SafeInt< size_t >( matDims[0 ]) * matDims[1 ] * matDims[2 ];
4142
42- if (strideA[0 ] != count) {
43+ if (static_cast < size_t >( strideA[0 ]) != count) {
4344 return ORT_MAKE_STATUS (ONNXRUNTIME , StatusCode::INVALID_ARGUMENT , " Stride is not packed"
4445 }
45- CUDA_CALL_THROW (cudaMemcpyAsync (mem_offset, pos + offset, count * sizeof (T), cudaMemcpyDeviceToDevice, cuda_stream));
46+ CUDA_RETURN_IF_ERROR (cudaMemcpyAsync (mem_offset, pos + offset, count * sizeof (T), cudaMemcpyDeviceToDevice, cuda_stream));
4647
4748 offset += count;
4849
@@ -57,9 +58,9 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
5758 const T* R_data,
5859 const T* B_data,
5960 cudaStream_t cuda_stream) const {
60- int w_offset = 0 ;
61- int r_offset = 0 ;
62- int bias_offset = 0 ;
61+ size_t w_offset = 0 ;
62+ size_t r_offset = 0 ;
63+ size_t bias_offset = 0 ;
6364 for (int layer = 0 ; layer < RNN_NUM_LAYERS * num_directions_; ++layer) {
6465 for (size_t idx = 0 ; idx < W_lin_layer_id_.size (); ++idx) {
6566 ORT_RETURN_IF_ERROR (SetWeightBias (
@@ -92,6 +93,12 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
9293 void * alloc_stream, cudaStream_t cuda_stream,
9394 cudnnHandle_t cudnn_handle) const {
9495 typedef typename ToCudaType<T>::MappedType CudaT;
96+ ORT_RETURN_IF (W->Shape ().NumDimensions () != 3 ,
97+ " Weight W must be 3-D [num_directions, hidden_size, input_size], got rank "
98+ W->Shape ().NumDimensions ());
99+ ORT_RETURN_IF (R->Shape ().NumDimensions () != 3 ,
100+ " Recurrence R must be 3-D [num_directions, hidden_size, hidden_size], got rank "
101+ R->Shape ().NumDimensions ());
95102 int64_t input_size = W->Shape ()[2 ];
96103 // RNN W[num_directions_, hidden_size_, input_size]
97104 // RNN R[num_directions_, hidden_size_, hidden_size_]
@@ -103,12 +110,12 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
103110 // LSTM R[num_directions_, 4*hidden_size_, hidden_size_]
104111 // LSTM B[num_directions_, 8*hidden_size_]
105112 size_t number = W_lin_layer_id_.size ();
106- int64_t w_size = num_directions_ * ( number * hidden_size_ * (input_size + hidden_size_ + 2 ) );
113+ int64_t w_size = SafeInt< int64_t >( num_directions_) * number * hidden_size_ * (input_size + hidden_size_ + 2 );
107114 TensorShapeVector dims_w ({w_size, 1 , 1 });
108115 ORT_RETURN_IF_ERROR (target_w_desc.Set (dims_w, CudnnTensor::GetDataType<CudaT>()));
109116
110117 // Prepare the weight data
111- reorganized_w_data_size_in_bytes = w_size * sizeof (T);
118+ reorganized_w_data_size_in_bytes = SafeInt< size_t >( w_size) * sizeof (T);
112119 reorganized_w_data = GetScratchBuffer<void >(reorganized_w_data_size_in_bytes, alloc_stream);
113120
114121 // In many cases, this allocation is bigger than needed, leaving part of
@@ -142,6 +149,8 @@ Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
142149 bool has_bias = B != nullptr ;
143150
144151 if (get_W && get_R) {
152+ ORT_RETURN_IF (W->Shape ().NumDimensions () != 3 ,
153+ " Constant W must be 3-D, got rank " Shape ().NumDimensions ());
145154 CudnnRNN tmp_rnn_desc;
146155 auto proj_size = hidden_size_;
147156 ORT_RETURN_IF_ERROR (tmp_rnn_desc.Set (W->Shape ()[2 ], // input_size
@@ -163,7 +172,7 @@ Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
163172 w_data_cache_size_in_bytes_, w_data_cache_, w_desc_cache_,
164173 tmp_rnn_desc, nullptr , nullptr , DefaultCudnnHandle ()));
165174 }
166- cudaStreamSynchronize (nullptr );
175+ CUDA_RETURN_IF_ERROR ( cudaStreamSynchronize (nullptr ) );
167176
168177 weight_cached_ = true ;
169178 }
@@ -177,6 +186,9 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
177186 // inputs
178187 const Tensor* X = ctx->Input <Tensor>(RNN_Input_Index::X); // inputs. [seq_length, batch_size, input_size]
179188 ORT_ENFORCE (nullptr != X);
189+ ORT_RETURN_IF (X->Shape ().NumDimensions () != 3 ,
190+ " Input X must be 3-D [seq_length, batch_size, input_size], got rank "
191+ X->Shape ().NumDimensions ());
180192
181193 // optional inputs
182194 // [batch_size]
@@ -187,6 +199,10 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
187199 if (rnn_mode_ == CUDNN_LSTM ) {
188200 // initial cell. [num_directions_, batch_size, hidden_size_]
189201 initial_c = ctx->Input <Tensor>(RNN_Input_Index::initial_c);
202+ // cuDNN LSTM does not support peephole weights (ONNX input P at index 7)
203+ const Tensor* P = ctx->Input <Tensor>(7 );
204+ ORT_RETURN_IF (P != nullptr ,
205+ " CUDA LSTM does not support peephole weights (input P). Use CPU EP instead."
190206 }
191207
192208 size_t proj_size = hidden_size_;
@@ -197,7 +213,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
197213 // we thread a single input as sequence_lens of length 1, require to expand to [batch_size]?
198214 std::vector<int32_t > sequence_lengths_temp;
199215 if (!sequence_lens) {
200- sequence_lengths_temp.resize (batch_size, gsl::narrow_cast <int32_t >(seq_length));
216+ sequence_lengths_temp.resize (batch_size, gsl::narrow <int32_t >(seq_length));
201217 }
202218
203219 const int32_t * sequence_lens_data = (sequence_lens == nullptr )
@@ -214,10 +230,10 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
214230
215231 // 0-len sequences are not supported by cuDNN.
216232 // Replace them by sequences of len 1 and mask them out with SetZeroSequences
217- for (int i = 0 ; i < batch_size; ++i) {
233+ for (int64_t i = 0 ; i < batch_size; ++i) {
218234 if (0 == sequence_lens_data[i]) {
219235 seq_len_array[i] = 1 ;
220- zero_seq_index_cache[zero_seq_count] = i ;
236+ zero_seq_index_cache[zero_seq_count] = gsl::narrow< int32_t >(i) ;
221237 ++zero_seq_count;
222238 } else {
223239 seq_len_array[i] = sequence_lens_data[i];
@@ -257,15 +273,15 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
257273 const T* x_data = X->Data <T>();
258274 if (reverse_) {
259275 // reverse input data
260- x_reversed_data = GetScratchBuffer<T>(seq_length * batch_size * input_size, GetComputeStream (ctx));
276+ x_reversed_data = GetScratchBuffer<T>(SafeInt< int64_t >( seq_length) * batch_size * input_size, GetComputeStream (ctx));
261277 ReverseBySequence (Stream (ctx),
262- gsl::narrow_cast <int32_t >(seq_length),
278+ gsl::narrow <int32_t >(seq_length),
263279 sequence_lens_buffer.GpuPtr (),
264- gsl::narrow_cast <int32_t >(batch_size),
265- gsl::narrow_cast <int32_t >(input_size),
280+ gsl::narrow <int32_t >(batch_size),
281+ gsl::narrow <int32_t >(input_size),
266282 reinterpret_cast <const CudaT*>(x_data),
267283 reinterpret_cast <CudaT*>(x_reversed_data.get ()),
268- seq_length * batch_size * input_size);
284+ SafeInt< int64_t >( seq_length) * batch_size * input_size);
269285 }
270286
271287 const T* x_data_input = reverse_ ? x_reversed_data.get () : x_data;
@@ -274,7 +290,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
274290 const T* cx_data = (initial_c == nullptr ) ? nullptr : initial_c->Data <T>();
275291 T* y_h_data = (Y_h == nullptr ) ? nullptr : Y_h->MutableData <T>();
276292 T* y_c_data = (Y_c == nullptr ) ? nullptr : Y_c->MutableData <T>();
277- int64_t output_size = seq_length * num_directions_ * batch_size * hidden_size_;
293+ int64_t output_size = SafeInt< int64_t >( seq_length) * num_directions_ * batch_size * hidden_size_;
278294 T* y_data = nullptr ;
279295 IAllocatorUniquePtr<T> y_alloc_data;
280296 if (Y != nullptr ) {
@@ -357,7 +373,8 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
357373 // Mask on output for 0 sequence batches
358374 if (zero_seq_count > 0 ) {
359375 // Mask on output for 0 sequence batches
360- SetZeroSequences (zero_seq_count, zero_seq_index_cache, y_data, y_h_data, y_c_data, GetComputeStream (ctx), Stream (ctx));
376+ SetZeroSequences (gsl::span<const int32_t >(zero_seq_index_cache.data (), zero_seq_count),
377+ y_data, y_h_data, y_c_data, GetComputeStream (ctx), Stream (ctx));
361378 }
362379 return Status::OK ();
363380 }
@@ -369,26 +386,26 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
369386 if (reverse_) {
370387 // reverse output data
371388 ReverseBySequence (Stream (ctx),
372- gsl::narrow_cast <int32_t >(seq_length),
389+ gsl::narrow <int32_t >(seq_length),
373390 sequence_lens_buffer.GpuPtr (),
374- gsl::narrow_cast <int32_t >(batch_size),
375- gsl::narrow_cast <int32_t >(hidden_size_),
391+ gsl::narrow <int32_t >(batch_size),
392+ gsl::narrow <int32_t >(hidden_size_),
376393 reinterpret_cast <CudaT*>(y_data),
377394 reinterpret_cast <CudaT*>(y_reorganized_data.get ()),
378395 output_size);
379396 } else {
380397 ReorderBidirectionalDataInSequence (Stream (ctx),
381- gsl::narrow_cast <int32_t >(seq_length),
382- gsl::narrow_cast <int32_t >(batch_size),
383- gsl::narrow_cast <int32_t >(hidden_size_),
398+ gsl::narrow <int32_t >(seq_length),
399+ gsl::narrow <int32_t >(batch_size),
400+ gsl::narrow <int32_t >(hidden_size_),
384401 reinterpret_cast <CudaT*>(y_data),
385402 reinterpret_cast <CudaT*>(y_reorganized_data.get ()),
386403 output_size);
387404 }
388405
389406 if (Y != nullptr ) {
390407 // User specified this optional output, so need to copy the reversed data to original place
391- CUDA_RETURN_IF_ERROR (cudaMemcpyAsync (y_data, y_reorganized_data.get (), output_size * sizeof (T),
408+ CUDA_RETURN_IF_ERROR (cudaMemcpyAsync (y_data, y_reorganized_data.get (), SafeInt< size_t >( output_size) * sizeof (T),
392409 cudaMemcpyDeviceToDevice, Stream (ctx)));
393410 } else {
394411 y_data = y_reorganized_data.get ();
@@ -397,26 +414,27 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
397414
398415 // Mask on output for 0 sequence batches
399416 if (zero_seq_count > 0 ) {
400- SetZeroSequences (zero_seq_count, zero_seq_index_cache, y_data, y_h_data, y_c_data, GetComputeStream (ctx), Stream (ctx));
417+ SetZeroSequences (gsl::span<const int32_t >(zero_seq_index_cache.data (), zero_seq_count),
418+ y_data, y_h_data, y_c_data, GetComputeStream (ctx), Stream (ctx));
401419 }
402420
403421 return Status::OK ();
404422}
405423
406424template <typename T>
407- void CudnnRnnBase<T>::SetZeroSequences(const int64_t zero_seq_index_cache_size,
408- const std::vector<int32_t > zero_seq_index_cache,
425+ void CudnnRnnBase<T>::SetZeroSequences(gsl::span<const int32_t > zero_seq_index_cache,
409426 T* y_data,
410427 T* y_h_data,
411428 T* y_c_data,
412429 void * alloc_stream, cudaStream_t cuda_stream) const {
413430 typedef typename ToCudaType<T>::MappedType CudaT;
431+ const int64_t zero_seq_index_cache_size = static_cast <int64_t >(zero_seq_index_cache.size ());
414432 CudaAsyncBuffer<int32_t > zero_seq_index_cache_async_buffer (this , zero_seq_index_cache_size);
415433 memcpy (zero_seq_index_cache_async_buffer.CpuPtr (), zero_seq_index_cache.data (),
416434 zero_seq_index_cache_size * sizeof (int32_t ));
417435 ORT_THROW_IF_ERROR (zero_seq_index_cache_async_buffer.CopyToGpu (alloc_stream));
418436 MaskZeroSequences (cuda_stream,
419- gsl::narrow_cast <int32_t >(hidden_size_),
437+ gsl::narrow <int32_t >(hidden_size_),
420438 reinterpret_cast <CudaT*>(y_data),
421439 reinterpret_cast <CudaT*>(y_h_data),
422440 reinterpret_cast <CudaT*>(y_c_data),
0 commit comments