From f07059ccc03bfc17510dbe2c5f7376afa3b46521 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Tue, 29 Sep 2020 13:33:38 -0700 Subject: [PATCH] Add weight prepacking to LSTM kernel (#5305) --- .../core/providers/cpu/rnn/deep_cpu_gru.cc | 2 +- .../core/providers/cpu/rnn/deep_cpu_lstm.cc | 131 +++++++++++++----- .../core/providers/cpu/rnn/deep_cpu_lstm.h | 12 +- onnxruntime/core/providers/cpu/rnn/rnn.cc | 2 +- .../core/providers/cpu/rnn/rnn_helpers.cc | 6 +- .../core/providers/cpu/rnn/rnn_helpers.h | 69 ++++++++- .../cpu/rnn/deep_cpu_lstm_op_test.cc | 82 ++++++----- 7 files changed, 213 insertions(+), 91 deletions(-) diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc index 2a5b3d9f9c..ca395808c3 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc @@ -272,7 +272,7 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const { int batch_size = gsl::narrow(X_shape[1]); int input_size = gsl::narrow(X_shape[2]); - auto status = ValidateCommonRnnInputs(X, W, R, B, 3, sequence_lens, initial_h, num_directions_, hidden_size_); + auto status = ValidateCommonRnnInputs(X, W.Shape(), R.Shape(), B, 3, sequence_lens, initial_h, num_directions_, hidden_size_); ORT_RETURN_IF_ERROR(status); // GRU outputs are optional but must be in the same order diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc index d7bb4df8cc..5df45aebb4 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc @@ -211,8 +211,8 @@ class UniDirectionalLstm { const ActivationFuncs::Entry& activation_func_h, float clip, concurrency::ThreadPool* thread_pool); void Compute(const gsl::span& inputs, const gsl::span& sequence_lengths, int num_directions, - const gsl::span& input_weights, const gsl::span& recurrent_weights, - gsl::span& outputs, gsl::span& final_hidden_state, gsl::span& final_cell_state); + const GemmWeights& input_weights, const GemmWeights& recurrent_weights, gsl::span& outputs, + gsl::span& final_hidden_state, gsl::span& final_cell_state); ~UniDirectionalLstm() = default; @@ -290,20 +290,74 @@ class UniDirectionalLstm { } // namespace detail +Status DeepCpuLstmOp::TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed) { + const auto& shape = weights.Shape(); + if (shape.NumDimensions() != 3) { + return Status::OK(); + } + + // weights: [num_directions, 4*hidden_size, input_size] + // recurrence weights: [num_directions, 4*hidden_size, hidden_size] + const size_t N = static_cast(shape[1]); + const size_t K = static_cast(shape[2]); + + if ((shape[0] != num_directions_) || (N != static_cast(hidden_size_ * 4))) { + return Status::OK(); + } + + const size_t packed_weights_size = MlasGemmPackBSize(N, K); + if (packed_weights_size == 0) { + return Status::OK(); + } + + auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); + auto* packed_weights_data = alloc->Alloc(SafeInt(packed_weights_size) * num_directions_); + packed_weights.buffer_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); + packed_weights.weights_size_ = packed_weights_size; + packed_weights.shape_ = shape; + + const auto* weights_data = weights.Data(); + for (int i = 0; i < num_directions_; i++) { + MlasGemmPackB(CblasTrans, N, K, weights_data, K, packed_weights_data); + packed_weights_data = static_cast(packed_weights_data) + packed_weights_size; + weights_data += N * K; + } + + is_packed = true; + return Status::OK(); +} + +#if !defined(USE_MKLML_FOR_BLAS) +Status DeepCpuLstmOp::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { + is_packed = false; + + if (tensor.IsDataType()) { + if (input_idx == 1) { + return TryPackWeights(tensor, packed_W_, is_packed); + } else if (input_idx == 2) { + return TryPackWeights(tensor, packed_R_, is_packed); + } + } + + return Status::OK(); +} +#endif + Status DeepCpuLstmOp::Compute(OpKernelContext* context) const { const Tensor& X = *context->Input(0); // inputs. [seq_length, batch_size, input_size] Status status; // auto& logger = context->Logger(); - if (X.IsDataType()) + if (X.IsDataType()) { status = ComputeImpl(*context); - else if (X.IsDataType()) { + } else if (X.IsDataType()) { /* Need to update all the helpers to support double... status = ComputeImpl(*context); */ ORT_NOT_IMPLEMENTED("LSTM operator does not support double yet"); - } else + } else { ORT_THROW("Invalid data type for LSTM operator of ", X.DataType()); + } return status; } @@ -322,8 +376,10 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { auto& logger = context.Logger(); const Tensor& X = *context.Input(0); // inputs. [seq_length, batch_size, input_size] - const Tensor& W = *context.Input(1); // weights. [num_directions, 4*hidden_size, input_size] - const Tensor& R = *context.Input(2); // recurrence weights. [num_directions, 4*hidden_size, hidden_size] + const Tensor* W = packed_W_.buffer_ ? nullptr : context.Input(1); + // weights. [num_directions, 4*hidden_size, input_size] + const Tensor* R = packed_R_.buffer_ ? nullptr : context.Input(2); + // recurrence weights. [num_directions, 4*hidden_size, hidden_size] // optional const Tensor* B = context.Input(3); // bias. [num_directions, 8*hidden_size] @@ -332,13 +388,16 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { const Tensor* initial_c = context.Input(6); // initial cell. [num_directions, batch_size, hidden_size] const Tensor* P = context.Input(7); // peephole weights. [num_directions, 3*hidden_size] - auto& X_shape = X.Shape(); + const auto& X_shape = X.Shape(); int seq_length = gsl::narrow(X_shape[0]); int batch_size = gsl::narrow(X_shape[1]); int input_size = gsl::narrow(X_shape[2]); - Status status = ValidateInputs(X, W, R, B, sequence_lens, initial_h, initial_c, P, batch_size); + const auto& W_shape = (W != nullptr) ? W->Shape() : packed_W_.shape_; + const auto& R_shape = (R != nullptr) ? R->Shape() : packed_R_.shape_; + + Status status = ValidateInputs(X, W_shape, R_shape, B, sequence_lens, initial_h, initial_c, P, batch_size); ORT_RETURN_IF_ERROR(status); // LSTM outputs are optional but must be in the same order @@ -370,8 +429,9 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { status = context.GetTempSpaceAllocator(&alloc); ORT_RETURN_IF_ERROR(status); - gsl::span input_weights = W.DataAsSpan(); - gsl::span recurrent_weights = R.DataAsSpan(); + const auto* input_weights = (W != nullptr) ? W->Data() : nullptr; + const auto* recurrent_weights = (R != nullptr) ? R->Data() : nullptr; + gsl::span bias = B != nullptr ? B->DataAsSpan() : gsl::span(); gsl::span peephole_weights = P != nullptr ? P->DataAsSpan() : gsl::span(); @@ -381,8 +441,9 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { const size_t bias_size_per_direction = 8 * hidden_size_; const size_t peephole_weights_size_per_direction = 3 * hidden_size_; - gsl::span input_weights_1 = input_weights.subspan(0, input_weights_size_per_direction); - gsl::span recurrent_weights_1 = recurrent_weights.subspan(0, hidden_weights_size_per_direction); + GemmWeights input_weights_1(0, input_weights, input_weights_size_per_direction, packed_W_); + GemmWeights recurrent_weights_1(0, recurrent_weights, hidden_weights_size_per_direction, packed_R_); + gsl::span bias_1 = bias.empty() ? bias : bias.subspan(0, bias_size_per_direction); gsl::span peephole_weights_1 = peephole_weights.empty() ? peephole_weights : peephole_weights.subspan(0, peephole_weights_size_per_direction); @@ -427,11 +488,10 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { gsl::span last_cell_1 = last_cell.subspan(0, last_cell_size_per_direction); if (direction_ == Direction::kBidirectional) { + GemmWeights input_weights_2(1, input_weights, input_weights_size_per_direction, packed_W_); + GemmWeights recurrent_weights_2(1, recurrent_weights, hidden_weights_size_per_direction, packed_R_); + // spans for second direction - gsl::span input_weights_2 = - input_weights.subspan(input_weights_size_per_direction, input_weights_size_per_direction); - gsl::span hidden_weights_2 = - recurrent_weights.subspan(hidden_weights_size_per_direction, hidden_weights_size_per_direction); gsl::span bias_2 = bias.empty() ? bias : bias.subspan(bias_size_per_direction, bias_size_per_direction); gsl::span peephole_weights_2 = peephole_weights.empty() ? peephole_weights : peephole_weights.subspan(peephole_weights_size_per_direction, peephole_weights_size_per_direction); @@ -459,8 +519,8 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { fw.Compute(input, sequence_lens_span, num_directions_, input_weights_1, recurrent_weights_1, output_1, hidden_output_1, last_cell_1); - bw.Compute(input, sequence_lens_span, num_directions_, input_weights_2, hidden_weights_2, output_2, hidden_output_2, - last_cell_2); + bw.Compute(input, sequence_lens_span, num_directions_, input_weights_2, recurrent_weights_2, output_2, + hidden_output_2, last_cell_2); } else { detail::UniDirectionalLstm fw(alloc, logger, seq_length, batch_size, input_size, hidden_size_, direction_, input_forget_, bias_1, peephole_weights_1, initial_hidden_1, initial_cell_1, @@ -481,11 +541,11 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const { return Status::OK(); } -Status DeepCpuLstmOp::ValidateInputs(const Tensor& X, const Tensor& W, const Tensor& R, const Tensor* B, - const Tensor* sequence_lens, const Tensor* initial_h, const Tensor* initial_c, - const Tensor* P, int batch_size) const { +Status DeepCpuLstmOp::ValidateInputs(const Tensor& X, const TensorShape& W_shape, const TensorShape& R_shape, + const Tensor* B, const Tensor* sequence_lens, const Tensor* initial_h, + const Tensor* initial_c, const Tensor* P, int batch_size) const { auto status = - rnn::detail::ValidateCommonRnnInputs(X, W, R, B, 4, sequence_lens, initial_h, num_directions_, hidden_size_); + rnn::detail::ValidateCommonRnnInputs(X, W_shape, R_shape, B, 4, sequence_lens, initial_h, num_directions_, hidden_size_); ORT_RETURN_IF_ERROR(status); if (initial_c != nullptr) { @@ -680,8 +740,8 @@ void UniDirectionalLstm::LoadBias(const gsl::span& WbRb_values) { template void UniDirectionalLstm::Compute(const gsl::span& inputs_arg, const gsl::span& sequence_lengths_arg, const int num_directions, - const gsl::span& input_weights, - const gsl::span& recurrent_weights, gsl::span& outputs, + const GemmWeights& input_weights, const GemmWeights& recurrent_weights, + gsl::span& outputs, gsl::span& final_hidden_state, gsl::span& final_cell_state) { // copy spans (just T* and size, not data in span) as we may change them gsl::span inputs = inputs_arg; @@ -736,9 +796,9 @@ void UniDirectionalLstm::Compute(const gsl::span& inputs_arg, const int total_rows = max_sequence_length * batch_size_; // apply the weights to all the inputs and save to output_IOFC - ComputeGemm(total_rows, hidden_size_x4, input_size_, alpha, inputs.cbegin(), inputs.cend(), input_size_, - input_weights.cbegin(), input_weights.cend(), // W[iofc] - input_size_, beta, output_iofc_.begin(), output_iofc_.end(), hidden_size_x4, thread_pool_); + ComputeGemm(total_rows, hidden_size_x4, input_size_, alpha, inputs.cbegin(), inputs.cend(), + input_weights, + beta, output_iofc_.begin(), output_iofc_.end(), hidden_size_x4, thread_pool_); DumpMatrix("Xt*(W[iofc]^T)", output_iofc_.data(), total_rows, hidden_size_x4); @@ -783,10 +843,10 @@ void UniDirectionalLstm::Compute(const gsl::span& inputs_arg, // calculate Xt*(W[iofc]^T) + Ht-t*R[iofc] // Do it sequentially to avoid nested parallelism - ComputeGemm(local_fused_hidden_rows, hidden_size_x4, hidden_size_, alpha, previous_state, - previous_state_end, // Ht-1 - hidden_size_, recurrent_weights.cbegin(), recurrent_weights.cend(), // R[iofc] - hidden_size_, beta, step_out_IOFC, output_iofc_.end(), // input contains Xt*(W[iofc]^T) + ComputeGemm(local_fused_hidden_rows, hidden_size_x4, hidden_size_, alpha, + previous_state, previous_state_end, // Ht-1 + recurrent_weights, // R[iofc] + beta, step_out_IOFC, output_iofc_.end(), // input contains Xt*(W[iofc]^T) hidden_size_x4, nullptr); DumpMatrix("Xt*(W[iofc]^T) + Ht-t*R[iofc]" + row_str, &*step_out_IOFC, local_fused_hidden_rows, hidden_size_x4); @@ -861,9 +921,10 @@ void UniDirectionalLstm::Compute(const gsl::span& inputs_arg, span_T_iter step_out_IOFC = output_iofc_.begin() + (step * batch_size_) * hidden_size_x4; // calculate Xt*(W[iofc]^T) + Ht-t*R[iofc] - ComputeGemm(batch_size_, hidden_size_x4, hidden_size_, alpha, previous_state, previous_state_end, // Ht-1 - hidden_size_, recurrent_weights.cbegin(), recurrent_weights.cend(), // R[iofc] - hidden_size_, beta, step_out_IOFC, output_iofc_.end(), // input contains Xt*(W[iofc]^T) + ComputeGemm(batch_size_, hidden_size_x4, hidden_size_, alpha, + previous_state, previous_state_end, // Ht-1 + recurrent_weights, // R[iofc] + beta, step_out_IOFC, output_iofc_.end(), // input contains Xt*(W[iofc]^T) hidden_size_x4, thread_pool_); span_T_iter batched_output; diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h index a3adeb7cd5..7c2f978b0c 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h @@ -50,17 +50,22 @@ class DeepCpuLstmOp final : public OpKernel { activation_func_betas); } +#if !defined(USE_MKLML_FOR_BLAS) + Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; +#endif Status Compute(OpKernelContext* context) const override; ~DeepCpuLstmOp() override = default; private: + Status TryPackWeights(const Tensor& weights, rnn::detail::PackedWeights& packed_weights, bool& is_packed); + template Status ComputeImpl(OpKernelContext& context) const; Status ValidateInputs(const Tensor& X, - const Tensor& W, - const Tensor& R, + const TensorShape& W, + const TensorShape& R, const Tensor* B, const Tensor* sequence_lens, const Tensor* initial_h, @@ -75,6 +80,9 @@ class DeepCpuLstmOp final : public OpKernel { float clip_; bool input_forget_ = false; + rnn::detail::PackedWeights packed_W_; + rnn::detail::PackedWeights packed_R_; + rnn::detail::ActivationFuncs activation_funcs_; }; diff --git a/onnxruntime/core/providers/cpu/rnn/rnn.cc b/onnxruntime/core/providers/cpu/rnn/rnn.cc index 631a6727f9..c6fa145090 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn.cc +++ b/onnxruntime/core/providers/cpu/rnn/rnn.cc @@ -119,7 +119,7 @@ Status RNN::Compute(OpKernelContext* ctx) const { int64_t batch_size = X.Shape()[1]; int64_t input_size = X.Shape()[2]; - auto status = rnn::detail::ValidateCommonRnnInputs(X, W, R, B, 1, sequence_lens, initial_h, + auto status = rnn::detail::ValidateCommonRnnInputs(X, W.Shape(), R.Shape(), B, 1, sequence_lens, initial_h, num_directions, hidden_size_); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc index 682c7e0b70..99578eb67a 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc +++ b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc @@ -24,8 +24,8 @@ namespace detail { using namespace ::onnxruntime::common; Status ValidateCommonRnnInputs(const Tensor& X, - const Tensor& W, - const Tensor& R, + const TensorShape& W_shape, + const TensorShape& R_shape, const Tensor* B, int WRB_dim_1_multipler, const Tensor* sequence_lens, @@ -33,8 +33,6 @@ Status ValidateCommonRnnInputs(const Tensor& X, int64_t num_directions, int64_t hidden_size) { auto& X_shape = X.Shape(); - auto& W_shape = W.Shape(); - auto& R_shape = R.Shape(); int64_t seq_length = X_shape[0]; int64_t batch_size = X_shape[1]; diff --git a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h index 5c221a6ccd..8b6f025628 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h +++ b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h @@ -12,7 +12,8 @@ #include "core/framework/allocator.h" #include "core/util/math.h" #include "core/util/math_cpuonly.h" - +#include "core/mlas/inc/mlas.h" +#include "core/common/safeint.h" #include "core/platform/threadpool.h" #include "gsl/gsl" @@ -70,8 +71,8 @@ gsl::span Allocate(std::shared_ptr allocator, // validate the common inputs to RNN, LSTM and GRU operators Status ValidateCommonRnnInputs(const Tensor& X, - const Tensor& W, - const Tensor& R, + const TensorShape& W_shape, + const TensorShape& R_shape, const Tensor* B, int WRB_dim_1_multipler, // multiplier used with hidden_size for W, R and B inputs const Tensor* sequence_lens, @@ -145,7 +146,8 @@ void ComputeGemm(const int M, const float beta, TSpanCIter C, TSpanCIter C_end, - const int ldc, concurrency::ThreadPool* tp) { + const int ldc, + concurrency::ThreadPool* thread_pool) { // validate all the inputs // need to use the lda/ldb/ldc strides which should be >= the columns for the span ORT_ENFORCE(lda >= K && ldb >= K && ldc >= N); @@ -158,7 +160,64 @@ void ComputeGemm(const int M, M, N, K, alpha, &*A, lda, &*B, ldb, beta, - &*C, ldc, tp); + &*C, ldc, thread_pool); +} + +struct PackedWeights { + BufferUniquePtr buffer_; + size_t weights_size_; + TensorShape shape_; +}; + +template +struct GemmWeights { + GemmWeights(int idx, const T* weights_data, size_t weights_size, const PackedWeights& packed_weights) { + if (packed_weights.buffer_) { + is_prepacked_ = true; + buffer_ = static_cast(packed_weights.buffer_.get()) + packed_weights.weights_size_ * idx; + } else { + is_prepacked_ = false; + buffer_ = weights_data + weights_size * idx; + } + } + + bool is_prepacked_; + const void* buffer_; +}; + +template +void ComputeGemm(const int M, + const int N, + const int K, + const float alpha, + TSpanAIter A, + TSpanAIter A_end, + const GemmWeights& weights, + const float beta, + TSpanCIter C, + TSpanCIter C_end, + const int ldc, + concurrency::ThreadPool* thread_pool) { + // validate all the inputs + // need to use the lda/ldb/ldc strides which should be >= the columns for the span + ORT_ENFORCE(A + (M * K) <= A_end); + ORT_ENFORCE(C + (M * ldc - (ldc - N)) <= C_end); + + if (weights.is_prepacked_) { + MlasGemm( + CblasNoTrans, + M, N, K, alpha, + &*A, K, + weights.buffer_, beta, + &*C, ldc, thread_pool); + } else { + ::onnxruntime::math::GemmEx( + CblasNoTrans, CblasTrans, + M, N, K, alpha, + &*A, K, + static_cast(weights.buffer_), K, beta, + &*C, ldc, thread_pool); + } } // helper to convert a span to a raw pointer diff --git a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc index 31306f3e03..b3ca514780 100644 --- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc @@ -25,7 +25,9 @@ T DuplicateContainer(const T& container) { static void RunLstmTest(const std::vector& X_data, const std::vector& W_data, + bool is_initializer_W, const std::vector& R_data, + bool is_initializer_R, const std::vector& Y_data, const std::vector& Y_h_data, const std::vector& Y_c_data, @@ -78,8 +80,8 @@ static void RunLstmTest(const std::vector& X_data, std::vector R_dims = {num_directions, 4 * hidden_size, hidden_size}; test.AddInput("X", X_dims, X_data); - test.AddInput("W", W_dims, W_data); - test.AddInput("R", R_dims, R_data); + test.AddInput("W", W_dims, W_data, is_initializer_W); + test.AddInput("R", R_dims, R_data, is_initializer_R); if (B_data) { std::vector B_dims = {num_directions, 8 * hidden_size}; @@ -169,14 +171,14 @@ void SimpleWeightsNoBiasTwoRows(std::string direction, W_data = DuplicateContainer(W_data); } - RunLstmTest(X_data, W_data, R_data, Y_data, Y_h_data, Y_c_data, + RunLstmTest(X_data, W_data, false, R_data, false, Y_data, Y_h_data, Y_c_data, input_size, batch_size, hidden_size, seq_length, nullptr, nullptr, nullptr, nullptr, seq_lengths, direction); // need at least one output, so we need Y_h or Y_c to be requested (non-empty output to compare against) in order // to test Y not being returned (output_sequence == false) if (!Y_h_data.empty() || !Y_c_data.empty()) - RunLstmTest(X_data, W_data, R_data, Y_data, Y_h_data, Y_c_data, + RunLstmTest(X_data, W_data, false, R_data, false, Y_data, Y_h_data, Y_c_data, input_size, batch_size, hidden_size, seq_length, nullptr, nullptr, nullptr, nullptr, seq_lengths, direction, 999.f, /* output_sequence*/ false); } @@ -367,8 +369,12 @@ TEST(LSTMTest, BatchParallelFalseSeqLengthGreaterThanOne) { std::vector Y_c_data{ 1.02721067f, 1.15254318f}; - RunLstmTest(X_data, W_data, R_data, Y_data, {}, Y_c_data, - input_size, batch_size, hidden_size, seq_length); + for (bool is_initializer_W : std::initializer_list{false, true}) { + for (bool is_initializer_R : std::initializer_list{false, true}) { + RunLstmTest(X_data, W_data, is_initializer_W, R_data, is_initializer_R, + Y_data, {}, Y_c_data, input_size, batch_size, hidden_size, seq_length); + } + } } // make sure GateComputations works correctly if batch_parallel_ is true due to large batch size @@ -393,9 +399,13 @@ static void LargeBatchWithClip(const std::vector& Y_h_data, float clip = std::vector R_data(num_directions * 4 * hidden_size * hidden_size, 0.1f); - RunLstmTest(X_data, W_data, R_data, {}, Y_h_data, {}, - input_size, batch_size, hidden_size, seq_length, - nullptr, nullptr, nullptr, nullptr, nullptr, direction, clip); + for (bool is_initializer_W : std::initializer_list{false, true}) { + for (bool is_initializer_R : std::initializer_list{false, true}) { + RunLstmTest(X_data, W_data, is_initializer_W, R_data, is_initializer_R, {}, + Y_h_data, {}, input_size, batch_size, hidden_size, seq_length, + nullptr, nullptr, nullptr, nullptr, nullptr, direction, clip); + } + } } TEST(LSTMTest, LargeBatchNoClipping) { @@ -613,37 +623,25 @@ class LstmOpContext2x1x2x2 { bool input_forget = false, bool hasClip = true) { // run with and without output_sequence to test UniDirectionalLstm handling when Y isn't returned - ::onnxruntime::test::RunLstmTest(X, input_weights_, recurrent_weights_, - expected_Y, expected_Y_h, expected_Y_c, - input_size_, batch_size, hidden_size_, seq_length, - use_bias ? &bias_ : nullptr, - use_peepholes ? &peephole_weights_ : nullptr, - initial_h, initial_c, - sequence_lens, - direction_, - clip, - /*output_sequence*/ true, - input_forget, - activation_func_names_, - activation_alphas_, - activation_betas_, - hasClip); - - ::onnxruntime::test::RunLstmTest(X, input_weights_, recurrent_weights_, - expected_Y, expected_Y_h, expected_Y_c, - input_size_, batch_size, hidden_size_, seq_length, - use_bias ? &bias_ : nullptr, - use_peepholes ? &peephole_weights_ : nullptr, - initial_h, initial_c, - sequence_lens, - direction_, - clip, - /*output_sequence*/ false, - input_forget, - activation_func_names_, - activation_alphas_, - activation_betas_, - hasClip); + for (bool output_sequence : std::initializer_list{false, true}) { + ::onnxruntime::test::RunLstmTest(X, + input_weights_, false, + recurrent_weights_, false, + expected_Y, expected_Y_h, expected_Y_c, + input_size_, batch_size, hidden_size_, seq_length, + use_bias ? &bias_ : nullptr, + use_peepholes ? &peephole_weights_ : nullptr, + initial_h, initial_c, + sequence_lens, + direction_, + clip, + output_sequence, + input_forget, + activation_func_names_, + activation_alphas_, + activation_betas_, + hasClip); + } } private: @@ -669,10 +667,8 @@ TEST(LSTMTest, ONNXRuntime_TestLSTMForwardPeepHole) { std::vector Y_h_data = {-0.03277518f, 0.05935364f}; std::vector Y_c_data = {-0.0780206f, 0.098829f}; - std::string direction = "forward"; - //Run Test - LstmOpContext2x1x2x2 context(direction); + LstmOpContext2x1x2x2 context("forward"); context.RunTest(input, batch_size, seq_len, nullptr, nullptr, Y_data, Y_h_data, Y_c_data); }