diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index c0594c1269..933b9a2ede 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -196,17 +196,24 @@ Status ReduceKernel::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(Y->template MutableData()))); - } else { CUDNN_RETURN_IF_ERROR(cudnnReduceTensor( CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes, - &one, input_tensor, reinterpret_cast(X->template Data()), + &one, input_tensor, input_data, &zero, output_tensor, reinterpret_cast(Y->template MutableData()))); + } 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() != X->template Data()) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), 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(X->template Data()), + &zero, output_tensor, reinterpret_cast(Y->template MutableData()))); + } } - } 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(output_count); CUDNN_RETURN_IF_ERROR(cudnnReduceTensor( diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc index 624d45b7a2..5dca49d8ac 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc +++ b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc @@ -47,21 +47,20 @@ Status CudnnRnnBase::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::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())); return Status::OK(); @@ -137,10 +135,14 @@ Status CudnnRnnBase::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; } diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h index 2aaced147e..4d44c96880 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.h +++ b/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 w_data_cache_; bool weight_cached_; diff --git a/onnxruntime/core/providers/cuda/rnn/gru.h b/onnxruntime/core/providers/cuda/rnn/gru.h index d69c122ee4..43a0ba4ab5 100644 --- a/onnxruntime/core/providers/cuda/rnn/gru.h +++ b/onnxruntime/core/providers/cuda/rnn/gru.h @@ -18,9 +18,9 @@ class GRU final : public CudnnRnnBase { CudnnRnnBase::rnn_mode_ = CUDNN_GRU; CudnnRnnBase::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::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::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 diff --git a/onnxruntime/core/providers/cuda/rnn/rnn_impl.cu b/onnxruntime/core/providers/cuda/rnn/rnn_impl.cu index fb7b9f4753..ae210ae681 100644 --- a/onnxruntime/core/providers/cuda/rnn/rnn_impl.cu +++ b/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(N) / GridDim::maxThreadsPerBlock)); _RnnMaskKernel<<>>( - 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) \ diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 42e5928afc..58028bada1 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/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{2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 1}, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddOutput("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{0}); diff --git a/onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc b/onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc index 22852972bc..e4143d529e 100644 --- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc @@ -321,7 +321,8 @@ class DeepCpuGruOpTestContext { const std::vector& sequence_length, const std::vector* initial_h, const std::vector& expected_Y, - const std::vector& expected_Y_h); + const std::vector& expected_Y_h, + const bool linear_before_reset = false); private: const int input_size_; @@ -456,7 +457,8 @@ void DeepCpuGruOpTestContext::RunTest(const std::vector& X, const std::vector& sequence_lens, const std::vector* initial_h, const std::vector& expected_Y, - const std::vector& expected_Y_h) { + const std::vector& 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& 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 activations = {"sigmoid", "tanh"}; + + DeepCpuGruOpTestContext ctx(direction, activations); + + const int batch_size = 2; + const int seq_length = 2; + std::vector X = {-0.455351f, -0.276391f, + -0.455351f, -0.276391f, + + -0.185934f, -0.269585f, + -0.185934f, -0.269585f}; + std::vector sequence_length = {2, 2}; + std::vector initial_h = {0.5f, -0.5f, 0.0f, 0.0f}; + std::vector expected_Y = {0.253942400f, -0.174207777f, + -0.0325528607f, 0.0774837881f, + + 0.0874997079f, -0.0485242009f, + -0.0577347837f, 0.0796165839f}; + + std::vector 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 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 activations = {"sigmoid", "tanh"}; + + DeepCpuGruOpTestContext ctx(direction, activations); + + const int batch_size = 1; + const int seq_length = 2; + std::vector X = {-0.455351f, -0.276391f, + -0.185934f, -0.269585f}; + std::vector sequence_length = {2}; + std::vector initial_h = {0.0f, 0.0f}; + std::vector expected_Y = {-0.0325528607f, 0.0774837881f, + -0.0577347837f, 0.0796165839f}; + std::vector 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 X2 = {-0.455351f, -0.276391f, + -0.455351f, -0.276391f, + + -0.185934f, -0.269585f, + 0.0f, 0.0f}; + std::vector sequence_length2 = {2, 1}; + std::vector initial_h2 = {0.5f, -0.5f, + 0.0f, 0.0f}; + std::vector expected_Y2 = {0.253942400f, -0.174207777f, + -0.0325528607f, 0.0774837881f, + + 0.0874997079f, -0.0485242009f, + 0.0f, 0.0f}; + + std::vector 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 activations = {"sigmoid", "tanh", "sigmoid", "tanh"}; + + DeepCpuGruOpTestContext ctx(direction, activations); + + const int batch_size = 1; + const int seq_length = 2; + std::vector X = {-0.455351f, -0.276391f, + -0.185934f, -0.269585f}; + std::vector sequence_length = {2}; + std::vector initial_h = {0.0f, 0.0f, 0.0f, 0.0f}; + std::vector expected_Y = {-0.0325528607f, 0.0774837881f, -0.0559310019f, 0.101836264f, + -0.0577347837f, 0.0796165839f, -0.0456649922f, 0.0462125242f}; + + std::vector 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 activations = {"sigmoid", "tanh", "sigmoid", "tanh"}; + + DeepCpuGruOpTestContext ctx(direction, activations); + + const int batch_size = 1; + const int seq_length = 2; + std::vector X = {0.855351f, 0.676391f, + 0.585934f, 0.669585f}; + std::vector sequence_length = {2}; + std::vector initial_h = {0.0f, 0.0f, 0.0f, 0.0f}; + std::vector expected_Y = {-0.275918573f, -0.0022855850f, -0.385578573f, 0.0370728001f, + -0.382134795f, 0.0607641526f, -0.248751760f, 0.0347689129f}; + std::vector 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 activations = {"sigmoid", "tanh", "sigmoid", "tanh"}; +// +// DeepCpuGruOpTestContext ctx(direction, activations); +// +// const int batch_size = 2; +// const int seq_length = 2; +// std::vector X = {-0.455351f, -0.276391f, +// 0.855351f, 0.676391f, +// -0.185934f, -0.269585f, +// 0.585934f, 0.669585f}; +// std::vector sequence_length = {2, 1}; +// std::vector initial_h = {0.0f, 0.0f, 0.0f, 0.0f, +// 0.0f, 0.0f, 0.0f, 0.0f}; +// std::vector 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 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 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 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 sequence_length = {2, 0}; std::vector initial_h = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};