Refactors LSTM step function for reuse. Bug: 113559542 Test: NeuralNetworksTest_static --gtest_filter=GeneratedTests.lstm* Change-Id: I427bc41bf81f4d0cd0021f333e79b5cab95d7105 Merged-In: I427bc41bf81f4d0cd0021f333e79b5cab95d7105 (cherry picked from commit 72fd4fe3fb705c9dbd543f1fe1c25d3d23545ad3)
diff --git a/common/operations/LSTM.cpp b/common/operations/LSTM.cpp index f1e01e2..5159952 100644 --- a/common/operations/LSTM.cpp +++ b/common/operations/LSTM.cpp
@@ -71,15 +71,15 @@ output_state_in_ = GetInput(operation, operands, kOutputStateInTensor); cell_state_in_ = GetInput(operation, operands, kCellStateInTensor); - params_.activation_ = static_cast<TfLiteFusedActivation>( + params_.activation = static_cast<TfLiteFusedActivation>( getScalarData<int32_t>(*GetInput(operation, operands, kActivationParam))); if (input_->type == OperandType::TENSOR_FLOAT32) { - params_.cell_clip_ = getScalarData<float>(*GetInput(operation, operands, kCellClipParam)); - params_.proj_clip_ = getScalarData<float>(*GetInput(operation, operands, kProjClipParam)); + params_.cell_clip = getScalarData<float>(*GetInput(operation, operands, kCellClipParam)); + params_.proj_clip = getScalarData<float>(*GetInput(operation, operands, kProjClipParam)); } else { - params_.cell_clip_ = static_cast<float>( + params_.cell_clip = static_cast<float>( getScalarData<_Float16>(*GetInput(operation, operands, kCellClipParam))); - params_.proj_clip_ = static_cast<float>( + params_.proj_clip = static_cast<float>( getScalarData<_Float16>(*GetInput(operation, operands, kProjClipParam))); } @@ -108,100 +108,126 @@ scratch_buffer_ = GetOutput(operation, operands, kScratchBufferTensor); } -bool LSTMCell::CheckInputTensorDimensions(const Operation& operation, - std::vector<RunTimeOperandInfo>& operands, - uint32_t n_input, uint32_t n_output, uint32_t n_cell) { +// static +bool LSTMCell::CheckInputTensorDimensions( + const RunTimeOperandInfo* input_, const RunTimeOperandInfo* input_to_input_weights, + const RunTimeOperandInfo* input_to_forget_weights, + const RunTimeOperandInfo* input_to_cell_weights, + const RunTimeOperandInfo* input_to_output_weights, + const RunTimeOperandInfo* recurrent_to_input_weights, + const RunTimeOperandInfo* recurrent_to_forget_weights, + const RunTimeOperandInfo* recurrent_to_cell_weights, + const RunTimeOperandInfo* recurrent_to_output_weights, + const RunTimeOperandInfo* cell_to_input_weights, + const RunTimeOperandInfo* cell_to_forget_weights, + const RunTimeOperandInfo* cell_to_output_weights, const RunTimeOperandInfo* input_gate_bias, + const RunTimeOperandInfo* forget_gate_bias, const RunTimeOperandInfo* cell_bias, + const RunTimeOperandInfo* output_gate_bias, const RunTimeOperandInfo* projection_weights, + const RunTimeOperandInfo* projection_bias, + const RunTimeOperandInfo* input_layer_norm_weights, + const RunTimeOperandInfo* forget_layer_norm_weights, + const RunTimeOperandInfo* cell_layer_norm_weights, + const RunTimeOperandInfo* output_layer_norm_weights, uint32_t n_input, uint32_t n_output, + uint32_t n_cell, LSTMParams* params) { // Making sure clipping parameters have valid values. // == 0 means no clipping // > 0 means clipping - NN_CHECK(params_.cell_clip_ >= 0); - NN_CHECK(params_.proj_clip_ >= 0); + NN_CHECK(params->cell_clip >= 0); + NN_CHECK(params->proj_clip >= 0); - if (!IsNullInput(input_to_input_weights_)) { - NN_CHECK_EQ(NumDimensions(input_to_input_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(input_to_input_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(input_to_input_weights_, 1), n_input); + if (!IsNullInput(input_to_input_weights)) { + NN_CHECK_EQ(NumDimensions(input_to_input_weights), 2); + NN_CHECK_EQ(SizeOfDimension(input_to_input_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(input_to_input_weights, 1), n_input); } - NN_CHECK_EQ(NumDimensions(input_to_forget_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(input_to_forget_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(input_to_forget_weights_, 1), n_input); + NN_CHECK_EQ(NumDimensions(input_to_forget_weights), 2); + NN_CHECK_EQ(SizeOfDimension(input_to_forget_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(input_to_forget_weights, 1), n_input); - NN_CHECK_EQ(NumDimensions(input_to_cell_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(input_to_cell_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(input_to_cell_weights_, 1), n_input); + NN_CHECK_EQ(NumDimensions(input_to_cell_weights), 2); + NN_CHECK_EQ(SizeOfDimension(input_to_cell_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(input_to_cell_weights, 1), n_input); - if (!IsNullInput(recurrent_to_input_weights_)) { - NN_CHECK_EQ(NumDimensions(recurrent_to_input_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_input_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_input_weights_, 1), n_output); + if (!IsNullInput(recurrent_to_input_weights)) { + NN_CHECK_EQ(NumDimensions(recurrent_to_input_weights), 2); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_input_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_input_weights, 1), n_output); } - NN_CHECK_EQ(NumDimensions(recurrent_to_forget_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_forget_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_forget_weights_, 1), n_output); + NN_CHECK_EQ(NumDimensions(recurrent_to_forget_weights), 2); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_forget_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_forget_weights, 1), n_output); - NN_CHECK_EQ(NumDimensions(recurrent_to_cell_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_cell_weights_, 0), n_cell); - NN_CHECK_EQ(SizeOfDimension(recurrent_to_cell_weights_, 1), n_output); + NN_CHECK_EQ(NumDimensions(recurrent_to_cell_weights), 2); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_cell_weights, 0), n_cell); + NN_CHECK_EQ(SizeOfDimension(recurrent_to_cell_weights, 1), n_output); // We make sure the input-gate's parameters are either both present (regular // LSTM) or not at all (CIFG-LSTM). const bool cifg_weights_all_or_none = - (!IsNullInput(input_to_input_weights_) && !IsNullInput(recurrent_to_input_weights_)) || - (IsNullInput(input_to_input_weights_) && IsNullInput(recurrent_to_input_weights_)); + (!IsNullInput(input_to_input_weights) && !IsNullInput(recurrent_to_input_weights)) || + (IsNullInput(input_to_input_weights) && IsNullInput(recurrent_to_input_weights)); NN_CHECK(cifg_weights_all_or_none); - if (!IsNullInput(cell_to_input_weights_)) { - NN_CHECK_EQ(NumDimensions(cell_to_input_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(cell_to_input_weights_, 0), n_cell); + if (!IsNullInput(cell_to_input_weights)) { + NN_CHECK_EQ(NumDimensions(cell_to_input_weights), 1); + NN_CHECK_EQ(SizeOfDimension(cell_to_input_weights, 0), n_cell); } - if (!IsNullInput(cell_to_forget_weights_)) { - NN_CHECK_EQ(NumDimensions(cell_to_forget_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(cell_to_forget_weights_, 0), n_cell); + if (!IsNullInput(cell_to_forget_weights)) { + NN_CHECK_EQ(NumDimensions(cell_to_forget_weights), 1); + NN_CHECK_EQ(SizeOfDimension(cell_to_forget_weights, 0), n_cell); } - if (!IsNullInput(cell_to_output_weights_)) { - NN_CHECK_EQ(NumDimensions(cell_to_output_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(cell_to_output_weights_, 0), n_cell); + if (!IsNullInput(cell_to_output_weights)) { + NN_CHECK_EQ(NumDimensions(cell_to_output_weights), 1); + NN_CHECK_EQ(SizeOfDimension(cell_to_output_weights, 0), n_cell); } // Making sure the peephole weights are there all or none. - const bool use_cifg = IsNullInput(input_to_input_weights_); + params->use_cifg = IsNullInput(input_to_input_weights); const bool peephole_weights_all_or_none = - ((!IsNullInput(cell_to_input_weights_) || use_cifg) && - !IsNullInput(cell_to_forget_weights_) && !IsNullInput(cell_to_output_weights_)) || - (IsNullInput(cell_to_input_weights_) && IsNullInput(cell_to_forget_weights_) && - IsNullInput(cell_to_output_weights_)); + ((!IsNullInput(cell_to_input_weights) || params->use_cifg) && + !IsNullInput(cell_to_forget_weights) && !IsNullInput(cell_to_output_weights)) || + (IsNullInput(cell_to_input_weights) && IsNullInput(cell_to_forget_weights) && + IsNullInput(cell_to_output_weights)); NN_CHECK(peephole_weights_all_or_none); + // Since we have already checked that weights are all there or none, we can + // check the existence of only one to the get the condition. + params->use_peephole = !IsNullInput(cell_to_output_weights); + params->use_layer_norm = !IsNullInput(input_layer_norm_weights); + + params->use_projection_weight = (projection_weights->lifetime != OperandLifeTime::NO_VALUE); + params->use_projection_bias = (projection_bias->lifetime != OperandLifeTime::NO_VALUE); + // Make sure the input gate bias is present only when not a CIFG-LSTM. - if (use_cifg) { - NN_CHECK(IsNullInput(input_gate_bias_)); + if (params->use_cifg) { + NN_CHECK(IsNullInput(input_gate_bias)); } else { - NN_CHECK_EQ(NumDimensions(input_gate_bias_), 1); - NN_CHECK_EQ(SizeOfDimension(input_gate_bias_, 0), n_cell); + NN_CHECK_EQ(NumDimensions(input_gate_bias), 1); + NN_CHECK_EQ(SizeOfDimension(input_gate_bias, 0), n_cell); } - NN_CHECK_EQ(NumDimensions(forget_gate_bias_), 1); - NN_CHECK_EQ(SizeOfDimension(forget_gate_bias_, 0), n_cell); + NN_CHECK_EQ(NumDimensions(forget_gate_bias), 1); + NN_CHECK_EQ(SizeOfDimension(forget_gate_bias, 0), n_cell); - NN_CHECK_EQ(NumDimensions(cell_bias_), 1); - NN_CHECK_EQ(SizeOfDimension(cell_bias_, 0), n_cell); + NN_CHECK_EQ(NumDimensions(cell_bias), 1); + NN_CHECK_EQ(SizeOfDimension(cell_bias, 0), n_cell); - NN_CHECK_EQ(NumDimensions(output_gate_bias_), 1); - NN_CHECK_EQ(SizeOfDimension(output_gate_bias_, 0), n_cell); + NN_CHECK_EQ(NumDimensions(output_gate_bias), 1); + NN_CHECK_EQ(SizeOfDimension(output_gate_bias, 0), n_cell); - if (!IsNullInput(projection_weights_)) { - NN_CHECK_EQ(NumDimensions(projection_weights_), 2); - NN_CHECK_EQ(SizeOfDimension(projection_weights_, 0), n_output); - NN_CHECK_EQ(SizeOfDimension(projection_weights_, 1), n_cell); + if (!IsNullInput(projection_weights)) { + NN_CHECK_EQ(NumDimensions(projection_weights), 2); + NN_CHECK_EQ(SizeOfDimension(projection_weights, 0), n_output); + NN_CHECK_EQ(SizeOfDimension(projection_weights, 1), n_cell); } - if (!IsNullInput(projection_bias_)) { - NN_CHECK_EQ(NumDimensions(projection_bias_), 1); - NN_CHECK_EQ(SizeOfDimension(projection_bias_, 0), n_output); + if (!IsNullInput(projection_bias)) { + NN_CHECK_EQ(NumDimensions(projection_bias), 1); + NN_CHECK_EQ(SizeOfDimension(projection_bias, 0), n_output); } // Making sure the projection tensors are consistent: @@ -210,31 +236,31 @@ // 2) If projection weight is present, then projection bias is optional. // TODO: make sure this is correct. const bool projecton_tensors_consistent = - (!IsNullInput(projection_weights_) || IsNullInput(projection_bias_)); + (!IsNullInput(projection_weights) || IsNullInput(projection_bias)); NN_CHECK(projecton_tensors_consistent == true); - if (!IsNullInput(input_layer_norm_weights_)) { - NN_CHECK_EQ(NumDimensions(input_layer_norm_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(input_layer_norm_weights_, 0), n_cell); + if (!IsNullInput(input_layer_norm_weights)) { + NN_CHECK_EQ(NumDimensions(input_layer_norm_weights), 1); + NN_CHECK_EQ(SizeOfDimension(input_layer_norm_weights, 0), n_cell); } - if (!IsNullInput(forget_layer_norm_weights_)) { - NN_CHECK_EQ(NumDimensions(forget_layer_norm_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(forget_layer_norm_weights_, 0), n_cell); + if (!IsNullInput(forget_layer_norm_weights)) { + NN_CHECK_EQ(NumDimensions(forget_layer_norm_weights), 1); + NN_CHECK_EQ(SizeOfDimension(forget_layer_norm_weights, 0), n_cell); } - if (!IsNullInput(cell_layer_norm_weights_)) { - NN_CHECK_EQ(NumDimensions(cell_layer_norm_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(cell_layer_norm_weights_, 0), n_cell); + if (!IsNullInput(cell_layer_norm_weights)) { + NN_CHECK_EQ(NumDimensions(cell_layer_norm_weights), 1); + NN_CHECK_EQ(SizeOfDimension(cell_layer_norm_weights, 0), n_cell); } - if (!IsNullInput(output_layer_norm_weights_)) { - NN_CHECK_EQ(NumDimensions(output_layer_norm_weights_), 1); - NN_CHECK_EQ(SizeOfDimension(output_layer_norm_weights_, 0), n_cell); + if (!IsNullInput(output_layer_norm_weights)) { + NN_CHECK_EQ(NumDimensions(output_layer_norm_weights), 1); + NN_CHECK_EQ(SizeOfDimension(output_layer_norm_weights, 0), n_cell); } const bool layer_norm_weights_all_or_none = - (IsNullInput(input_layer_norm_weights_) && IsNullInput(forget_layer_norm_weights_) && - IsNullInput(cell_layer_norm_weights_) && IsNullInput(input_layer_norm_weights_)) || - (!IsNullInput(input_layer_norm_weights_) && !IsNullInput(forget_layer_norm_weights_) && - !IsNullInput(cell_layer_norm_weights_) && !IsNullInput(input_layer_norm_weights_)); + (IsNullInput(input_layer_norm_weights) && IsNullInput(forget_layer_norm_weights) && + IsNullInput(cell_layer_norm_weights) && IsNullInput(input_layer_norm_weights)) || + (!IsNullInput(input_layer_norm_weights) && !IsNullInput(forget_layer_norm_weights) && + !IsNullInput(cell_layer_norm_weights) && !IsNullInput(input_layer_norm_weights)); NN_CHECK(layer_norm_weights_all_or_none); return true; @@ -263,7 +289,15 @@ const uint32_t n_output = SizeOfDimension(recurrent_to_output_weights_, 1); // Check that input tensor dimensions matches with each other. - if (!CheckInputTensorDimensions(operation, operands, n_input, n_output, n_cell)) { + if (!CheckInputTensorDimensions( + input_, input_to_input_weights_, input_to_forget_weights_, input_to_cell_weights_, + input_to_output_weights_, recurrent_to_input_weights_, recurrent_to_forget_weights_, + recurrent_to_cell_weights_, recurrent_to_output_weights_, cell_to_input_weights_, + cell_to_forget_weights_, cell_to_output_weights_, input_gate_bias_, + forget_gate_bias_, cell_bias_, output_gate_bias_, projection_weights_, + projection_bias_, input_layer_norm_weights_, forget_layer_norm_weights_, + cell_layer_norm_weights_, output_layer_norm_weights_, n_input, n_output, n_cell, + ¶ms_)) { return false; } @@ -285,8 +319,7 @@ cellStateShape->offset = inputShape.offset; cellStateShape->scale = inputShape.scale; - const bool use_cifg = IsNullInput(input_to_input_weights_); - if (use_cifg) { + if (params_.use_cifg) { // Reserving space for Cell, Forget, Output gates scratchShape->dimensions = {n_batch, n_cell * 3}; } else { @@ -300,13 +333,16 @@ return true; } -bool LSTMCell::EvalFloat32( - const float* input_buffer, const float* input_to_input_weights_buffer, - const float* input_to_forget_weights_buffer, const float* input_to_cell_weights_buffer, - const float* input_to_output_weights_buffer, const float* recurrent_to_input_weights_buffer, +// static +bool LSTMCell::LSTMStep( + const LSTMParams& params, const float* input_buffer, const Shape& input_shape, + const float* input_to_input_weights_buffer, const float* input_to_forget_weights_buffer, + const float* input_to_cell_weights_buffer, const float* input_to_output_weights_buffer, + const Shape& input_to_output_weights_shape, const float* recurrent_to_input_weights_buffer, const float* recurrent_to_forget_weights_buffer, const float* recurrent_to_cell_weights_buffer, - const float* recurrent_to_output_weights_buffer, const float* cell_to_input_weights_buffer, + const float* recurrent_to_output_weights_buffer, + const Shape& recurrent_to_output_weights_shape, const float* cell_to_input_weights_buffer, const float* cell_to_forget_weights_buffer, const float* cell_to_output_weights_buffer, const float* input_gate_bias_buffer, const float* forget_gate_bias_buffer, const float* cell_bias_buffer, const float* output_gate_bias_buffer, @@ -316,26 +352,20 @@ const float* cell_layer_norm_weights_buffer, const float* output_layer_norm_weights_buffer, float* output_state_out_buffer, float* cell_state_out_buffer, float* output_buffer, float* scratch_buffer_buffer) { - NNTRACE_COMP("LSTMCell::Eval"); + NNTRACE_COMP("LSTMCell::LSTMStep"); - const uint32_t n_batch = input_->shape().dimensions[0]; - const uint32_t n_input = input_->shape().dimensions[1]; + const uint32_t n_batch = input_shape.dimensions[0]; + const uint32_t n_input = input_shape.dimensions[1]; // n_cell and n_output will be the same size when there is no projection. - const uint32_t n_cell = input_to_output_weights_->shape().dimensions[0]; - const uint32_t n_output = recurrent_to_output_weights_->shape().dimensions[1]; - - // Since we have already checked that weights are all there or none, we can - // check the existence of only one to the get the condition. - const bool use_cifg = IsNullInput(input_to_input_weights_); - const bool use_peephole = !IsNullInput(cell_to_output_weights_); - const bool use_layer_norm = !IsNullInput(input_layer_norm_weights_); + const uint32_t n_cell = input_to_output_weights_shape.dimensions[0]; + const uint32_t n_output = recurrent_to_output_weights_shape.dimensions[1]; // Index the scratch buffers pointers to the global scratch buffer. float* input_gate_scratch = nullptr; float* cell_scratch = nullptr; float* forget_gate_scratch = nullptr; float* output_gate_scratch = nullptr; - if (use_cifg) { + if (params.use_cifg) { cell_scratch = scratch_buffer_buffer; forget_gate_scratch = cell_scratch + n_cell * n_batch; output_gate_scratch = cell_scratch + 2 * n_cell * n_batch; @@ -346,9 +376,9 @@ output_gate_scratch = input_gate_scratch + 3 * n_cell * n_batch; } - if (!use_layer_norm) { + if (!params.use_layer_norm) { // Initialize scratch buffers with bias. - if (!use_cifg) { + if (!params.use_cifg) { tflite::tensor_utils::VectorBatchVectorAssign(input_gate_bias_buffer, n_cell, n_batch, input_gate_scratch); } @@ -360,7 +390,7 @@ output_gate_scratch); } else { // Initialize scratch buffers with zeroes. - if (!use_cifg) { + if (!params.use_cifg) { tflite::tensor_utils::ZeroVector(input_gate_scratch, n_cell * n_batch); } tflite::tensor_utils::ZeroVector(forget_gate_scratch, n_cell * n_batch); @@ -369,7 +399,7 @@ } // For each batch and cell: compute input_weight * input. - if (!use_cifg) { + if (!params.use_cifg) { tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( input_to_input_weights_buffer, n_cell, n_input, input_buffer, n_batch, input_gate_scratch, /*result_stride*/ 1); @@ -385,7 +415,7 @@ output_gate_scratch, /*result_stride*/ 1); // For each batch and cell: compute recurrent_weight * output_state. - if (!use_cifg) { + if (!params.use_cifg) { tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( recurrent_to_input_weights_buffer, n_cell, n_output, output_state_in_buffer, n_batch, input_gate_scratch, @@ -402,13 +432,13 @@ output_gate_scratch, /*result_stride*/ 1); // For each batch and cell: update input gate. - if (!use_cifg) { - if (use_peephole) { + if (!params.use_cifg) { + if (params.use_peephole) { tflite::tensor_utils::VectorBatchVectorCwiseProductAccumulate( cell_to_input_weights_buffer, n_cell, cell_state_in_buffer, n_batch, input_gate_scratch); } - if (use_layer_norm) { + if (params.use_layer_norm) { tflite::tensor_utils::MeanStddevNormalization(input_gate_scratch, input_gate_scratch, n_cell, n_batch, kLayerNormEpsilon); tflite::tensor_utils::VectorBatchVectorCwiseProduct(input_layer_norm_weights_buffer, @@ -422,12 +452,12 @@ } // For each batch and cell: update forget gate. - if (use_peephole) { + if (params.use_peephole) { tflite::tensor_utils::VectorBatchVectorCwiseProductAccumulate(cell_to_forget_weights_buffer, n_cell, cell_state_in_buffer, n_batch, forget_gate_scratch); } - if (use_layer_norm) { + if (params.use_layer_norm) { tflite::tensor_utils::MeanStddevNormalization(forget_gate_scratch, forget_gate_scratch, n_cell, n_batch, kLayerNormEpsilon); tflite::tensor_utils::VectorBatchVectorCwiseProduct(forget_layer_norm_weights_buffer, @@ -440,7 +470,7 @@ forget_gate_scratch); // For each batch and cell: update the cell. - if (use_layer_norm) { + if (params.use_layer_norm) { tflite::tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell, n_batch, kLayerNormEpsilon); tflite::tensor_utils::VectorBatchVectorCwiseProduct(cell_layer_norm_weights_buffer, n_cell, @@ -449,9 +479,9 @@ } tflite::tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_in_buffer, n_batch * n_cell, cell_state_out_buffer); - tflite::tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, - params_.activation_, cell_scratch); - if (use_cifg) { + tflite::tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, params.activation, + cell_scratch); + if (params.use_cifg) { tflite::tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell, forget_gate_scratch); tflite::tensor_utils::VectorVectorCwiseProductAccumulate( @@ -460,18 +490,18 @@ tflite::tensor_utils::VectorVectorCwiseProductAccumulate( cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_out_buffer); } - if (params_.cell_clip_ > 0.0) { - tflite::tensor_utils::ClipVector(cell_state_out_buffer, n_batch * n_cell, - params_.cell_clip_, cell_state_out_buffer); + if (params.cell_clip > 0.0) { + tflite::tensor_utils::ClipVector(cell_state_out_buffer, n_batch * n_cell, params.cell_clip, + cell_state_out_buffer); } // For each batch and cell: update the output gate. - if (use_peephole) { + if (params.use_peephole) { tflite::tensor_utils::VectorBatchVectorCwiseProductAccumulate(cell_to_output_weights_buffer, n_cell, cell_state_out_buffer, n_batch, output_gate_scratch); } - if (use_layer_norm) { + if (params.use_layer_norm) { tflite::tensor_utils::MeanStddevNormalization(output_gate_scratch, output_gate_scratch, n_cell, n_batch, kLayerNormEpsilon); tflite::tensor_utils::VectorBatchVectorCwiseProduct(output_layer_norm_weights_buffer, @@ -483,15 +513,13 @@ tflite::tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, output_gate_scratch); tflite::tensor_utils::ApplyActivationToVector(cell_state_out_buffer, n_batch * n_cell, - params_.activation_, cell_scratch); + params.activation, cell_scratch); tflite::tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch, n_batch * n_cell, output_gate_scratch); // For each batch: update the projection and output_state. - const bool use_projection_weight = (projection_weights_->lifetime != OperandLifeTime::NO_VALUE); - const bool use_projection_bias = (projection_bias_->lifetime != OperandLifeTime::NO_VALUE); - if (use_projection_weight) { - if (use_projection_bias) { + if (params.use_projection_weight) { + if (params.use_projection_bias) { tflite::tensor_utils::VectorBatchVectorAssign(projection_bias_buffer, n_output, n_batch, output_buffer); } else { @@ -501,8 +529,8 @@ projection_weights_buffer, n_output, n_cell, output_gate_scratch, n_batch, output_buffer, /*result_stride*/ 1); - if (params_.proj_clip_ > 0.0) { - tflite::tensor_utils::ClipVector(output_buffer, n_batch * n_output, params_.proj_clip_, + if (params.proj_clip > 0.0) { + tflite::tensor_utils::ClipVector(output_buffer, n_batch * n_output, params.proj_clip, output_buffer); } } else { @@ -515,37 +543,33 @@ bool LSTMCell::Eval() { switch (input_->type) { case OperandType::TENSOR_FLOAT32: { - // clang-format off - EvalFloat32( - GetBuffer<const float>(input_), - GetBuffer<const float>(input_to_input_weights_), - GetBuffer<const float>(input_to_forget_weights_), - GetBuffer<const float>(input_to_cell_weights_), - GetBuffer<const float>(input_to_output_weights_), - GetBuffer<const float>(recurrent_to_input_weights_), - GetBuffer<const float>(recurrent_to_forget_weights_), - GetBuffer<const float>(recurrent_to_cell_weights_), - GetBuffer<const float>(recurrent_to_output_weights_), - GetBuffer<const float>(cell_to_input_weights_), - GetBuffer<const float>(cell_to_forget_weights_), - GetBuffer<const float>(cell_to_output_weights_), - GetBuffer<const float>(input_gate_bias_), - GetBuffer<const float>(forget_gate_bias_), - GetBuffer<const float>(cell_bias_), - GetBuffer<const float>(output_gate_bias_), - GetBuffer<const float>(projection_weights_), - GetBuffer<const float>(projection_bias_), - GetBuffer<const float>(output_state_in_), - GetBuffer<const float>(cell_state_in_), - GetBuffer<const float>(input_layer_norm_weights_), - GetBuffer<const float>(forget_layer_norm_weights_), - GetBuffer<const float>(cell_layer_norm_weights_), - GetBuffer<const float>(output_layer_norm_weights_), - GetBuffer<float>(output_state_out_), - GetBuffer<float>(cell_state_out_), - GetBuffer<float>(output_), - GetBuffer<float>(scratch_buffer_)); - // clang-format on + LSTMStep(params_, GetBuffer<const float>(input_), input_->shape(), + GetBuffer<const float>(input_to_input_weights_), + GetBuffer<const float>(input_to_forget_weights_), + GetBuffer<const float>(input_to_cell_weights_), + GetBuffer<const float>(input_to_output_weights_), + input_to_output_weights_->shape(), + GetBuffer<const float>(recurrent_to_input_weights_), + GetBuffer<const float>(recurrent_to_forget_weights_), + GetBuffer<const float>(recurrent_to_cell_weights_), + GetBuffer<const float>(recurrent_to_output_weights_), + recurrent_to_output_weights_->shape(), + GetBuffer<const float>(cell_to_input_weights_), + GetBuffer<const float>(cell_to_forget_weights_), + GetBuffer<const float>(cell_to_output_weights_), + GetBuffer<const float>(input_gate_bias_), + GetBuffer<const float>(forget_gate_bias_), GetBuffer<const float>(cell_bias_), + GetBuffer<const float>(output_gate_bias_), + GetBuffer<const float>(projection_weights_), + GetBuffer<const float>(projection_bias_), + GetBuffer<const float>(output_state_in_), + GetBuffer<const float>(cell_state_in_), + GetBuffer<const float>(input_layer_norm_weights_), + GetBuffer<const float>(forget_layer_norm_weights_), + GetBuffer<const float>(cell_layer_norm_weights_), + GetBuffer<const float>(output_layer_norm_weights_), + GetBuffer<float>(output_state_out_), GetBuffer<float>(cell_state_out_), + GetBuffer<float>(output_), GetBuffer<float>(scratch_buffer_)); } break; case OperandType::TENSOR_FLOAT16: { std::vector<float> input_float32(getNumberOfElements(input_->shape())); @@ -699,37 +723,22 @@ getNumberOfElements(scratch_buffer_->shape())); convertFloat16ToFloat32(GetBuffer<_Float16>(scratch_buffer_), &scratch_buffer_float32); - // clang-format off - EvalFloat32( - input_float32.data(), - input_to_input_weights_buffer, - input_to_forget_weights_float32.data(), - input_to_cell_weights_float32.data(), - input_to_output_weights_float32.data(), - recurrent_to_input_weights_buffer, - recurrent_to_forget_weights_float32.data(), - recurrent_to_cell_weights_float32.data(), - recurrent_to_output_weights_float32.data(), - cell_to_input_weights_buffer, - cell_to_forget_weights_buffer, - cell_to_output_weights_buffer, - input_gate_bias_buffer, - forget_gate_bias_float32.data(), - cell_bias_float32.data(), - output_gate_bias_float32.data(), - projection_weights_buffer, - projection_bias_buffer, - output_state_in_float32.data(), - cell_state_in_float32.data(), - input_layer_norm_weights_buffer, - forget_layer_norm_weights_buffer, - cell_layer_norm_weights_buffer, - output_layer_norm_weights_buffer, - output_state_out_float32.data(), - cell_state_out_float32.data(), - output_float32.data(), - scratch_buffer_float32.data()); - // clang-format on + LSTMStep(params_, input_float32.data(), input_->shape(), input_to_input_weights_buffer, + input_to_forget_weights_float32.data(), input_to_cell_weights_float32.data(), + input_to_output_weights_float32.data(), input_to_output_weights_->shape(), + recurrent_to_input_weights_buffer, recurrent_to_forget_weights_float32.data(), + recurrent_to_cell_weights_float32.data(), + recurrent_to_output_weights_float32.data(), + recurrent_to_output_weights_->shape(), cell_to_input_weights_buffer, + cell_to_forget_weights_buffer, cell_to_output_weights_buffer, + input_gate_bias_buffer, forget_gate_bias_float32.data(), + cell_bias_float32.data(), output_gate_bias_float32.data(), + projection_weights_buffer, projection_bias_buffer, + output_state_in_float32.data(), cell_state_in_float32.data(), + input_layer_norm_weights_buffer, forget_layer_norm_weights_buffer, + cell_layer_norm_weights_buffer, output_layer_norm_weights_buffer, + output_state_out_float32.data(), cell_state_out_float32.data(), + output_float32.data(), scratch_buffer_float32.data()); convertFloat32ToFloat16(output_state_out_float32, GetBuffer<_Float16>(output_state_out_));