diff --git a/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc b/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc index ea0e2dd9306..b89130b8cce 100644 --- a/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc +++ b/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc @@ -49,7 +49,7 @@ LstmTensors::~LstmTensors() { } // Verify the LSTM internal tensor properties (e.g., type checks) -// Input/output/states/fc weights tensors are required for kernel evaluation. +// Input/output/states/fc weights tensors are required for kernel evaulation. // The state tensors should be variables. Variants of the standard LSTM // are not supported here, therefore their corresponding tensors should be // invalid @@ -105,14 +105,14 @@ TfLiteStatus LstmTensors::ValidateTensorStatus(TfLiteContext* context) const { namespace lstm_internal { -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) const int32_t kInt16Max = std::numeric_limits::max(); const int32_t kInt16Min = std::numeric_limits::min(); #endif void AddElementWise(const int16_t* input_1, const int16_t* input_2, int n_batch, int n_input, int16_t* output) { -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) for (int batch = 0; batch < n_batch; ++batch) { for (int i = 0; i < n_input; ++i) { const int index = batch * n_input + i; @@ -122,21 +122,29 @@ void AddElementWise(const int16_t* input_1, const int16_t* input_2, int n_batch, } } #else - xa_nn_elm_add_16x16_16(output, input_1, input_2, n_batch * n_input); + WORD32 err; + err = xa_nn_elm_add_16x16_16(output, input_1, input_2, n_batch * n_input); + (void)err; #endif } void AddElementWise(const float* input_1, const float* input_2, int n_batch, int n_input, float* output) { +#if defined(INCLUDE_FLOAT_OPT) && (defined(HIFI5) || defined(HIFI4)) + WORD32 err; + err = xa_nn_elm_add_f32xf32_f32(output, input_1, input_2, n_batch * n_input); + (void)err; +#else for (int batch = 0; batch < n_batch; ++batch) { for (int i = 0; i < n_input; ++i) { const int index = batch * n_input + i; output[index] = input_1[index] + input_2[index]; } } +#endif } -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void Sigmoid(const RuntimeShape& data_shape, int16_t* data) { reference_integer_ops::Logistic( 0 /*data->input_multiplier*/, 0 /*data->input_left_shift */, @@ -224,15 +232,23 @@ void FullyConnected(const FullyConnectedParams& params, params, input_shape, input_data, filter_shape, filter_data, bias_shape, bias_data, output_shape, output_data); } -#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void Sigmoid(int16_t* data, int32_t data_size) { - xa_nn_vec_sigmoid_sym16s_sym16s(data, data, 0, 0, data_size); + WORD32 err; + err = xa_nn_vec_sigmoid_sym16s_sym16s(data, data, 0, 0, data_size); + (void)err; } void Sigmoid(float* data, int32_t data_size) { +#if defined(INCLUDE_FLOAT_OPT) && !(defined(HIFI_IQ)) + WORD32 err; + err = xa_nn_vec_sigmoid_f32_f32(data, data, data_size); + (void)err; +#else int data_dims[2] = {1, data_size}; RuntimeShape data_shape(2, reinterpret_cast(data_dims)); reference_ops::Logistic(data_shape, data, data_shape, data); +#endif } void Tanh(int32_t cell_state_scale_power, int16_t* input_data, @@ -249,8 +265,10 @@ void Tanh(int32_t cell_state_scale_power, int16_t* input_data, input_multiplier = 3; #endif } - xa_nn_vec_tanh_sym16s_sym16s(output_data, input_data, input_multiplier, + WORD32 err; + err = xa_nn_vec_tanh_sym16s_sym16s(output_data, input_data, input_multiplier, tanh_input_left_shift, data_size); + (void)err; } void Tanh(int32_t cell_state_scale_power, float* input_data, float* output_data, @@ -263,20 +281,24 @@ void Tanh(int32_t cell_state_scale_power, float* input_data, float* output_data, // Input and output have the same shape in LSTM void Mul(const ArithmeticParams& params, const int16_t* input1_data, const int16_t* input2_data, int8_t* output_data, int32_t data_size) { - xa_nn_elm_mul_sym16sxsym16s_asym8s( + WORD32 err; + err = xa_nn_elm_mul_sym16sxsym16s_asym8s( output_data, params.output_offset, params.output_shift, params.output_multiplier, params.quantized_activation_min, params.quantized_activation_max, input1_data, input2_data, data_size); + (void)err; } // Input and output have the same shape in LSTM void Mul(const ArithmeticParams& params, const int16_t* input1_data, const int16_t* input2_data, int16_t* output_data, int32_t data_size) { int dims_4D[4] = {1, 1, 1, data_size}; - xa_nn_elm_mul_broadcast_4D_sym16sxsym16s_sym16s( + WORD32 err; + err = xa_nn_elm_mul_broadcast_4D_sym16sxsym16s_sym16s( output_data, dims_4D, params.output_shift, params.output_multiplier, params.quantized_activation_min, params.quantized_activation_max, input1_data, dims_4D, input2_data, dims_4D); + (void)err; return; } @@ -294,42 +316,54 @@ void FullyConnected(const FullyConnectedParams& params, const int32_t* bias_data, int16_t* output_data, const int num_batches, const int output_depth, const int accum_depth) { -#pragma loop_count min = 1 - for (int b = 0; b < num_batches; b++) { - xa_nn_matXvec_out_stride_sym8sxasym8s_16( - output_data + b * output_depth, filter_data, - input_data + b * accum_depth, bias_data, output_depth, accum_depth, - accum_depth, 1, params.input_offset, params.output_multiplier, + WORD32 err; + if(num_batches == 1) { + err = xa_nn_matXvec_out_stride_sym8sxasym8s_16( + output_data, filter_data, input_data, bias_data, + output_depth, accum_depth, accum_depth, 1, + params.input_offset, params.output_multiplier, params.output_shift); } + else{ + err = xa_nn_matmul_sym8sxasym8s_sym16s( + output_data, filter_data, input_data, bias_data, + output_depth, accum_depth, accum_depth, num_batches, + accum_depth, output_depth, 1, params.input_offset, + params.output_multiplier, params.output_shift); + } + (void)err; return; } -#define ARG_CHK_ALIGN(_ptr, _align) \ - (((unsigned int)(_ptr) & ((_align) - 1)) == 0) - void FullyConnected(const FullyConnectedParams& params, const int16_t* input_data, const int8_t* filter_data, const int64_t* bias_data, int16_t* output_data, const int num_batches, const int output_depth, const int accum_depth) { WORD32 err; - if (num_batches == 1 && ARG_CHK_ALIGN(output_data, sizeof(WORD16) * 8) && - ARG_CHK_ALIGN(filter_data, sizeof(WORD8) * 16) && - ARG_CHK_ALIGN(input_data, sizeof(WORD16) * 8) && - ARG_CHK_ALIGN(bias_data, sizeof(WORD64) * 2)) { - err = xa_nn_matXvec_v2_sym8sxsym16s_sym16s( - output_data, filter_data, input_data, bias_data, output_depth, - accum_depth, accum_depth, params.output_multiplier, params.output_shift, - -32768, 32767, NULL); - } else { - err = xa_nn_matmul_sym8sxsym16s_sym16s( - output_data, filter_data, input_data, bias_data, output_depth, - accum_depth, accum_depth, num_batches, accum_depth, output_depth, 1, - params.input_offset, params.output_multiplier, params.output_shift, - params.output_offset); + // The optimized single-batch fully_connected_v2 kernel requires all of its + // pointer arguments to be 16-byte aligned (see xa_nn_matXvec_v2_sym8sxsym16s + // alignment checks). Fall back to the matmul kernel, which tolerates lower + // alignment, whenever any pointer does not meet this requirement. + const bool inputs_16byte_aligned = + ((reinterpret_cast(output_data) & 15) == 0) && + ((reinterpret_cast(filter_data) & 15) == 0) && + ((reinterpret_cast(input_data) & 15) == 0) && + ((reinterpret_cast(bias_data) & 15) == 0); + if(num_batches == 1 && inputs_16byte_aligned) { + err = xa_nn_fully_connected_v2_sym8sxsym16s_sym16s( + output_data, filter_data, input_data, bias_data, + accum_depth, output_depth, params.output_multiplier, params.output_shift, + -32768, 32767, NULL); } - (void)err; + else{ + err = xa_nn_matmul_sym8sxsym16s_sym16s( + output_data, filter_data, input_data, bias_data, output_depth, + accum_depth, accum_depth, num_batches, accum_depth, output_depth, 1, + params.input_offset, params.output_multiplier, params.output_shift, + params.output_offset); + } + (void)err; return; } @@ -348,26 +382,32 @@ void FullyConnected(const FullyConnectedParams& params, const float* input_data, params, input_shape, input_data, filter_shape, filter_data, bias_shape, bias_data, output_shape, output_data); } -#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void Clipping(const int v_size, const CellStateInfo& cell_state_info, int16_t* vector) { - for (int i = 0; i < v_size; i++) { - vector[i] = - std::max(std::min(cell_state_info.quantized_cell_clip, vector[i]), - static_cast(-cell_state_info.quantized_cell_clip)); - } + WORD32 err; + err = xa_nn_vec_activation_min_max_16_16(vector, vector, -cell_state_info.quantized_cell_clip, + cell_state_info.quantized_cell_clip , v_size); + (void)err; } void Clipping(const int v_size, const CellStateInfo& cell_state_info, float* vector) { +#if defined(INCLUDE_FLOAT_OPT) + WORD32 err; + err = xa_nn_vec_activation_min_max_f32_f32(vector, vector, -cell_state_info.cell_clip, + cell_state_info.cell_clip , v_size); + (void)err; +#else for (int i = 0; i < v_size; i++) { vector[i] = std::max(std::min(cell_state_info.cell_clip, vector[i]), -cell_state_info.cell_clip); } +#endif } -#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) void UpdateLstmCell(const LstmStepManager& step_info, TfLiteEvalTensor* cell_state, // Gate outputs @@ -383,14 +423,16 @@ void UpdateLstmCell(const LstmStepManager& step_info, TFLITE_DCHECK_LE(step_info.CellStateOffset() + cell_state_shape.FlatSize(), tflite::micro::GetTensorShape(cell_state).FlatSize()); + WORD32 err; // Multiplier is equivalent to 0.5 here so adding 1 to shifts - xa_nn_lstm_cell_state_update_16( + err = xa_nn_lstm_cell_state_update_16( tflite::micro::GetTensorData(cell_state) + step_info.CellStateOffset(), forget_gate_output, cell_gate_output, input_gate_output, forget_cell_mul_params.output_shift - 1, input_mul_params.output_shift - 1, cell_state_info.quantized_cell_clip, cell_state_shape.FlatSize()); + (void)err; } void UpdateLstmCell(const LstmStepManager& step_info, @@ -434,7 +476,7 @@ void UpdateLstmCell(const LstmStepManager& step_info, step_info.CellStateOffset()); } } -#endif // #if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) +#endif // #if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) // Increment the data offset so the sigle time step invocation call can access // the corresponding input/output tensor data at the time step @@ -473,8 +515,7 @@ void LstmStepManager::UpdateBatch() { // Multi-batch for time_major input RuntimeShape LstmStepManager::InputShape() const { int batch_size = 1; - if (size_info_.time_major || - (size_info_.batch_size > 1 && size_info_.time_steps == 1)) { + if (size_info_.time_major) { batch_size = size_info_.batch_size; } const int dims[2] = {batch_size, size_info_.input_dimension}; @@ -486,8 +527,7 @@ RuntimeShape LstmStepManager::InputShape() const { // Multi-batch for time_major input RuntimeShape LstmStepManager::StateShape() const { int batch_size = 1; - if (size_info_.time_major || - (size_info_.batch_size > 1 && size_info_.time_steps == 1)) { + if (size_info_.time_major) { batch_size = size_info_.batch_size; } const int dims[2] = {batch_size, size_info_.state_dimension}; @@ -496,4 +536,96 @@ RuntimeShape LstmStepManager::StateShape() const { } } // namespace lstm_internal + +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) +template <> +TfLiteStatus EvalLstm( + const OpDataLSTM& op_data, LSTMKernelContents& kernel_content, + const LSTMBuffers& buffers) { + const LstmSizeInfo& size_info = op_data.size_info; + void* p_scratch = buffers.buffer0; + + // Gate weights: input (W) and recurrent (U) for i/f/c/o gates. + lstm_weights_ptrs weights; + weights.p_ig_W = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmInputToInputWeightsTensor)); + weights.p_fg_W = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmInputToForgetWeightsTensor)); + weights.p_cg_W = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmInputToCellWeightsTensor)); + weights.p_og_W = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmInputToOutputWeightsTensor)); + weights.p_ig_U = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmRecurrentToInputWeightsTensor)); + weights.p_fg_U = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmRecurrentToForgetWeightsTensor)); + weights.p_cg_U = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmRecurrentToCellWeightsTensor)); + weights.p_og_U = (void*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmRecurrentToOutputWeightsTensor)); + + lstm_bias_ptrs biases = {}; + biases.p_ig_W_bias = (void*)tflite::micro::GetOptionalTensorData( + kernel_content.GetInternalTensor(kLstmInputGateBiasTensor)); + biases.p_fg_W_bias = (void*)tflite::micro::GetOptionalTensorData( + kernel_content.GetInternalTensor(kLstmForgetGateBiasTensor)); + biases.p_cg_W_bias = (void*)tflite::micro::GetOptionalTensorData( + kernel_content.GetInternalTensor(kLstmCellGateBiasTensor)); + biases.p_og_W_bias = (void*)tflite::micro::GetOptionalTensorData( + kernel_content.GetInternalTensor(kLstmOutputGateBiasTensor)); + + const GateParameters& ig = op_data.input_gate_parameters; + const GateParameters& fg = op_data.forget_gate_parameters; + const GateParameters& cg = op_data.cell_gate_parameters; + const GateParameters& og = op_data.output_gate_parameters; + + lstm_quant_params qp; + qp.ig_W_out_multiplier = ig.input_fc_params.output_multiplier; + qp.fg_W_out_multiplier = fg.input_fc_params.output_multiplier; + qp.cg_W_out_multiplier = cg.input_fc_params.output_multiplier; + qp.og_W_out_multiplier = og.input_fc_params.output_multiplier; + qp.ig_U_out_multiplier = ig.recurrent_fc_params.output_multiplier; + qp.fg_U_out_multiplier = fg.recurrent_fc_params.output_multiplier; + qp.cg_U_out_multiplier = cg.recurrent_fc_params.output_multiplier; + qp.og_U_out_multiplier = og.recurrent_fc_params.output_multiplier; + qp.ig_W_out_shift = ig.input_fc_params.output_shift; + qp.fg_W_out_shift = fg.input_fc_params.output_shift; + qp.cg_W_out_shift = cg.input_fc_params.output_shift; + qp.og_W_out_shift = og.input_fc_params.output_shift; + qp.ig_U_out_shift = ig.recurrent_fc_params.output_shift; + qp.fg_U_out_shift = fg.recurrent_fc_params.output_shift; + qp.cg_U_out_shift = cg.recurrent_fc_params.output_shift; + qp.og_U_out_shift = og.recurrent_fc_params.output_shift; + qp.quantized_cell_clip = op_data.cell_state_info.quantized_cell_clip; + qp.cell_state_scale = op_data.cell_state_info.cell_state_scale_power; + qp.hidden_multiplier = + op_data.inter_gate_parameters.output_mul_params.output_multiplier; + qp.hidden_shift = op_data.inter_gate_parameters.output_mul_params.output_shift; + // input_zero_bias = -input_zero_point (matches input FC input_offset). + qp.input_zero_bias = fg.input_fc_params.input_offset; + // hidden_zero_bias = hidden_zero_point (output mul offset). + qp.hidden_zero_bias = + op_data.inter_gate_parameters.output_mul_params.output_offset; + + // Zero-initialize so any newer nnlib flag fields (e.g. `back`) default to 0 + // without referencing them explicitly (keeps this header-version agnostic). + lstm_flags flags = {}; + flags.time_major = size_info.time_major; + flags.use_cifg = 0; + + WORD32 err = xa_nn_lstm_sym8sxasym8s_16( + tflite::micro::GetTensorData(kernel_content.output_tensor), + tflite::micro::GetTensorData(kernel_content.HiddenStateTensor()), + tflite::micro::GetTensorData(kernel_content.CellStateTensor()), + &weights, &biases, + (WORD8*)tflite::micro::GetTensorData( + kernel_content.GetInternalTensor(kLstmInputTensor)), + size_info.input_dimension, size_info.state_dimension, + size_info.state_dimension, size_info.batch_size, size_info.time_steps, + size_info.state_dimension, &qp, &flags, p_scratch); + + return (err == 0) ? kTfLiteOk : kTfLiteError; +} +#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + } // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/lstm_eval.h b/tensorflow/lite/micro/kernels/xtensa/lstm_eval.h index a0338b02251..d7595e680ef 100644 --- a/tensorflow/lite/micro/kernels/xtensa/lstm_eval.h +++ b/tensorflow/lite/micro/kernels/xtensa/lstm_eval.h @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// Functions to perform integer evaluation for standard LSTM (e.g., defined in +// Functions to perform integer evaulation for standard LSTM (e.g., defined in // the keras lstm layer, no peephole etc.). Currently used by the 16 bits // activation case only @@ -44,7 +44,7 @@ class LstmTensors { ~LstmTensors(); // Verify the LSTM internal tensor properties (e.g., type checks) - // Input/output/states/fc weights tensors are required for kernel evaluation. + // Input/output/states/fc weights tensors are required for kernel evaulation. // The state tensors should be variables. Variants of the standard LSTM // are not supported here, therefore their corresponding tensors should be // invalid @@ -158,7 +158,7 @@ LSTMBuffers CreateLSTMBuffers(TfLiteContext* context, // namespace to expose them for testing namespace lstm_internal { -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void Sigmoid(const RuntimeShape& data_shape, int16_t* data); void Sigmoid(const RuntimeShape& data_shape, float* data); @@ -200,7 +200,7 @@ void FullyConnected(const FullyConnectedParams& params, const RuntimeShape& filter_shape, const float* filter_data, const RuntimeShape& bias_shape, const float* bias_data, const RuntimeShape& output_shape, float* output_data); -#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void Sigmoid(int16_t* data, int32_t data_size); void Sigmoid(float* data, int32_t data_size); @@ -236,7 +236,7 @@ void FullyConnected(const FullyConnectedParams& params, const float* input_data, const float* filter_data, const float* bias_data, float* output_data, const int num_batches, const int output_depth, const int accum_depth); -#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) void AddElementWise(const int16_t* input_1, const int16_t* input_2, int n_batch, int n_input, int16_t* output); @@ -272,7 +272,7 @@ class LstmStepManager { int OutputOffset() const { return output_offset_; } int HiddenStateOffset() const { return hidden_state_offset_; } int CellStateOffset() const { return cell_state_offset_; } -#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) int time_major() const { return size_info_.time_major; } int batch_size() const { return size_info_.batch_size; } @@ -290,7 +290,7 @@ class LstmStepManager { int hidden_state_offset_ = 0; int cell_state_offset_ = 0; // Sizeinfo is from LstmOpData, which reside in the memory arena - // (guarantee to outlast LSTMStepManager, which reside in stack) + // (guarante to outlast LSTMStepManager, which reside in stack) const LstmSizeInfo& size_info_; }; @@ -298,7 +298,7 @@ class LstmStepManager { // Implements the following formula: // gate = activate(FC(input) + FC(recurrent)) // Activation is sigmoid except for the "cell" gate (configurable, usually tanh) -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) template void CalculateLstmGate( @@ -406,7 +406,7 @@ void UpdateLstmCell(const LstmStepManager& step_info, step_info.CellStateOffset()); } } -#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) template void CalculateLstmGate( @@ -496,7 +496,7 @@ void UpdateLstmCell(const LstmStepManager& step_info, const ArithmeticParams& forget_cell_mul_params, const ArithmeticParams& input_mul_params, const CellStateInfo& cell_state_info, float* buffer); -#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) // Update the hidden state of the LSTM kernel using the following formula: // updated_hidden_state = Tanh(updated_cell_state) * output_gate_output, * means @@ -521,7 +521,7 @@ void UpdateLstmHidden(const LstmStepManager& step_info, tflite::micro::GetTensorData(cell_state) + step_info.CellStateOffset(); // Tanh(cell_state) -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) Tanh(cell_state_scale_power, cell_state_shape, cell_state_data, cell_state_shape, buffer); // Update the hidden state @@ -539,7 +539,7 @@ void UpdateLstmHidden(const LstmStepManager& step_info, #endif } -#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) template void LstmStep(const LstmStepManager& step_info, const OpDataLSTM& op_data, @@ -651,7 +651,7 @@ void LstmStep(const LstmStepManager& step_info, const OpDataLSTM& op_data, step_info.HiddenStateOffset(), step_info.StateShape().FlatSize() * sizeof(ActivationType)); } -#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#else // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) template void LstmStep(const LstmStepManager& step_info, const OpDataLSTM& op_data, @@ -661,14 +661,10 @@ void LstmStep(const LstmStepManager& step_info, const OpDataLSTM& op_data, kernel_content.GetInternalTensor(tflite::kLstmInputTensor); TfLiteEvalTensor* recurrent = kernel_content.HiddenStateTensor(); - const auto& size_info = op_data.size_info; - const int time_major = step_info.time_major(); - const int batch_size = size_info.batch_size; - const int time_steps = size_info.time_steps; - const int num_batches = time_major == 0 ? (time_steps == 1 ? batch_size : 1) - : step_info.batch_size(); - const int input_dimension = step_info.input_dimension(); - const int state_dimension = step_info.state_dimension(); + int time_major = step_info.time_major(); + int num_batches = time_major == 0 ? 1 : step_info.batch_size(); + int input_dimension = step_info.input_dimension(); + int state_dimension = step_info.state_dimension(); // Check offset validity to avoid memory overflow TFLITE_DCHECK_LE(step_info.InputOffset() + num_batches * input_dimension, @@ -786,11 +782,11 @@ void LstmStep(const LstmStepManager& step_info, const OpDataLSTM& op_data, step_info.HiddenStateOffset(), step_info.StateShape().FlatSize() * sizeof(ActivationType)); } -#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5)) +#endif // #if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)) } // namespace lstm_internal -// Evaluate the LSTM kernel with (potential) multi-steps and multi-batch input +// Evaulate the LSTM kernel with (potential) multi-steps and multi-batch input // Since template @@ -799,7 +795,7 @@ TfLiteStatus EvalLstm(const OpDataLSTM& op_data, const LSTMBuffers& buffers) { lstm_internal::LstmStepManager step_info(&op_data.size_info); const auto& size_info = op_data.size_info; - // time is the first dimension, enable batch computation + // time is the first dimention, enable batch computation if (size_info.time_major) { for (int t = 0; t < size_info.time_steps; t++) { lstm_internal::LstmStep( @@ -807,10 +803,8 @@ TfLiteStatus EvalLstm(const OpDataLSTM& op_data, // prepare for the next time step step_info.UpdateTime(); } - } else if (size_info.batch_size > 1 && size_info.time_steps == 1) { - lstm_internal::LstmStep( - step_info, op_data, kernel_content, buffers); } else { + // batch first, unable to size the input data. single batch inference for (int b = 0; b < size_info.batch_size; b++) { for (int t = 0; t < size_info.time_steps; t++) { lstm_internal::LstmStep( @@ -825,6 +819,14 @@ TfLiteStatus EvalLstm(const OpDataLSTM& op_data, } return kTfLiteOk; } + +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) +template <> +TfLiteStatus EvalLstm( + const OpDataLSTM& op_data, LSTMKernelContents& kernel_content, + const LSTMBuffers& buffers); +#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + } // namespace tflite #endif // TENSORFLOW_LITE_MICRO_KERNELS_LSTM_EVAL_16ACT_H_ diff --git a/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.cc b/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.cc index 0f6a02eaf09..175bb277258 100644 --- a/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.cc +++ b/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.cc @@ -84,12 +84,32 @@ TfLiteStatus UnidirectionalSequenceLstmPrepare(TfLiteContext* context, return kTfLiteError; } // request buffers (four buffers) - for (size_t i = 0; i < 4; i++) { + size_t default_buffer_size = op_data->size_info.batch_size * + op_data->size_info.state_dimension * + TfLiteTypeGetSize(cell_state_type); + size_t buffer0_size = default_buffer_size; +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + { + const TfLiteType activation_type = + lstm_tensors.GetInternalTensor(kLstmInputTensor)->type; + const TfLiteType weight_type = + lstm_tensors.GetInternalTensor(kLstmInputToForgetWeightsTensor)->type; + if (activation_type == kTfLiteInt8 && weight_type == kTfLiteInt8) { + int fused_scratch_size = xa_nn_lstm_getsize( + op_data->size_info.batch_size, op_data->size_info.time_steps, + op_data->size_info.state_dimension, /*cell_state_precision=*/16); + if (static_cast(fused_scratch_size) > buffer0_size) { + buffer0_size = static_cast(fused_scratch_size); + } + } + } +#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + TF_LITE_ENSURE_OK(context, + context->RequestScratchBufferInArena( + context, buffer0_size, &(op_data->buffer_indices[0]))); + for (size_t i = 1; i < 4; i++) { TF_LITE_ENSURE_OK(context, context->RequestScratchBufferInArena( - context, - op_data->size_info.batch_size * - op_data->size_info.state_dimension * - TfLiteTypeGetSize(cell_state_type), + context, default_buffer_size, &(op_data->buffer_indices[i]))); } return kTfLiteOk; @@ -110,7 +130,7 @@ TfLiteStatus UnidirectionalSequenceLstmEval(TfLiteContext* context, case kTfLiteFloat32: { LSTMBuffers buffers = CreateLSTMBuffers(context, op_data.buffer_indices); - EvalLstm(op_data, kernel_content, buffers); + return EvalLstm(op_data, kernel_content, buffers); break; } case kTfLiteInt8: { @@ -119,7 +139,7 @@ TfLiteStatus UnidirectionalSequenceLstmEval(TfLiteContext* context, // 8(activation)x8(weight)->16(cell) LSTM with 32 bits bias LSTMBuffers buffers = CreateLSTMBuffers(context, op_data.buffer_indices); - EvalLstm(op_data, kernel_content, + return EvalLstm(op_data, kernel_content, buffers); break; } @@ -137,7 +157,7 @@ TfLiteStatus UnidirectionalSequenceLstmEval(TfLiteContext* context, // 16(activation)x8(weight)->16(cell) LSTM with 64 bits bias LSTMBuffers buffers = CreateLSTMBuffers(context, op_data.buffer_indices); - EvalLstm(op_data, kernel_content, + return EvalLstm(op_data, kernel_content, buffers); break; } diff --git a/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.h b/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.h new file mode 100644 index 00000000000..16aa23b9163 --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/unidirectional_sequence_lstm.h @@ -0,0 +1,47 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +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 TENSORFLOW_LITE_MICRO_KERNELS_UNIDIRECTIONAL_SEQUENCE_LSTM_H_ +#define TENSORFLOW_LITE_MICRO_KERNELS_UNIDIRECTIONAL_SEQUENCE_LSTM_H_ + +#include + +#include "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/types.h" + +namespace tflite { + +// This is the most generic TFLMRegistration. The actual supported types +// may still be target dependent. The only requirement is that every +// implementation (reference or optimized) must define this function. +// TODO(b/230666079): resolve conflict with xtensa implementation +TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); + +#if defined(CMSIS_NN) +// Returns a TFLMRegistration struct for kernel variant that only supports +// int8 activations and int8 weights and uses the latency optimized +// implementations. +TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT8(); + +#else +inline TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT8() { + return Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); +} +#endif + +} // namespace tflite + +#endif // TENSORFLOW_LITE_MICRO_KERNELS_UNIDIRECTIONAL_SEQUENCE_LSTM_H_ diff --git a/tensorflow/lite/micro/kernels/xtensa/xtensa.h b/tensorflow/lite/micro/kernels/xtensa/xtensa.h index 0e7e51b0cb6..cdcb1f75eb3 100644 --- a/tensorflow/lite/micro/kernels/xtensa/xtensa.h +++ b/tensorflow/lite/micro/kernels/xtensa/xtensa.h @@ -22,7 +22,7 @@ limitations under the License. #include "tensorflow/lite/micro/kernels/xtensa/fixedpoint_utils_hifimini.h" #endif // defined(HIFMINI) -#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) #include "include/nnlib/xa_nnlib_api.h" #include "include/nnlib/xa_nnlib_standards.h"