Enable int16_t support for CMSIS-NN LSTM kernel (#2551) BUG=New feature in CMSIS-NN
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/unidirectional_sequence_lstm.cc b/tensorflow/lite/micro/kernels/cmsis_nn/unidirectional_sequence_lstm.cc index 75ba5ea..49da4d9 100644 --- a/tensorflow/lite/micro/kernels/cmsis_nn/unidirectional_sequence_lstm.cc +++ b/tensorflow/lite/micro/kernels/cmsis_nn/unidirectional_sequence_lstm.cc
@@ -56,6 +56,12 @@ arm_vector_sum_s8(kernel_sum, size1, size2, weights, offset, biases); } +void CMSIS_NN_VectorSum(int64_t* kernel_sum, const int32_t size1, + const int32_t size2, const int8_t* weights, + const int32_t offset, const int64_t* biases) { + arm_vector_sum_s8_s64(kernel_sum, size1, size2, weights, offset, biases); +} + template <typename BiasType> TfLiteStatus CMSIS_NN_PortOpData(TfLiteContext* context, OpDataLSTM* params_ref, const LSTMKernelContents& kernel_content, @@ -289,6 +295,32 @@ return kTfLiteOk; } +TfLiteStatus CMSIS_NN_EvalInteger16x8_16Lstm( + const OpData& op_data, const LSTMKernelContents& kernel_content, + const LSTMBuffers<int16_t>& buffers) { + TFLITE_DCHECK( + kernel_content.GetInternalTensor(tflite::kLstmInputTensor)->dims->size >= + 2 && + kernel_content.GetInternalTensor(tflite::kLstmInputTensor)->dims->size <= + 3); + + const int16_t* input = tflite::micro::GetOptionalTensorData<int16_t>( + kernel_content.GetInternalTensor(tflite::kLstmInputTensor)); + int16_t* output = + tflite::micro::GetTensorData<int16_t>(kernel_content.output_tensor); + + // Create lstm buffer struct + cmsis_nn_lstm_context cmsis_buffers; + cmsis_buffers.temp1 = reinterpret_cast<int16_t*>(buffers.buffer0); + cmsis_buffers.temp2 = reinterpret_cast<int16_t*>(buffers.buffer1); + cmsis_buffers.cell_state = reinterpret_cast<int16_t*>(buffers.buffer2); + + arm_lstm_unidirectional_s16(input, output, &op_data.params_cmsis_nn, + &cmsis_buffers); + + return kTfLiteOk; +} + /*Kernel functions*/ void* UnidirectionalSequenceLstmInit(TfLiteContext* context, const char* buffer, size_t length) { @@ -351,6 +383,12 @@ number_of_buffers = 3; CMSIS_NN_PortOpData<int32_t>(context, op_data_lstm, kernel_content, &op_data->params_cmsis_nn); + } else if (activation_type == kTfLiteInt16 && + cell_state_type == kTfLiteInt16) { + auto kernel_content = CreateLSTMKernelContent(context, node); + number_of_buffers = 3; + CMSIS_NN_PortOpData<int64_t>(context, op_data_lstm, kernel_content, + &op_data->params_cmsis_nn); } else { number_of_buffers = 4; } @@ -394,8 +432,7 @@ // 8(activation)x8(weight)->16(cell) LSTM with 32 bits bias LSTMBuffers<int16_t> buffers = CMSIS_NN_CreateLSTMBuffers(context, op_data_lstm.buffer_indices); - return CMSIS_NN_EvalInteger8x8_16Lstm(op_data, kernel_content, - buffers); + CMSIS_NN_EvalInteger8x8_16Lstm(op_data, kernel_content, buffers); break; } default: { @@ -411,9 +448,8 @@ case kTfLiteInt8: { // 16(activation)x8(weight)->16(cell) LSTM with 64 bits bias LSTMBuffers<int16_t> buffers = - CreateLSTMBuffers<int16_t>(context, op_data_lstm.buffer_indices); - EvalLstm<int16_t, int8_t, int16_t, int64_t>(op_data_lstm, - kernel_content, buffers); + CMSIS_NN_CreateLSTMBuffers(context, op_data_lstm.buffer_indices); + CMSIS_NN_EvalInteger16x8_16Lstm(op_data, kernel_content, buffers); break; } default: { @@ -460,6 +496,33 @@ return kTfLiteOk; } +TfLiteStatus UnidirectionalSequenceLstmEvalInt16(TfLiteContext* context, + TfLiteNode* node) { + TFLITE_DCHECK(node->user_data != nullptr); + const OpData& op_data = *reinterpret_cast<const OpData*>(node->user_data); + const OpDataLSTM& op_data_lstm = op_data.params_ref; + auto kernel_content = CreateLSTMKernelContent(context, node); + const auto activation_type = + kernel_content.internal_tensors[kLstmInputTensor]->type; + const auto weight_type = + kernel_content.internal_tensors[kLstmInputToInputWeightsTensor]->type; + + TFLITE_DCHECK(weight_type == kTfLiteInt16 && + "Only int16 filter type supported."); + + if (activation_type == kTfLiteInt16) { + LSTMBuffers<int16_t> buffers = + CMSIS_NN_CreateLSTMBuffers(context, op_data_lstm.buffer_indices); + + return CMSIS_NN_EvalInteger16x8_16Lstm(op_data, kernel_content, buffers); + } else { + MicroPrintf("Input type %s (%d) not supported.", + TfLiteTypeGetName(activation_type), activation_type); + return kTfLiteError; + } + return kTfLiteOk; +} + } // namespace TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM() { @@ -474,4 +537,10 @@ UnidirectionalSequenceLstmEvalInt8); } +TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT16() { + return tflite::micro::RegisterOp(UnidirectionalSequenceLstmInit, + UnidirectionalSequenceLstmPrepare, + UnidirectionalSequenceLstmEvalInt16); +} + } // namespace tflite
diff --git a/tensorflow/lite/micro/kernels/unidirectional_sequence_lstm.h b/tensorflow/lite/micro/kernels/unidirectional_sequence_lstm.h index 16aa23b..46f6b2d 100644 --- a/tensorflow/lite/micro/kernels/unidirectional_sequence_lstm.h +++ b/tensorflow/lite/micro/kernels/unidirectional_sequence_lstm.h
@@ -1,4 +1,4 @@ -/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2024 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. @@ -36,10 +36,19 @@ // implementations. TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT8(); +// Returns a TFLMRegistration struct for kernel variant that only supports +// int16 activations and int8 weights and uses the latency optimized +// implementations. +TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT16(); + #else inline TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT8() { return Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); } + +inline TFLMRegistration Register_UNIDIRECTIONAL_SEQUENCE_LSTM_INT16() { + return Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); +} #endif } // namespace tflite
diff --git a/tensorflow/lite/micro/tools/make/ext_libs/cmsis_nn_download.sh b/tensorflow/lite/micro/tools/make/ext_libs/cmsis_nn_download.sh index fae77aa..482766c 100755 --- a/tensorflow/lite/micro/tools/make/ext_libs/cmsis_nn_download.sh +++ b/tensorflow/lite/micro/tools/make/ext_libs/cmsis_nn_download.sh
@@ -47,9 +47,9 @@ echo >&2 "${DOWNLOADED_CMSIS_NN_PATH} already exists, skipping the download." else - ZIP_PREFIX_NN="6cc31fb36fa330325b2bb0ffde3a7288384e58ab" - CMSIS_NN_URL="http://github.com/ARM-software/CMSIS-NN/archive/6cc31fb36fa330325b2bb0ffde3a7288384e58ab.zip" - CMSIS_NN_MD5="42000f264b93b7b6cd60c1b507792daf" + ZIP_PREFIX_NN="8492d82a1a81651977c5f5128492b4a0f0cf6715" + CMSIS_NN_URL="http://github.com/ARM-software/CMSIS-NN/archive/${ZIP_PREFIX_NN}.zip" + CMSIS_NN_MD5="2cb03e4f044b78af6751009cd53247a8" # wget is much faster than git clone of the entire repo. So we wget a specific # version and can then apply a patch, as needed.