Fix some issue in CUDA GRU and ReduceSum (#845)

* Fix issues in GRU GPU implementation. The cudnnGetRNNWorkspaceSize could failed because some descriptor are defined as local variable and are destroyed.

* Fix the issue for ReduceSum.  cudnnReduceTensor for ReduceSum has issue if input and output has same size, we just need to copy the data for this case.

onnxruntime/core/providers/cuda/reduction/reduction_ops.cc

@@ -196,17 +196,24 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
       return Status::OK();
     }
     if (calculate_sqt_) {
-      CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
-            CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
-            &one, input_tensor, input_data,
-            &zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
-    } else {
       CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
           CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
-          &one, input_tensor, reinterpret_cast<const CudaT*>(X->template Data<T>()),
+          &one, input_tensor, input_data,
           &zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
+    } else {
+      // cudnnReduceTensor for ReduceSum has issue if input and output has same size, we just need to copy the data for this case
+      if (input_count == output_count) {
+        if (Y->template MutableData<T>() != X->template Data<T>()) {
+          CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<T>(), X->template Data<T>(), input_count * sizeof(T), cudaMemcpyDeviceToDevice));
+        }
+      } else {
+        CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
+            CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
+            &one, input_tensor, reinterpret_cast<const CudaT*>(X->template Data<T>()),
+            &zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
+      }
     }
-  } else { // For ArgMax & ArgMin ops, use the indicies as the output with int64 type
+  } else {  // For ArgMax & ArgMin ops, use the indicies as the output with int64 type
     if (temp_X) {
       auto temp_output = GetScratchBuffer<float>(output_count);
       CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(

onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc

@@ -47,21 +47,20 @@ Status CudnnRnnBase<T>::SetCudnnRnnWeightBias(const cudnnHandle_t cudnn_handle,
                                               const T* R_data,
                                               const T* B_data) const {
   //Onnx only support 1 layer
-  CudnnFilterDescriptor filter_desc;
   int w_offset = 0;
   int r_offset = 0;
   int bias_offset = 0;
   for (int layer = 0; layer < num_layers_ * num_directions_; ++layer) {
     for (int 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_, 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_, w_data, W_lin_layer_id_[idx], B_data, bias_offset, false);
       }
     }
     for (int 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_, 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_, w_data, R_lin_layer_id_[idx], B_data, bias_offset, false);
       }
     }
   }
@@ -84,9 +83,8 @@ Status CudnnRnnBase<T>::SetCudnnRnnDesc() {
     reverse_ = true;
   }
 
-  CudnnDropout cudnn_dropout_desc;
-  cudnn_dropout_desc.Set(CudnnHandle());
-  ORT_RETURN_IF_ERROR(rnn_desc_.Set(CudnnHandle(), hidden_size_, num_layers_, cudnn_dropout_desc,
+  cudnn_dropout_desc_.Set(CudnnHandle());
+  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();
@@ -137,10 +135,14 @@ Status CudnnRnnBase<T>::CacheCudnnRnnWeights(const OpKernelInfo& info) {
   const Tensor* B;
   bool get_W = info.TryGetConstantInput(Input_Index::W, &W);
   bool get_R = info.TryGetConstantInput(Input_Index::R, &R);
+  bool get_B = info.TryGetConstantInput(Input_Index::B, &B);
 
   if (get_W && get_R) {
-    info.TryGetConstantInput(Input_Index::B, &B);
-    ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_));
+    if (get_B) {
+      ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, B, w_data_cache_, w_desc_cache_));
+    } else {
+      ORT_RETURN_IF_ERROR(ReorganizeWeights(W, R, nullptr, w_data_cache_, w_desc_cache_));
+    }
     weight_cached_ = true;
   }
 

onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h

@@ -153,6 +153,8 @@ class CudnnRnnBase : public CudaKernel {
   // optional
   std::string direction_;
   CudnnFilterDescriptor w_desc_cache_;
+  CudnnDropout cudnn_dropout_desc_;
+  CudnnFilterDescriptor filter_desc_;
   IAllocatorUniquePtr<void> w_data_cache_;
   bool weight_cached_;
 

onnxruntime/core/providers/cuda/rnn/gru.h

@@ -18,9 +18,9 @@ class GRU final : public CudnnRnnBase<T> {
     CudnnRnnBase<T>::rnn_mode_ = CUDNN_GRU;
     CudnnRnnBase<T>::SetCudnnRnnDesc();
 
-    // ONNX W layout is Wzrh, WBzrh, mapping to RNNLinLayerMatrixParams the linLayerID is 0, 3, 1, 2
+    // 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});
-    // ONNX R layout is Rzrh, RBzrh, mapping to RNNLinLayerMatrixParams the linLayerID is 4, 7, 5, 6
+    // ONNX R layout is Rzrh, RBzrh, mapping to RNNLinLayerMatrixParams the linLayerID is 4, 3, 5
     CudnnRnnBase<T>::R_lin_layer_id_.assign({4, 3, 5});
     // ONNX B layout is Wbzrh, Rbzrh, mapping to RNNLinLayerMatrixParams
     // the linLayerID is 1, 0, 2, 4, 3, 5, we can reuse it from W_lin_layer_id & R_lin_layer_id

onnxruntime/core/providers/cuda/rnn/rnn_impl.cu

@@ -86,15 +86,17 @@ __global__ void _RnnMaskKernel(const int32_t seq_length,
                                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, offset, batch_id, batch_offset;
+  int seq_id, direction_id, batch_id, offset;
   div_seq_block.divmod(id, seq_id, offset);
-  div_batch_block.divmod(offset, batch_id, batch_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) {
@@ -106,8 +108,9 @@ __global__ void _RnnMaskKernel(const int32_t seq_length,
     return;
   }
 
-  if ((y_h_output_data != nullptr) && (batch_seq_length != seq_length) && ((seq_id + 1) == batch_seq_length)) {
-    int hy_idx = batch_id * hidden_size + batch_offset;
+  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];
   }
 }
@@ -122,11 +125,12 @@ void RnnMaskImpl(const int32_t num_directions,
                  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>>>(
-      seq_length, batch_size, hidden_size, sequence_lens,
-      div_seq_block, div_batch_block, y_output_data, y_h_output_data, (CUDA_LONG)N);
+      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);
 }
 
 #define SPECIALIZED_RNN_IMPL(T)                                                 \

onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc

@@ -259,6 +259,18 @@ TEST(ReductionOpTest, ReduceLogSum) {
   test.Run();
 }
 
+TEST(ReductionOpTest, ReduceLogSum_samesize) {
+  OpTester test("ReduceLogSum");
+  test.AddAttribute("axes", std::vector<int64_t>{2});
+  test.AddAttribute("keepdims", (int64_t)1);
+  test.AddInput<float>("data", {3, 2, 1}, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  test.AddOutput<float>("reduced", {3, 2, 1},
+                        {0.0f, 0.6931471f,
+                         1.09861230f, 1.38629436f,
+                         1.60943794f, 1.79175949f});
+  test.Run();
+}
+
 TEST(ReductionOpTest, ReduceLogSum_do_not_keepdims_2) {
   OpTester test("ReduceLogSum");
   test.AddAttribute("axes", std::vector<int64_t>{0});

onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc

@@ -321,7 +321,8 @@ class DeepCpuGruOpTestContext {
                const std::vector<int>& sequence_length,
                const std::vector<float>* initial_h,
                const std::vector<float>& expected_Y,
-               const std::vector<float>& expected_Y_h);
+               const std::vector<float>& expected_Y_h,
+               const bool linear_before_reset = false);
 
  private:
   const int input_size_;
@@ -456,7 +457,8 @@ void DeepCpuGruOpTestContext::RunTest(const std::vector<float>& X,
                                       const std::vector<int>& sequence_lens,
                                       const std::vector<float>* initial_h,
                                       const std::vector<float>& expected_Y,
-                                      const std::vector<float>& expected_Y_h) {
+                                      const std::vector<float>& expected_Y_h,
+                                      const bool linear_before_reset) {
    //run with and without output_sequence
   ::onnxruntime::test::RunGruTest(X, gru_input_weights_, gru_recurrent_weights_,
                                   expected_Y, expected_Y_h,
@@ -467,24 +469,24 @@ void DeepCpuGruOpTestContext::RunTest(const std::vector<float>& X,
                                   direction_,
                                   9999999999.f,
                                   /*output_sequence*/ true,
-                                  false,
+                                   linear_before_reset,
                                   activation_func_names_,
                                   alphas_,
                                   betas_);
 
-  ::onnxruntime::test::RunGruTest(X, gru_input_weights_, gru_recurrent_weights_,
-                                  expected_Y, expected_Y_h,
-                                  input_size_, batch_size, hidden_dim_, seq_length,
-                                  use_bias_ ? &gru_bias_ : nullptr,
-                                  initial_h,
-                                  &sequence_lens,
-                                  direction_,
-                                  9999999999.f,
-                                  /*output_sequence*/ false,
-                                  false,
-                                  activation_func_names_,
-                                  alphas_,
-                                  betas_);
+  //::onnxruntime::test::RunGruTest(X, gru_input_weights_, gru_recurrent_weights_,
+  //                                expected_Y, expected_Y_h,
+  //                                input_size_, batch_size, hidden_dim_, seq_length,
+  //                                use_bias_ ? &gru_bias_ : nullptr,
+  //                                initial_h,
+  //                                &sequence_lens,
+  //                                direction_,
+  //                                9999999999.f,
+  //                                /*output_sequence*/ false,
+  //                                linear_before_reset,
+  //                                activation_func_names_,
+  //                                alphas_,
+  //                                betas_);
 }
 
 TEST(GRUTest, ONNXRuntime_TestGRUOpForwardBasic) {
@@ -612,6 +614,33 @@ TEST(GRUTest, ONNXRuntime_TestGRUOpForwardBatch) {
   ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h);
 }
 
+TEST(GRUTest, ONNXRuntime_TestGRUOpForwardBatchLinearBeforeReset) {
+  const std::string direction = "forward";
+  const std::vector<std::string> activations = {"sigmoid", "tanh"};
+
+  DeepCpuGruOpTestContext ctx(direction, activations);
+
+  const int batch_size = 2;
+  const int seq_length = 2;
+  std::vector<float> X = {-0.455351f, -0.276391f,
+                          -0.455351f, -0.276391f,
+
+                          -0.185934f, -0.269585f,
+                          -0.185934f, -0.269585f};
+  std::vector<int> sequence_length = {2, 2};
+  std::vector<float> initial_h = {0.5f, -0.5f, 0.0f, 0.0f};
+  std::vector<float> expected_Y = {0.253942400f, -0.174207777f,
+                                   -0.0325528607f, 0.0774837881f,
+
+                                   0.0874997079f, -0.0485242009f,
+                                   -0.0577347837f, 0.0796165839f};
+
+  std::vector<float> expected_Y_h = {0.0874997079f, -0.0485242009f,
+                                     -0.0577347837f, 0.0796165839f};
+
+  ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h, true);
+}
+
 TEST(GRUTest, ONNXRuntime_TestGRUOpGrowBatchSequenceLength) {
   const std::string direction = "forward";
   const std::vector<std::string> activations = {"sigmoid", "tanh"};
@@ -652,6 +681,116 @@ TEST(GRUTest, ONNXRuntime_TestGRUOpGrowBatchSequenceLength) {
   ctx.RunTest(X2, batch2, seq_length2, sequence_length2, &initial_h2, expected_Y2, expected_Y_h2);
 }
 
+TEST(GRUTest, ONNXRuntime_TestGRUOpGrowBatchSequenceLengthLinearBeforeReset) {
+  const std::string direction = "forward";
+  const std::vector<std::string> activations = {"sigmoid", "tanh"};
+
+  DeepCpuGruOpTestContext ctx(direction, activations);
+
+  const int batch_size = 1;
+  const int seq_length = 2;
+  std::vector<float> X = {-0.455351f, -0.276391f,
+                          -0.185934f, -0.269585f};
+  std::vector<int> sequence_length = {2};
+  std::vector<float> initial_h = {0.0f, 0.0f};
+  std::vector<float> expected_Y = {-0.0325528607f, 0.0774837881f,
+                                   -0.0577347837f, 0.0796165839f};
+  std::vector<float> expected_Y_h = {-0.0577347837f, 0.0796165839f};
+
+  ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h, true);
+
+  const int batch2 = 2;
+  const int seq_length2 = 2;
+  std::vector<float> X2 = {-0.455351f, -0.276391f,
+                           -0.455351f, -0.276391f,
+
+                           -0.185934f, -0.269585f,
+                           0.0f, 0.0f};
+  std::vector<int> sequence_length2 = {2, 1};
+  std::vector<float> initial_h2 = {0.5f, -0.5f,
+                                   0.0f, 0.0f};
+  std::vector<float> expected_Y2 = {0.253942400f, -0.174207777f,
+                                    -0.0325528607f, 0.0774837881f,
+
+                                    0.0874997079f, -0.0485242009f,
+                                    0.0f, 0.0f};
+
+  std::vector<float> expected_Y_h2 = {0.0874997079f, -0.0485242009f,
+                                      -0.0325528607f, 0.0774837881f};
+
+  ctx.RunTest(X2, batch2, seq_length2, sequence_length2, &initial_h2, expected_Y2, expected_Y_h2, true);
+}
+
+TEST(GRUTest, ONNXRuntime_TestGRUOpSequenceLengthWithBidirectionalLinearBeforeResetB1) {
+  const std::string direction = "bidirectional";
+  const std::vector<std::string> activations = {"sigmoid", "tanh", "sigmoid", "tanh"};
+
+  DeepCpuGruOpTestContext ctx(direction, activations);
+
+  const int batch_size = 1;
+  const int seq_length = 2;
+  std::vector<float> X = {-0.455351f, -0.276391f,
+                          -0.185934f, -0.269585f};
+  std::vector<int> sequence_length = {2};
+  std::vector<float> initial_h = {0.0f, 0.0f, 0.0f, 0.0f};
+  std::vector<float> expected_Y = {-0.0325528607f, 0.0774837881f, -0.0559310019f, 0.101836264f,
+                                   -0.0577347837f, 0.0796165839f, -0.0456649922f, 0.0462125242f};
+
+  std::vector<float> expected_Y_h = {-0.0577347837f, 0.0796165839f,
+                                     -0.0559310019f, 0.101836264f};
+
+  ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h, true);
+}
+
+TEST(GRUTest, ONNXRuntime_TestGRUOpSequenceLengthWithBidirectionalLinearBeforeResetB2) {
+  const std::string direction = "bidirectional";
+  const std::vector<std::string> activations = {"sigmoid", "tanh", "sigmoid", "tanh"};
+
+  DeepCpuGruOpTestContext ctx(direction, activations);
+
+  const int batch_size = 1;
+  const int seq_length = 2;
+  std::vector<float> X = {0.855351f, 0.676391f,
+                          0.585934f, 0.669585f};
+  std::vector<int> sequence_length = {2};
+  std::vector<float> initial_h = {0.0f, 0.0f, 0.0f, 0.0f};
+  std::vector<float> expected_Y = {-0.275918573f, -0.0022855850f, -0.385578573f, 0.0370728001f,
+                                   -0.382134795f, 0.0607641526f, -0.248751760f, 0.0347689129f};
+  std::vector<float> expected_Y_h = {-0.382134795f, 0.0607641526f,
+                                     -0.385578573f, 0.0370728001f};
+
+  ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h, true);
+}
+
+// Need CPU fix
+//TEST(GRUTest, ONNXRuntime_TestGRUOpSequenceLengthWithBidirectionalLinearBeforeReset) {
+//  const std::string direction = "bidirectional";
+//  const std::vector<std::string> activations = {"sigmoid", "tanh", "sigmoid", "tanh"};
+//
+//  DeepCpuGruOpTestContext ctx(direction, activations);
+//
+//  const int batch_size = 2;
+//  const int seq_length = 2;
+//  std::vector<float> X = {-0.455351f, -0.276391f,
+//                          0.855351f, 0.676391f,
+//                          -0.185934f, -0.269585f,
+//                          0.585934f, 0.669585f};
+//  std::vector<int> sequence_length = {2, 1};
+//  std::vector<float> initial_h = {0.0f, 0.0f, 0.0f, 0.0f,
+//                                  0.0f, 0.0f, 0.0f, 0.0f};
+//  std::vector<float> expected_Y = {-0.0325528607f, 0.0774837881f, -0.275918573f, -0.00228558504f,
+//                                   -0.0559310019f, 0.101836264f, -0.385578573f, 0.0370728001f,
+//
+//                                   -0.0577347837f, 0.0796165839f, 0.0f, 0.0f,
+//                                   -0.0456649922f, 0.0462125242f, 0.0f, 0.0f};
+//  std::vector<float> expected_Y_h = {-0.0577347837f, 0.0796165839f,
+//                                     -0.275918573f, -0.00228558504f,
+//                                     -0.0559310019f, 0.101836264f,
+//                                     -0.385578573f, 0.0370728001f};
+//
+//  ctx.RunTest(X, batch_size, seq_length, sequence_length, &initial_h, expected_Y, expected_Y_h, true);
+//}
+
 TEST(GRUTest, ONNXRuntime_TestGRUOpSequenceLengthWithPartialZero) {
   const std::string direction = "bidirectional";
   const std::vector<std::string> activations = {"sigmoid", "tanh", "sigmoid", "tanh"};
@@ -661,9 +800,9 @@ TEST(GRUTest, ONNXRuntime_TestGRUOpSequenceLengthWithPartialZero) {
   const int batch_size = 2;
   const int seq_length = 2;
   std::vector<float> X = {-0.455351f, -0.276391f,
-                          -0.455351f, -0.276391f,
+                          0.455351f, 0.276391f,
                           -0.185934f, -0.269585f,
-                          -0.185934f, -0.269585f};
+                          0.185934f, 0.269585f};
   std::vector<int> sequence_length = {2, 0};
   std::vector<float> initial_h = {0.0f, 0.0f, 0.0f, 0.0f,
                                   0.0f, 0.0f, 0.0f, 0.0f};