Implements BIDIRECTIONAL_SEQUENCE_LSTM operation. Test: NeuralNetworksTest_static --gtest_filter=GeneratedTests.*lstm Bug: 113559542 Change-Id: If30e31c851bfbd97445710d8e1998306a551ac08 Merged-In: If30e31c851bfbd97445710d8e1998306a551ac08 (cherry picked from commit be339f503051e72ed63e88a5645af15b02a44478)
diff --git a/common/operations/BidirectionalSequenceLSTM.cpp b/common/operations/BidirectionalSequenceLSTM.cpp new file mode 100644 index 0000000..06356f9 --- /dev/null +++ b/common/operations/BidirectionalSequenceLSTM.cpp
@@ -0,0 +1,428 @@ +/* + * Copyright (C) 2019 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "BidirectionalSequenceLSTM.h" + +#include "CpuExecutor.h" +#include "CpuOperationUtils.h" +#include "HalInterfaces.h" + +#include "Tracing.h" + +namespace android { +namespace nn { + +namespace { + +template <typename T> +inline T* GetBuffer(RunTimeOperandInfo* operand) { + return reinterpret_cast<T*>(operand->buffer); +} + +template <typename T> +inline const T* GetBuffer(const RunTimeOperandInfo* operand) { + return reinterpret_cast<const T*>(operand->buffer); +} + +template <typename T> +inline const T* GetOptionalBuffer(const RunTimeOperandInfo* operand) { + return !IsNullInput(operand) ? reinterpret_cast<const T*>(operand->buffer) : nullptr; +} + +} // anonymous namespace + +BidirectionalSequenceLSTM::BidirectionalSequenceLSTM(const Operation& operation, + std::vector<RunTimeOperandInfo>& operands) { + input_ = GetInput(operation, operands, kInputTensor); + + fw_input_to_input_weights_ = + GetInput(operation, operands, kFwInputToInputWeightsTensor); // optional + fw_input_to_forget_weights_ = GetInput(operation, operands, kFwInputToForgetWeightsTensor); + fw_input_to_cell_weights_ = GetInput(operation, operands, kFwInputToCellWeightsTensor); + fw_input_to_output_weights_ = GetInput(operation, operands, kFwInputToOutputWeightsTensor); + + fw_recurrent_to_input_weights_ = + GetInput(operation, operands, kFwRecurrentToInputWeightsTensor); // optional + fw_recurrent_to_forget_weights_ = + GetInput(operation, operands, kFwRecurrentToForgetWeightsTensor); + fw_recurrent_to_cell_weights_ = GetInput(operation, operands, kFwRecurrentToCellWeightsTensor); + fw_recurrent_to_output_weights_ = + GetInput(operation, operands, kFwRecurrentToOutputWeightsTensor); + + fw_cell_to_input_weights_ = + GetInput(operation, operands, kFwCellToInputWeightsTensor); // optional + fw_cell_to_forget_weights_ = + GetInput(operation, operands, kFwCellToForgetWeightsTensor); // optional + fw_cell_to_output_weights_ = + GetInput(operation, operands, kFwCellToOutputWeightsTensor); // optional + + fw_input_gate_bias_ = GetInput(operation, operands, kFwInputGateBiasTensor); + fw_forget_gate_bias_ = GetInput(operation, operands, kFwForgetGateBiasTensor); + fw_cell_bias_ = GetInput(operation, operands, kFwCellGateBiasTensor); + fw_output_gate_bias_ = GetInput(operation, operands, kFwOutputGateBiasTensor); + + fw_projection_weights_ = GetInput(operation, operands, kFwProjectionWeightsTensor); // optional + fw_projection_bias_ = GetInput(operation, operands, kFwProjectionBiasTensor); // optional + + fw_activation_state_ = GetInput(operation, operands, kFwInputActivationStateTensor); + fw_cell_state_ = GetInput(operation, operands, kFwInputCellStateTensor); + + bw_input_to_input_weights_ = + GetInput(operation, operands, kBwInputToInputWeightsTensor); // optional + bw_input_to_forget_weights_ = GetInput(operation, operands, kBwInputToForgetWeightsTensor); + bw_input_to_cell_weights_ = GetInput(operation, operands, kBwInputToCellWeightsTensor); + bw_input_to_output_weights_ = GetInput(operation, operands, kBwInputToOutputWeightsTensor); + + bw_recurrent_to_input_weights_ = + GetInput(operation, operands, kBwRecurrentToInputWeightsTensor); // optional + bw_recurrent_to_forget_weights_ = + GetInput(operation, operands, kBwRecurrentToForgetWeightsTensor); + bw_recurrent_to_cell_weights_ = GetInput(operation, operands, kBwRecurrentToCellWeightsTensor); + bw_recurrent_to_output_weights_ = + GetInput(operation, operands, kBwRecurrentToOutputWeightsTensor); + + bw_cell_to_input_weights_ = + GetInput(operation, operands, kBwCellToInputWeightsTensor); // optional + bw_cell_to_forget_weights_ = + GetInput(operation, operands, kBwCellToForgetWeightsTensor); // optional + bw_cell_to_output_weights_ = + GetInput(operation, operands, kBwCellToOutputWeightsTensor); // optional + + bw_input_gate_bias_ = GetInput(operation, operands, kBwInputGateBiasTensor); + bw_forget_gate_bias_ = GetInput(operation, operands, kBwForgetGateBiasTensor); + bw_cell_bias_ = GetInput(operation, operands, kBwCellGateBiasTensor); + bw_output_gate_bias_ = GetInput(operation, operands, kBwOutputGateBiasTensor); + + bw_projection_weights_ = GetInput(operation, operands, kBwProjectionWeightsTensor); // optional + bw_projection_bias_ = GetInput(operation, operands, kBwProjectionBiasTensor); // optional + + bw_activation_state_ = GetInput(operation, operands, kBwInputActivationStateTensor); + bw_cell_state_ = GetInput(operation, operands, kBwInputCellStateTensor); + + aux_input_ = GetInput(operation, operands, kAuxInputTensor); + fw_aux_input_to_input_weights_ = GetInput(operation, operands, kFwAuxInputToInputWeightsTensor); + fw_aux_input_to_forget_weights_ = + GetInput(operation, operands, kFwAuxInputToForgetWeightsTensor); + fw_aux_input_to_cell_weights_ = GetInput(operation, operands, kFwAuxInputToCellWeightsTensor); + fw_aux_input_to_output_weights_ = + GetInput(operation, operands, kFwAuxInputToOutputWeightsTensor); + bw_aux_input_to_input_weights_ = GetInput(operation, operands, kBwAuxInputToInputWeightsTensor); + bw_aux_input_to_forget_weights_ = + GetInput(operation, operands, kBwAuxInputToForgetWeightsTensor); + bw_aux_input_to_cell_weights_ = GetInput(operation, operands, kBwAuxInputToCellWeightsTensor); + bw_aux_input_to_output_weights_ = + GetInput(operation, operands, kBwAuxInputToOutputWeightsTensor); + + 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)); + } else { + params_.cell_clip = static_cast<float>( + getScalarData<_Float16>(*GetInput(operation, operands, kCellClipParam))); + params_.proj_clip = static_cast<float>( + getScalarData<_Float16>(*GetInput(operation, operands, kProjClipParam))); + } + params_.merge_outputs = getScalarData<bool>(*GetInput(operation, operands, kMergeOutputsParam)); + params_.time_major = getScalarData<bool>(*GetInput(operation, operands, kTimeMajorParam)); + params_.use_layer_norm = false; + + fw_output_ = GetOutput(operation, operands, kFwOutputTensor); + bw_output_ = GetOutput(operation, operands, kBwOutputTensor); +} + +bool BidirectionalSequenceLSTM::Prepare(const Operation& operation, + std::vector<RunTimeOperandInfo>& operands, + Shape* fwOutputShape, Shape* bwOutputShape) { + // Inferring batch size, number of outputs and number of cells from the + // input tensors. + NN_CHECK(NumDimensions(input_) == 3); + const uint32_t max_time = SizeOfDimension(input_, 0); + const uint32_t n_batch = SizeOfDimension(input_, 1); + const uint32_t n_input = SizeOfDimension(input_, 2); + + const uint32_t n_fw_cell = SizeOfDimension(fw_input_to_output_weights_, 0); + NN_CHECK_EQ(NumDimensions(fw_input_to_output_weights_), 2); + NN_CHECK_EQ(SizeOfDimension(fw_input_to_output_weights_, 1), n_input); + + NN_CHECK_EQ(NumDimensions(fw_recurrent_to_output_weights_), 2); + NN_CHECK_EQ(SizeOfDimension(fw_recurrent_to_output_weights_, 0), n_fw_cell); + const uint32_t n_fw_output = SizeOfDimension(fw_recurrent_to_output_weights_, 1); + + RunTimeOperandInfo nullOpInfo; + nullOpInfo.lifetime = OperandLifeTime::NO_VALUE; + // Check that input tensor dimensions matches with each other. + if (!LSTMCell::CheckInputTensorDimensions( + input_, fw_input_to_input_weights_, fw_input_to_forget_weights_, + fw_input_to_cell_weights_, fw_input_to_output_weights_, + fw_recurrent_to_input_weights_, fw_recurrent_to_forget_weights_, + fw_recurrent_to_cell_weights_, fw_recurrent_to_output_weights_, + fw_cell_to_input_weights_, fw_cell_to_forget_weights_, fw_cell_to_output_weights_, + fw_input_gate_bias_, fw_forget_gate_bias_, fw_cell_bias_, fw_output_gate_bias_, + fw_projection_weights_, fw_projection_bias_, + /*input_layer_norm_weights=*/&nullOpInfo, + /*forget_layer_norm_weights=*/&nullOpInfo, /*cell_layer_norm_weights=*/&nullOpInfo, + /*output_layer_norm_weights=*/&nullOpInfo, n_input, n_fw_output, n_fw_cell, + ¶ms_)) { + return false; + } + + const bool aux_inputs_all_or_none = + (!IsNullInput(aux_input_) && !IsNullInput(fw_aux_input_to_cell_weights_) && + !IsNullInput(fw_aux_input_to_forget_weights_) && + !IsNullInput(fw_aux_input_to_output_weights_) && + !IsNullInput(bw_aux_input_to_cell_weights_) && + !IsNullInput(bw_aux_input_to_forget_weights_) && + !IsNullInput(bw_aux_input_to_output_weights_)) || + (IsNullInput(fw_aux_input_to_cell_weights_) && + IsNullInput(fw_aux_input_to_forget_weights_) && + IsNullInput(fw_aux_input_to_output_weights_) && + IsNullInput(bw_aux_input_to_cell_weights_) && + IsNullInput(bw_aux_input_to_forget_weights_) && + IsNullInput(bw_aux_input_to_output_weights_)); + NN_CHECK(aux_inputs_all_or_none); + if (!IsNullInput(aux_input_)) { + // Check that aux_input has the same dimensions (except last) as the input. + NN_CHECK_EQ(aux_input_->shape().dimensions[0], input_->shape().dimensions[0]); + NN_CHECK_EQ(aux_input_->shape().dimensions[1], input_->shape().dimensions[1]); + } + + const uint32_t n_bw_cell = SizeOfDimension(bw_input_to_output_weights_, 0); + NN_CHECK_EQ(NumDimensions(bw_input_to_output_weights_), 2); + NN_CHECK_EQ(SizeOfDimension(bw_input_to_output_weights_, 1), n_input); + + NN_CHECK_EQ(NumDimensions(bw_recurrent_to_output_weights_), 2); + NN_CHECK_EQ(SizeOfDimension(bw_recurrent_to_output_weights_, 0), n_bw_cell); + const uint32_t n_bw_output = SizeOfDimension(bw_recurrent_to_output_weights_, 1); + + const Shape& inputShape = input_->shape(); + fwOutputShape->type = inputShape.type; + fwOutputShape->offset = inputShape.offset; + fwOutputShape->scale = inputShape.scale; + fwOutputShape->dimensions.resize(3); + fwOutputShape->dimensions[0] = params_.time_major ? max_time : n_batch; + fwOutputShape->dimensions[1] = params_.time_major ? n_batch : max_time; + fwOutputShape->dimensions[2] = params_.merge_outputs ? n_fw_output + n_bw_output : n_fw_output; + + // Check that input tensor dimensions matches with each other. + if (!LSTMCell::CheckInputTensorDimensions( + input_, bw_input_to_input_weights_, bw_input_to_forget_weights_, + bw_input_to_cell_weights_, bw_input_to_output_weights_, + bw_recurrent_to_input_weights_, bw_recurrent_to_forget_weights_, + bw_recurrent_to_cell_weights_, bw_recurrent_to_output_weights_, + bw_cell_to_input_weights_, bw_cell_to_forget_weights_, bw_cell_to_output_weights_, + bw_input_gate_bias_, bw_forget_gate_bias_, bw_cell_bias_, bw_output_gate_bias_, + bw_projection_weights_, bw_projection_bias_, + /*input_layer_norm_weights=*/&nullOpInfo, + /*forget_layer_norm_weights=*/&nullOpInfo, /*cell_layer_norm_weights=*/&nullOpInfo, + /*output_layer_norm_weights=*/&nullOpInfo, n_input, n_bw_output, n_bw_cell, + ¶ms_)) { + return false; + } + + if (!params_.merge_outputs) { + bwOutputShape->type = inputShape.type; + bwOutputShape->offset = inputShape.offset; + bwOutputShape->scale = inputShape.scale; + bwOutputShape->dimensions.resize(3); + bwOutputShape->dimensions[0] = params_.time_major ? max_time : n_batch; + bwOutputShape->dimensions[1] = params_.time_major ? n_batch : max_time; + bwOutputShape->dimensions[2] = n_bw_output; + } + + if (params_.use_cifg) { + fw_scratch_shape_.dimensions = {n_batch, n_fw_cell * 3}; + bw_scratch_shape_.dimensions = {n_batch, n_bw_cell * 3}; + } else { + fw_scratch_shape_.dimensions = {n_batch, n_fw_cell * 4}; + bw_scratch_shape_.dimensions = {n_batch, n_bw_cell * 4}; + } + fw_scratch_shape_.type = bw_scratch_shape_.type = inputShape.type; + fw_scratch_shape_.offset = bw_scratch_shape_.offset = inputShape.offset; + fw_scratch_shape_.scale = bw_scratch_shape_.scale = inputShape.scale; + + return true; +} + +bool BidirectionalSequenceLSTM::Eval() { + const uint32_t n_fw_output = SizeOfDimension(fw_recurrent_to_output_weights_, 1); + switch (input_->type) { + case OperandType::TENSOR_FLOAT32: { + std::vector<float> fw_scratch_buffer(getNumberOfElements(fw_scratch_shape_)); + const bool kForwardSequence = true; + LSTMCell::LSTMEvalFloat32( + params_, GetBuffer<const float>(input_), input_->shape(), + GetBuffer<const float>(fw_input_to_input_weights_), + GetBuffer<const float>(fw_input_to_forget_weights_), + GetBuffer<const float>(fw_input_to_cell_weights_), + GetBuffer<const float>(fw_input_to_output_weights_), + fw_input_to_output_weights_->shape(), + GetBuffer<const float>(fw_recurrent_to_input_weights_), + GetBuffer<const float>(fw_recurrent_to_forget_weights_), + GetBuffer<const float>(fw_recurrent_to_cell_weights_), + GetBuffer<const float>(fw_recurrent_to_output_weights_), + fw_recurrent_to_output_weights_->shape(), + GetBuffer<const float>(fw_cell_to_input_weights_), + GetBuffer<const float>(fw_cell_to_forget_weights_), + GetBuffer<const float>(fw_cell_to_output_weights_), + GetOptionalBuffer<const float>(aux_input_), aux_input_->shape(), + GetOptionalBuffer<const float>(fw_aux_input_to_input_weights_), + GetOptionalBuffer<const float>(fw_aux_input_to_forget_weights_), + GetOptionalBuffer<const float>(fw_aux_input_to_cell_weights_), + GetOptionalBuffer<const float>(fw_aux_input_to_output_weights_), + GetBuffer<const float>(fw_input_gate_bias_), + GetBuffer<const float>(fw_forget_gate_bias_), + GetBuffer<const float>(fw_cell_bias_), + GetBuffer<const float>(fw_output_gate_bias_), + GetBuffer<const float>(fw_projection_weights_), + GetBuffer<const float>(fw_projection_bias_), + GetBuffer<const float>(fw_activation_state_), + GetBuffer<const float>(fw_cell_state_), + /*input_layer_norm_weights=*/nullptr, + /*forget_layer_norm_weights=*/nullptr, + /*cell_layer_norm_weights=*/nullptr, + /*output_layer_norm_weights=*/nullptr, GetBuffer<float>(fw_activation_state_), + GetBuffer<float>(fw_cell_state_), GetBuffer<float>(fw_output_), + fw_scratch_buffer.data(), kForwardSequence, params_.time_major); + + std::vector<float> bw_scratch_buffer(getNumberOfElements(bw_scratch_shape_)); + const bool kBackwardSequence = false; + LSTMCell::LSTMEvalFloat32( + params_, GetBuffer<const float>(input_), input_->shape(), + GetBuffer<const float>(bw_input_to_input_weights_), + GetBuffer<const float>(bw_input_to_forget_weights_), + GetBuffer<const float>(bw_input_to_cell_weights_), + GetBuffer<const float>(bw_input_to_output_weights_), + bw_input_to_output_weights_->shape(), + GetBuffer<const float>(bw_recurrent_to_input_weights_), + GetBuffer<const float>(bw_recurrent_to_forget_weights_), + GetBuffer<const float>(bw_recurrent_to_cell_weights_), + GetBuffer<const float>(bw_recurrent_to_output_weights_), + bw_recurrent_to_output_weights_->shape(), + GetBuffer<const float>(bw_cell_to_input_weights_), + GetBuffer<const float>(bw_cell_to_forget_weights_), + GetBuffer<const float>(bw_cell_to_output_weights_), + GetOptionalBuffer<const float>(aux_input_), aux_input_->shape(), + GetOptionalBuffer<const float>(bw_aux_input_to_input_weights_), + GetOptionalBuffer<const float>(bw_aux_input_to_forget_weights_), + GetOptionalBuffer<const float>(bw_aux_input_to_cell_weights_), + GetOptionalBuffer<const float>(bw_aux_input_to_output_weights_), + GetBuffer<const float>(bw_input_gate_bias_), + GetBuffer<const float>(bw_forget_gate_bias_), + GetBuffer<const float>(bw_cell_bias_), + GetBuffer<const float>(bw_output_gate_bias_), + GetBuffer<const float>(bw_projection_weights_), + GetBuffer<const float>(bw_projection_bias_), + GetBuffer<const float>(bw_activation_state_), + GetBuffer<const float>(bw_cell_state_), + /*input_layer_norm_weights=*/nullptr, + /*forget_layer_norm_weights=*/nullptr, + /*cell_layer_norm_weights=*/nullptr, + /*output_layer_norm_weights=*/nullptr, GetBuffer<float>(bw_activation_state_), + GetBuffer<float>(bw_cell_state_), + params_.merge_outputs ? GetBuffer<float>(fw_output_) + n_fw_output + : GetBuffer<float>(bw_output_), + bw_scratch_buffer.data(), kBackwardSequence, params_.time_major); + } break; + case OperandType::TENSOR_FLOAT16: { + std::vector<_Float16> fw_scratch_buffer(getNumberOfElements(fw_scratch_shape_)); + const bool kForwardSequence = true; + LSTMCell::LSTMEvalFloat16( + params_, GetBuffer<const _Float16>(input_), input_->shape(), + GetBuffer<const _Float16>(fw_input_to_input_weights_), + GetBuffer<const _Float16>(fw_input_to_forget_weights_), + GetBuffer<const _Float16>(fw_input_to_cell_weights_), + GetBuffer<const _Float16>(fw_input_to_output_weights_), + fw_input_to_output_weights_->shape(), + GetBuffer<const _Float16>(fw_recurrent_to_input_weights_), + GetBuffer<const _Float16>(fw_recurrent_to_forget_weights_), + GetBuffer<const _Float16>(fw_recurrent_to_cell_weights_), + GetBuffer<const _Float16>(fw_recurrent_to_output_weights_), + fw_recurrent_to_output_weights_->shape(), + GetBuffer<const _Float16>(fw_cell_to_input_weights_), + GetBuffer<const _Float16>(fw_cell_to_forget_weights_), + GetBuffer<const _Float16>(fw_cell_to_output_weights_), + GetOptionalBuffer<const _Float16>(aux_input_), aux_input_->shape(), + GetOptionalBuffer<const _Float16>(fw_aux_input_to_input_weights_), + GetOptionalBuffer<const _Float16>(fw_aux_input_to_forget_weights_), + GetOptionalBuffer<const _Float16>(fw_aux_input_to_cell_weights_), + GetOptionalBuffer<const _Float16>(fw_aux_input_to_output_weights_), + GetBuffer<const _Float16>(fw_input_gate_bias_), + GetBuffer<const _Float16>(fw_forget_gate_bias_), + GetBuffer<const _Float16>(fw_cell_bias_), + GetBuffer<const _Float16>(fw_output_gate_bias_), + GetBuffer<const _Float16>(fw_projection_weights_), + GetBuffer<const _Float16>(fw_projection_bias_), + GetBuffer<const _Float16>(fw_activation_state_), + GetBuffer<const _Float16>(fw_cell_state_), + /*input_layer_norm_weights=*/nullptr, + /*forget_layer_norm_weights=*/nullptr, + /*cell_layer_norm_weights=*/nullptr, + /*output_layer_norm_weights=*/nullptr, + GetBuffer<_Float16>(fw_activation_state_), GetBuffer<_Float16>(fw_cell_state_), + GetBuffer<_Float16>(fw_output_), fw_scratch_buffer.data(), kForwardSequence, + params_.time_major); + + std::vector<_Float16> bw_scratch_buffer(getNumberOfElements(bw_scratch_shape_)); + const bool kBackwardSequence = false; + LSTMCell::LSTMEvalFloat16( + params_, GetBuffer<const _Float16>(input_), input_->shape(), + GetBuffer<const _Float16>(bw_input_to_input_weights_), + GetBuffer<const _Float16>(bw_input_to_forget_weights_), + GetBuffer<const _Float16>(bw_input_to_cell_weights_), + GetBuffer<const _Float16>(bw_input_to_output_weights_), + bw_input_to_output_weights_->shape(), + GetBuffer<const _Float16>(bw_recurrent_to_input_weights_), + GetBuffer<const _Float16>(bw_recurrent_to_forget_weights_), + GetBuffer<const _Float16>(bw_recurrent_to_cell_weights_), + GetBuffer<const _Float16>(bw_recurrent_to_output_weights_), + bw_recurrent_to_output_weights_->shape(), + GetBuffer<const _Float16>(bw_cell_to_input_weights_), + GetBuffer<const _Float16>(bw_cell_to_forget_weights_), + GetBuffer<const _Float16>(bw_cell_to_output_weights_), + GetOptionalBuffer<const _Float16>(aux_input_), aux_input_->shape(), + GetOptionalBuffer<const _Float16>(bw_aux_input_to_input_weights_), + GetOptionalBuffer<const _Float16>(bw_aux_input_to_forget_weights_), + GetOptionalBuffer<const _Float16>(bw_aux_input_to_cell_weights_), + GetOptionalBuffer<const _Float16>(bw_aux_input_to_output_weights_), + GetBuffer<const _Float16>(bw_input_gate_bias_), + GetBuffer<const _Float16>(bw_forget_gate_bias_), + GetBuffer<const _Float16>(bw_cell_bias_), + GetBuffer<const _Float16>(bw_output_gate_bias_), + GetBuffer<const _Float16>(bw_projection_weights_), + GetBuffer<const _Float16>(bw_projection_bias_), + GetBuffer<const _Float16>(bw_activation_state_), + GetBuffer<const _Float16>(bw_cell_state_), + /*input_layer_norm_weights=*/nullptr, + /*forget_layer_norm_weights=*/nullptr, + /*cell_layer_norm_weights=*/nullptr, + /*output_layer_norm_weights=*/nullptr, + GetBuffer<_Float16>(bw_activation_state_), GetBuffer<_Float16>(bw_cell_state_), + params_.merge_outputs ? GetBuffer<_Float16>(fw_output_) + n_fw_output + : GetBuffer<_Float16>(bw_output_), + bw_scratch_buffer.data(), kBackwardSequence, params_.time_major); + } break; + default: { + LOG(ERROR) << "Unsupported data type: " << static_cast<int>(input_->type); + return false; + } + } + return true; +} + +} // namespace nn +} // namespace android
diff --git a/common/operations/BidirectionalSequenceLSTM.h b/common/operations/BidirectionalSequenceLSTM.h new file mode 100644 index 0000000..8eeb88b --- /dev/null +++ b/common/operations/BidirectionalSequenceLSTM.h
@@ -0,0 +1,213 @@ +/* + * Copyright (C) 2019 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FRAMEWORKS_ML_NN_BIDIRECTIONAL_SEQUENCE_LSTM_H +#define FRAMEWORKS_ML_NN_BIDIRECTIONAL_SEQUENCE_LSTM_H + +#include "ActivationFunctor.h" +#include "HalOperation.h" +#include "LSTM.h" +#include "OperationsUtils.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" + +#include <algorithm> +#include <cmath> + +namespace android { +namespace nn { + +struct RunTimeOperandInfo; + +class BidirectionalSequenceLSTM { + public: + BidirectionalSequenceLSTM(const Operation& operation, + std::vector<RunTimeOperandInfo>& operands); + + bool Prepare(const Operation& operation, std::vector<RunTimeOperandInfo>& operands, + Shape* fwOutputShape, Shape* bwOutputShape); + bool Eval(); + + // Input Tensors of size {max_time, n_batch, n_input} + static constexpr int kInputTensor = 0; + + // Forward LSTM cell tensors. + // Input weight tensors of size: {n_cell, n_input} + static constexpr int kFwInputToInputWeightsTensor = 1; // Optional + static constexpr int kFwInputToForgetWeightsTensor = 2; + static constexpr int kFwInputToCellWeightsTensor = 3; + static constexpr int kFwInputToOutputWeightsTensor = 4; + + // Recurrent weight tensors of size {n_cell, n_output} + static constexpr int kFwRecurrentToInputWeightsTensor = 5; // Optional + static constexpr int kFwRecurrentToForgetWeightsTensor = 6; + static constexpr int kFwRecurrentToCellWeightsTensor = 7; + static constexpr int kFwRecurrentToOutputWeightsTensor = 8; + + // Peephole weights tensors of size {n_cell}, representing a diagonal matrix. + static constexpr int kFwCellToInputWeightsTensor = 9; // Optional + static constexpr int kFwCellToForgetWeightsTensor = 10; // Optional + static constexpr int kFwCellToOutputWeightsTensor = 11; // Optional + + // Gates bias tensors of size {n_cell} + static constexpr int kFwInputGateBiasTensor = 12; // Optional + static constexpr int kFwForgetGateBiasTensor = 13; + static constexpr int kFwCellGateBiasTensor = 14; + static constexpr int kFwOutputGateBiasTensor = 15; + + // Projection weight tensor of size {n_output, n_cell} + static constexpr int kFwProjectionWeightsTensor = 16; // Optional + // Projection bias tensor of size {n_output} + static constexpr int kFwProjectionBiasTensor = 17; // Optional + + // Backward LSTM cell tensors. + // Input weight tensors of size: {n_cell, n_input} + static constexpr int kBwInputToInputWeightsTensor = 18; // Optional + static constexpr int kBwInputToForgetWeightsTensor = 19; + static constexpr int kBwInputToCellWeightsTensor = 20; + static constexpr int kBwInputToOutputWeightsTensor = 21; + + // Recurrent weight tensors of size {n_cell, n_output} + static constexpr int kBwRecurrentToInputWeightsTensor = 22; // Optional + static constexpr int kBwRecurrentToForgetWeightsTensor = 23; + static constexpr int kBwRecurrentToCellWeightsTensor = 24; + static constexpr int kBwRecurrentToOutputWeightsTensor = 25; + + // Peephole weights tensors of size {n_cell}, representing a diagonal matrix. + static constexpr int kBwCellToInputWeightsTensor = 26; // Optional + static constexpr int kBwCellToForgetWeightsTensor = 27; // Optional + static constexpr int kBwCellToOutputWeightsTensor = 28; // Optional + + // Gates bias tensors of size {n_cell} + static constexpr int kBwInputGateBiasTensor = 29; // Optional + static constexpr int kBwForgetGateBiasTensor = 30; + static constexpr int kBwCellGateBiasTensor = 31; + static constexpr int kBwOutputGateBiasTensor = 32; + + // Projection weight tensor of size {n_output, n_cell} + static constexpr int kBwProjectionWeightsTensor = 33; // Optional + // Projection bias tensor of size {n_output} + static constexpr int kBwProjectionBiasTensor = 34; // Optional + + // Stateful input tensors that are variables and will be modified by the Op. + // Activation state tensors of size {n_batch, n_output} + static constexpr int kFwInputActivationStateTensor = 35; + // Cell state tensors of size {n_batch, n_cell} + static constexpr int kFwInputCellStateTensor = 36; + // Activation state tensors of size {n_batch, n_output} + static constexpr int kBwInputActivationStateTensor = 37; + // Cell state tensors of size {n_batch, n_cell} + static constexpr int kBwInputCellStateTensor = 38; + + // Used as auxiliary input and weights when stacking for + // tf.contrib.rnn.stack_bidirectional_rnn case (with cross links); Used as input + // to the backward cell when stacking for tf.nn.static_bidirectional_rnn case + // (without cross links). + static constexpr int kAuxInputTensor = 39; // Optional + // Forward weights. + static constexpr int kFwAuxInputToInputWeightsTensor = 40; // Optional + static constexpr int kFwAuxInputToForgetWeightsTensor = 41; // Optional + static constexpr int kFwAuxInputToCellWeightsTensor = 42; // Optional + static constexpr int kFwAuxInputToOutputWeightsTensor = 43; // Optional + // Backward weights. + static constexpr int kBwAuxInputToInputWeightsTensor = 44; // Optional + static constexpr int kBwAuxInputToForgetWeightsTensor = 45; // Optional + static constexpr int kBwAuxInputToCellWeightsTensor = 46; // Optional + static constexpr int kBwAuxInputToOutputWeightsTensor = 47; // Optional + + static constexpr int kActivationParam = 48; + static constexpr int kCellClipParam = 49; + static constexpr int kProjClipParam = 50; + static constexpr int kMergeOutputsParam = 51; + static constexpr int kTimeMajorParam = 52; + + // Output tensors. + static constexpr int kFwOutputTensor = 0; + static constexpr int kBwOutputTensor = 1; // Ignored if merge_outputs is set. + + private: + LSTMParams params_; + Shape fw_scratch_shape_; + Shape bw_scratch_shape_; + + const RunTimeOperandInfo* input_; + + const RunTimeOperandInfo* aux_input_; + const RunTimeOperandInfo* fw_aux_input_to_input_weights_; + const RunTimeOperandInfo* fw_aux_input_to_forget_weights_; + const RunTimeOperandInfo* fw_aux_input_to_cell_weights_; + const RunTimeOperandInfo* fw_aux_input_to_output_weights_; + const RunTimeOperandInfo* bw_aux_input_to_input_weights_; + const RunTimeOperandInfo* bw_aux_input_to_forget_weights_; + const RunTimeOperandInfo* bw_aux_input_to_cell_weights_; + const RunTimeOperandInfo* bw_aux_input_to_output_weights_; + + const RunTimeOperandInfo* fw_input_to_input_weights_; + const RunTimeOperandInfo* fw_input_to_forget_weights_; + const RunTimeOperandInfo* fw_input_to_cell_weights_; + const RunTimeOperandInfo* fw_input_to_output_weights_; + + const RunTimeOperandInfo* fw_recurrent_to_input_weights_; + const RunTimeOperandInfo* fw_recurrent_to_forget_weights_; + const RunTimeOperandInfo* fw_recurrent_to_cell_weights_; + const RunTimeOperandInfo* fw_recurrent_to_output_weights_; + + const RunTimeOperandInfo* fw_cell_to_input_weights_; + const RunTimeOperandInfo* fw_cell_to_forget_weights_; + const RunTimeOperandInfo* fw_cell_to_output_weights_; + + const RunTimeOperandInfo* fw_input_gate_bias_; + const RunTimeOperandInfo* fw_forget_gate_bias_; + const RunTimeOperandInfo* fw_cell_bias_; + const RunTimeOperandInfo* fw_output_gate_bias_; + + const RunTimeOperandInfo* fw_projection_weights_; + const RunTimeOperandInfo* fw_projection_bias_; + + RunTimeOperandInfo* fw_activation_state_; + RunTimeOperandInfo* fw_cell_state_; + RunTimeOperandInfo* fw_output_; + + const RunTimeOperandInfo* bw_input_to_input_weights_; + const RunTimeOperandInfo* bw_input_to_forget_weights_; + const RunTimeOperandInfo* bw_input_to_cell_weights_; + const RunTimeOperandInfo* bw_input_to_output_weights_; + + const RunTimeOperandInfo* bw_recurrent_to_input_weights_; + const RunTimeOperandInfo* bw_recurrent_to_forget_weights_; + const RunTimeOperandInfo* bw_recurrent_to_cell_weights_; + const RunTimeOperandInfo* bw_recurrent_to_output_weights_; + + const RunTimeOperandInfo* bw_cell_to_input_weights_; + const RunTimeOperandInfo* bw_cell_to_forget_weights_; + const RunTimeOperandInfo* bw_cell_to_output_weights_; + + const RunTimeOperandInfo* bw_input_gate_bias_; + const RunTimeOperandInfo* bw_forget_gate_bias_; + const RunTimeOperandInfo* bw_cell_bias_; + const RunTimeOperandInfo* bw_output_gate_bias_; + + const RunTimeOperandInfo* bw_projection_weights_; + const RunTimeOperandInfo* bw_projection_bias_; + + RunTimeOperandInfo* bw_activation_state_; + RunTimeOperandInfo* bw_cell_state_; + RunTimeOperandInfo* bw_output_; +}; + +} // namespace nn +} // namespace android + +#endif // FRAMEWORKS_ML_NN_BIDIRECTIONAL_SEQUENCE_LSTM_H
diff --git a/common/operations/LSTM.cpp b/common/operations/LSTM.cpp index 293b23e..1b36574 100644 --- a/common/operations/LSTM.cpp +++ b/common/operations/LSTM.cpp
@@ -350,14 +350,19 @@ 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, - const float* projection_weights_buffer, const float* projection_bias_buffer, - const float* output_state_in_buffer, const float* cell_state_in_buffer, - const float* input_layer_norm_weights_buffer, const float* forget_layer_norm_weights_buffer, - 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, bool timeMajor) { + const float* aux_input_buffer, const Shape& aux_input_shape, + const float* aux_input_to_input_weights_buffer, + const float* aux_input_to_forget_weights_buffer, + const float* aux_input_to_cell_weights_buffer, + const float* aux_input_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, const float* projection_weights_buffer, + const float* projection_bias_buffer, const float* output_state_in_buffer, + const float* cell_state_in_buffer, const float* input_layer_norm_weights_buffer, + const float* forget_layer_norm_weights_buffer, 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, + bool forwardSequence, bool timeMajor) { NNTRACE_COMP("LSTMCell::LSTMEvalFloat32"); const uint32_t inputRank = getNumberOfDimensions(input_shape); @@ -388,13 +393,20 @@ transposedOutputShape = transposedInputShape; transposedOutputShape.dimensions[2] = outputSize; } - const float* inputCurrentTimeStep = timeMajor ? input_buffer : transposedInput.data(); - float* outputCurrentTimeStep = timeMajor ? output_buffer : transposedOutput.data(); + const float* inputData = timeMajor ? input_buffer : transposedInput.data(); + float* outputData = timeMajor ? output_buffer : transposedOutput.data(); std::vector<float> outputStateInCurrentTimeStep( output_state_in_buffer, output_state_in_buffer + batchSize * outputSize); std::vector<float> cellStateInCurrentTimeStep(cell_state_in_buffer, cell_state_in_buffer + batchSize * numCells); + const float* inputCurrentTimeStep = + inputData + (forwardSequence ? 0 : batchInputSize * (maxTime - 1)); + float* outputCurrentTimeStep = + outputData + (forwardSequence ? 0 : batchOutputSize * (maxTime - 1)); + const int batchInputDelta = forwardSequence ? batchInputSize : -batchInputSize; + const int batchOutputDelta = forwardSequence ? batchOutputSize : -batchOutputSize; + for (int t = 0; t < maxTime; ++t) { LSTMStep(params, inputCurrentTimeStep, batchInputShape, input_to_input_weights_buffer, input_to_forget_weights_buffer, input_to_cell_weights_buffer, @@ -402,16 +414,18 @@ recurrent_to_input_weights_buffer, recurrent_to_forget_weights_buffer, recurrent_to_cell_weights_buffer, recurrent_to_output_weights_buffer, 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_buffer, cell_bias_buffer, - output_gate_bias_buffer, projection_weights_buffer, projection_bias_buffer, - outputStateInCurrentTimeStep.data(), cellStateInCurrentTimeStep.data(), - input_layer_norm_weights_buffer, forget_layer_norm_weights_buffer, - cell_layer_norm_weights_buffer, output_layer_norm_weights_buffer, - output_state_out_buffer, cell_state_out_buffer, outputCurrentTimeStep, - scratch_buffer_buffer); - inputCurrentTimeStep += batchInputSize; - outputCurrentTimeStep += batchOutputSize; + cell_to_forget_weights_buffer, cell_to_output_weights_buffer, aux_input_buffer, + aux_input_shape, aux_input_to_input_weights_buffer, + aux_input_to_forget_weights_buffer, aux_input_to_cell_weights_buffer, + aux_input_to_output_weights_buffer, input_gate_bias_buffer, + forget_gate_bias_buffer, cell_bias_buffer, output_gate_bias_buffer, + projection_weights_buffer, projection_bias_buffer, output_state_in_buffer, + cell_state_in_buffer, input_layer_norm_weights_buffer, + forget_layer_norm_weights_buffer, cell_layer_norm_weights_buffer, + output_layer_norm_weights_buffer, output_state_out_buffer, cell_state_out_buffer, + outputCurrentTimeStep, scratch_buffer_buffer); + inputCurrentTimeStep += batchInputDelta; + outputCurrentTimeStep += batchOutputDelta; outputStateInCurrentTimeStep.assign(output_state_out_buffer, output_state_out_buffer + batchSize * outputSize); cellStateInCurrentTimeStep.assign(cell_state_out_buffer, @@ -439,7 +453,11 @@ const _Float16* recurrent_to_output_weights_buffer, const Shape& recurrent_to_output_weights_shape, const _Float16* cell_to_input_weights_buffer, const _Float16* cell_to_forget_weights_buffer, - const _Float16* cell_to_output_weights_buffer, const _Float16* input_gate_bias_buffer, + const _Float16* cell_to_output_weights_buffer, const _Float16* aux_input_buffer, + const Shape& aux_input_shape, const _Float16* aux_input_to_input_weights_buffer, + const _Float16* aux_input_to_forget_weights_buffer, + const _Float16* aux_input_to_cell_weights_buffer, + const _Float16* aux_input_to_output_weights_buffer, const _Float16* input_gate_bias_buffer, const _Float16* forget_gate_bias_buffer, const _Float16* cell_bias_buffer, const _Float16* output_gate_bias_buffer, const _Float16* projection_weights_buffer, const _Float16* projection_bias_buffer, const _Float16* output_state_in_buffer, @@ -448,7 +466,7 @@ const _Float16* cell_layer_norm_weights_buffer, const _Float16* output_layer_norm_weights_buffer, _Float16* output_state_out_buffer, _Float16* cell_state_out_buffer, _Float16* output_buffer, _Float16* scratch_buffer_buffer, - bool timeMajor) { + bool forwardSequence, bool timeMajor) { NNTRACE_COMP("LSTMCell::LSTMEvalFloat16"); const uint32_t inputRank = getNumberOfDimensions(input_shape); @@ -507,6 +525,22 @@ convertFloat16ToFloat32(cell_to_output_weights_buffer, &cell_to_output_weights_float32); } + std::vector<float> aux_input_float32(maxTime * batchInputSize); + convertFloat16ToFloat32(aux_input_buffer, &aux_input_float32); + std::vector<float> aux_input_to_input_weights_float32(numCells * inputSize); + if (aux_input_to_input_weights_buffer != nullptr) { + convertFloat16ToFloat32(aux_input_to_input_weights_buffer, + &aux_input_to_input_weights_float32); + } + std::vector<float> aux_input_to_forget_weights_float32(numCells * inputSize); + convertFloat16ToFloat32(aux_input_to_forget_weights_buffer, + &aux_input_to_forget_weights_float32); + std::vector<float> aux_input_to_cell_weights_float32(numCells * inputSize); + convertFloat16ToFloat32(aux_input_to_cell_weights_buffer, &aux_input_to_cell_weights_float32); + std::vector<float> aux_input_to_output_weights_float32(numCells * inputSize); + convertFloat16ToFloat32(aux_input_to_output_weights_buffer, + &aux_input_to_output_weights_float32); + std::vector<float> input_gate_bias_float32(numCells); if (input_gate_bias_buffer != nullptr) { convertFloat16ToFloat32(input_gate_bias_buffer, &input_gate_bias_float32); @@ -570,13 +604,21 @@ transposedOutputShape = transposedInputShape; transposedOutputShape.dimensions[2] = outputSize; } - const float* inputCurrentTimeStep = timeMajor ? input_float32.data() : transposedInput.data(); - float* outputCurrentTimeStep = timeMajor ? output_float32.data() : transposedOutput.data(); + const float* inputData = timeMajor ? input_float32.data() : transposedInput.data(); + float* outputData = timeMajor ? output_float32.data() : transposedOutput.data(); std::vector<float> outputStateInCurrentTimeStep(batchSize * outputSize); convertFloat16ToFloat32(output_state_in_buffer, &outputStateInCurrentTimeStep); std::vector<float> cellStateInCurrentTimeStep(batchSize * numCells); convertFloat16ToFloat32(cell_state_in_buffer, &cellStateInCurrentTimeStep); + + const float* inputCurrentTimeStep = + inputData + (forwardSequence ? 0 : batchInputSize * (maxTime - 1)); + float* outputCurrentTimeStep = + outputData + (forwardSequence ? 0 : batchOutputSize * (maxTime - 1)); + const int batchInputDelta = forwardSequence ? batchInputSize : -batchInputSize; + const int batchOutputDelta = forwardSequence ? batchOutputSize : -batchOutputSize; + for (int t = 0; t < maxTime; ++t) { LSTMStep(params, inputCurrentTimeStep, batchInputShape, input_to_input_weights_float32.data(), input_to_forget_weights_float32.data(), @@ -586,7 +628,11 @@ recurrent_to_cell_weights_float32.data(), recurrent_to_output_weights_float32.data(), recurrent_to_output_weights_shape, cell_to_input_weights_float32.data(), cell_to_forget_weights_float32.data(), - cell_to_output_weights_float32.data(), input_gate_bias_float32.data(), + cell_to_output_weights_float32.data(), aux_input_float32.data(), aux_input_shape, + aux_input_to_input_weights_float32.data(), + aux_input_to_forget_weights_float32.data(), + aux_input_to_cell_weights_float32.data(), + aux_input_to_output_weights_float32.data(), input_gate_bias_float32.data(), forget_gate_bias_float32.data(), cell_bias_float32.data(), output_gate_bias_float32.data(), projection_weights_float32.data(), projection_bias_float32.data(), outputStateInCurrentTimeStep.data(), @@ -595,8 +641,8 @@ output_layer_norm_weights_float32.data(), output_state_out_float32.data(), cell_state_out_float32.data(), outputCurrentTimeStep, scratch_buffer_float32.data()); - inputCurrentTimeStep += batchInputSize; - outputCurrentTimeStep += batchOutputSize; + inputCurrentTimeStep += batchInputDelta; + outputCurrentTimeStep += batchOutputDelta; outputStateInCurrentTimeStep = output_state_out_float32; cellStateInCurrentTimeStep = cell_state_out_float32; } @@ -624,14 +670,18 @@ 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, - const float* projection_weights_buffer, const float* projection_bias_buffer, - const float* output_state_in_buffer, const float* cell_state_in_buffer, - const float* input_layer_norm_weights_buffer, const float* forget_layer_norm_weights_buffer, - 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) { + const float* aux_input_buffer, const Shape& aux_input_shape, + const float* aux_input_to_input_weights_buffer, + const float* aux_input_to_forget_weights_buffer, + const float* aux_input_to_cell_weights_buffer, + const float* aux_input_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, const float* projection_weights_buffer, + const float* projection_bias_buffer, const float* output_state_in_buffer, + const float* cell_state_in_buffer, const float* input_layer_norm_weights_buffer, + const float* forget_layer_norm_weights_buffer, 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::LSTMStep"); const uint32_t n_batch = input_shape.dimensions[0]; @@ -639,6 +689,7 @@ // 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]; + const uint32_t n_aux_input = aux_input_buffer == nullptr ? 0 : n_input; // Index the scratch buffers pointers to the global scratch buffer. float* input_gate_scratch = nullptr; @@ -694,6 +745,26 @@ input_to_output_weights_buffer, n_cell, n_input, input_buffer, n_batch, output_gate_scratch, /*result_stride*/ 1); + // If auxiliary input is available then compute aux_input_weight * aux_input + if (aux_input_buffer != nullptr) { + if (!params.use_cifg) { + tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( + aux_input_to_input_weights_buffer, n_cell, n_aux_input, aux_input_buffer, + n_batch, input_gate_scratch, + /*result_stride=*/1); + } + + tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( + aux_input_to_forget_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, + forget_gate_scratch, /*result_stride=*/1); + tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( + aux_input_to_cell_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, + cell_scratch, /*result_stride=*/1); + tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( + aux_input_to_output_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, + output_gate_scratch, /*result_stride=*/1); + } + // For each batch and cell: compute recurrent_weight * output_state. if (!params.use_cifg) { tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate( @@ -837,6 +908,11 @@ GetBuffer<const float>(cell_to_input_weights_), GetBuffer<const float>(cell_to_forget_weights_), GetBuffer<const float>(cell_to_output_weights_), + /*aux_input_buffer=*/nullptr, input_->shape(), + /*aux_input_to_input_weights_buffer=*/nullptr, + /*aux_input_to_forget_weights_buffer=*/nullptr, + /*aux_input_to_cell_weights_buffer=*/nullptr, + /*aux_input_to_output_weights_buffer=*/nullptr, GetBuffer<const float>(input_gate_bias_), GetBuffer<const float>(forget_gate_bias_), GetBuffer<const float>(cell_bias_), @@ -867,6 +943,11 @@ GetOptionalBuffer<const _Float16>(cell_to_input_weights_), GetOptionalBuffer<const _Float16>(cell_to_forget_weights_), GetOptionalBuffer<const _Float16>(cell_to_output_weights_), + /*aux_input_buffer=*/nullptr, input_->shape(), + /*aux_input_to_input_weights_buffer=*/nullptr, + /*aux_input_to_forget_weights_buffer=*/nullptr, + /*aux_input_to_cell_weights_buffer=*/nullptr, + /*aux_input_to_output_weights_buffer=*/nullptr, GetOptionalBuffer<const _Float16>(input_gate_bias_), GetBuffer<const _Float16>(forget_gate_bias_), GetBuffer<const _Float16>(cell_bias_),
diff --git a/common/operations/LSTM.h b/common/operations/LSTM.h index 8caae90..6345a1d 100644 --- a/common/operations/LSTM.h +++ b/common/operations/LSTM.h
@@ -36,6 +36,8 @@ bool use_layer_norm; bool use_projection_weight; bool use_projection_bias; + bool merge_outputs; + bool time_major; }; struct RunTimeOperandInfo; @@ -113,7 +115,10 @@ 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* cell_to_output_weights_buffer, const float* aux_input_buffer, + const Shape& aux_input_shape, const float* aux_input_to_input_weights, + const float* aux_input_to_forget_weights, const float* aux_input_to_cell_weights, + const float* aux_input_to_output_weights, const float* input_gate_bias_buffer, const float* forget_gate_bias_buffer, const float* cell_bias_buffer, const float* output_gate_bias_buffer, const float* projection_weights_buffer, const float* projection_bias_buffer, const float* output_state_in_buffer, @@ -122,7 +127,7 @@ 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, - bool timeMajor = true); + bool forwardSequence = true, bool timeMajor = true); static bool LSTMEvalFloat16( const LSTMParams& params, const _Float16* input_buffer, const Shape& input_shape, @@ -138,7 +143,10 @@ const Shape& recurrent_to_output_weights_shape, const _Float16* cell_to_input_weights_buffer, const _Float16* cell_to_forget_weights_buffer, - const _Float16* cell_to_output_weights_buffer, const _Float16* input_gate_bias_buffer, + const _Float16* cell_to_output_weights_buffer, const _Float16* aux_input_buffer, + const Shape& aux_input_shape, const _Float16* aux_input_to_input_weights, + const _Float16* aux_input_to_forget_weights, const _Float16* aux_input_to_cell_weights, + const _Float16* aux_input_to_output_weights, const _Float16* input_gate_bias_buffer, const _Float16* forget_gate_bias_buffer, const _Float16* cell_bias_buffer, const _Float16* output_gate_bias_buffer, const _Float16* projection_weights_buffer, const _Float16* projection_bias_buffer, const _Float16* output_state_in_buffer, @@ -147,7 +155,7 @@ const _Float16* cell_layer_norm_weights_buffer, const _Float16* output_layer_norm_weights_buffer, _Float16* output_state_out_buffer, _Float16* cell_state_out_buffer, _Float16* output_buffer, - _Float16* scratch_buffer_buffer, bool timeMajor = true); + _Float16* scratch_buffer_buffer, bool forwardSequence = true, bool timeMajor = true); static bool LSTMStep( const LSTMParams& params, const float* input_buffer, const Shape& input_shape, @@ -160,7 +168,10 @@ 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* cell_to_output_weights_buffer, const float* aux_input_buffer, + const Shape& aux_input_shape, const float* aux_input_to_input_weights, + const float* aux_input_to_forget_weights, const float* aux_input_to_cell_weights, + const float* aux_input_to_output_weights, const float* input_gate_bias_buffer, const float* forget_gate_bias_buffer, const float* cell_bias_buffer, const float* output_gate_bias_buffer, const float* projection_weights_buffer, const float* projection_bias_buffer, const float* output_state_in_buffer,
diff --git a/common/operations/UnidirectionalSequenceLSTM.cpp b/common/operations/UnidirectionalSequenceLSTM.cpp index 6ba5bc5..d2edce3 100644 --- a/common/operations/UnidirectionalSequenceLSTM.cpp +++ b/common/operations/UnidirectionalSequenceLSTM.cpp
@@ -325,6 +325,11 @@ context->getInputBuffer<float>(kCellToInputWeightsTensor), context->getInputBuffer<float>(kCellToForgetWeightsTensor), context->getInputBuffer<float>(kCellToOutputWeightsTensor), + /*aux_input_buffer=*/nullptr, context->getInputShape(kInputTensor), + /*aux_input_to_input_weights_buffer=*/nullptr, + /*aux_input_to_forget_weights_buffer=*/nullptr, + /*aux_input_to_cell_weights_buffer=*/nullptr, + /*aux_input_to_output_weights_buffer=*/nullptr, context->getInputBuffer<float>(kInputGateBiasTensor), context->getInputBuffer<float>(kForgetGateBiasTensor), context->getInputBuffer<float>(kCellGateBiasTensor), @@ -360,6 +365,11 @@ context->getInputBuffer<_Float16>(kCellToInputWeightsTensor), context->getInputBuffer<_Float16>(kCellToForgetWeightsTensor), context->getInputBuffer<_Float16>(kCellToOutputWeightsTensor), + /*aux_input_buffer=*/nullptr, context->getInputShape(kInputTensor), + /*aux_input_to_input_weights_buffer=*/nullptr, + /*aux_input_to_forget_weights_buffer=*/nullptr, + /*aux_input_to_cell_weights_buffer=*/nullptr, + /*aux_input_to_output_weights_buffer=*/nullptr, context->getInputBuffer<_Float16>(kInputGateBiasTensor), context->getInputBuffer<_Float16>(kForgetGateBiasTensor), context->getInputBuffer<_Float16>(kCellGateBiasTensor),