Optimization in 8x16 lstm layer FC for HiFi. (#3227)

pramods-cad · veblush · web-flow · commit cee9550a3770 · 2025-10-17T17:29:10.000Z
* Optimization in 8x16 lstm layer for HiFi by using xa_nn_matXvec_v2_sym8sxsym16s_sym16s with appropriate conditions for FC.

* Updated HiFi4/HiFi5 NNLibs to 30 Sept 2025 release, made necessary changes in kernel integration code.

* Fixed formatting errors.

---------

Co-authored-by: Esun Kim &lt;veblush@google.com&gt;
diff --git a/tensorflow/lite/micro/kernels/xtensa/conv_hifi.cc b/tensorflow/lite/micro/kernels/xtensa/conv_hifi.cc
@@ -81,6 +81,7 @@ TfLiteStatus ConvPrepareHifi(TfLiteContext* context, TfLiteNode* node) {
   const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
   const int filter_height = filter_shape.Dims(1);
   const int filter_width = filter_shape.Dims(2);
+  const int filter_depth = filter_shape.Dims(3);
   const int output_height = output_shape.Dims(1);
   const int output_width = output_shape.Dims(2);
   const int output_channels = output_shape.Dims(3);
@@ -94,43 +95,19 @@ TfLiteStatus ConvPrepareHifi(TfLiteContext* context, TfLiteNode* node) {
   if ((params->dilation_width_factor == 1) &&
       (params->dilation_height_factor == 1)) {
     if (input->type == kTfLiteInt8) {
-      if (input_height == 1 && filter_height == 1 && output_height == 1) {
-        int inp_h, filt_h, filt_w, str_h, pad_h, out_h;
-        inp_h = input_width;
-        filt_h = filter_width;
-        filt_w = filter_height;
-        str_h = stride_width;
-        pad_h = pad_width;
-        out_h = output_width;
-        required_scratch = xa_nn_conv2d_std_getsize(
-            inp_h, input_depth, filt_h, filt_w, str_h, pad_h, out_h,
-            output_channels, PREC_ASYM8S);
-      } else {
-        required_scratch = xa_nn_conv2d_std_getsize(
-            input_height, input_depth, filter_height, filter_width,
-            stride_height, pad_height, output_height, output_channels,
-            PREC_ASYM8S);
-      }
+      required_scratch = xa_nn_conv2d_std_getsize(
+          input_height, input_width, input_depth, filter_height, filter_width,
+          filter_depth, stride_height, pad_height, stride_width, pad_width,
+          output_height, output_width, output_channels, PREC_ASYM8S, PREC_SYM8S,
+          1, 1, 0);
       TF_LITE_ENSURE(context, required_scratch > 0);
     }
     if (input->type == kTfLiteInt16) {
-      if (input_height == 1 && filter_height == 1 && output_height == 1) {
-        int inp_h, filt_h, filt_w, str_h, pad_h, out_h;
-        inp_h = input_width;
-        filt_h = filter_width;
-        filt_w = filter_height;
-        str_h = stride_width;
-        pad_h = pad_width;
-        out_h = output_width;
-        required_scratch = xa_nn_conv2d_std_getsize(
-            inp_h, input_depth, filt_h, filt_w, str_h, pad_h, out_h,
-            output_channels, PREC_SYM16S);
-      } else {
-        required_scratch = xa_nn_conv2d_std_getsize(
-            input_height, input_depth, filter_height, filter_width,
-            stride_height, pad_height, output_height, output_channels,
-            PREC_SYM16S);
-      }
+      required_scratch = xa_nn_conv2d_std_getsize(
+          input_height, input_width, input_depth, filter_height, filter_width,
+          filter_depth, stride_height, pad_height, stride_width, pad_width,
+          output_height, output_width, output_channels, PREC_SYM16S, PREC_SYM8S,
+          1, 1, 0);
       TF_LITE_ENSURE(context, required_scratch > 0);
     }
   }
diff --git a/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc b/tensorflow/lite/micro/kernels/xtensa/lstm_eval.cc
@@ -305,16 +305,31 @@ void FullyConnected(const FullyConnectedParams& params,
   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) {
-  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);
+  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);
+  }
+  (void)err;
   return;
 }
 
diff --git a/tensorflow/lite/micro/kernels/xtensa/transpose_conv.cc b/tensorflow/lite/micro/kernels/xtensa/transpose_conv.cc
@@ -204,7 +204,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     scratch_buffer_size = xa_nn_transpose_conv_getsize(
         input_height, input_width, input_depth, filter_height, filter_width,
         stride_width, stride_height, output_height, output_width, num_channels,
-        PREC_SYM8S, PREC_ASYM8S);
+        1 /* num_groups */, PREC_SYM8S, PREC_ASYM8S);
     TFLITE_DCHECK(context->RequestScratchBufferInArena(
                       context, scratch_buffer_size,
                       &(data->scratch_buffer_index)) == kTfLiteOk);
@@ -232,7 +232,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     scratch_buffer_size = xa_nn_transpose_conv_getsize(
         input_height, input_width, input_depth, filter_height, filter_width,
         stride_width, stride_height, output_height, output_width, num_channels,
-        PREC_SYM8S, PREC_SYM16S);
+        1 /* num_groups */, PREC_SYM8S, PREC_SYM16S);
     TFLITE_DCHECK(context->RequestScratchBufferInArena(
                       context, scratch_buffer_size,
                       &(data->scratch_buffer_index)) == kTfLiteOk);
@@ -404,9 +404,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               stride_width, stride_height, pad_width, pad_height, input_depth,
               output_depth, input_height, input_width, filter_height,
               filter_width, output_height, output_width, num_elements / batches,
-              data.params.input_offset, data.params.output_offset,
-              data.per_channel_output_shift, data.per_channel_output_multiplier,
-              scratch_buffer);
+              1 /* num_groups */, data.params.input_offset,
+              data.params.output_offset, data.per_channel_output_shift,
+              data.per_channel_output_multiplier, scratch_buffer);
         }
       } else {
         reference_integer_ops::TransposeConv(
@@ -539,8 +539,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               stride_width, stride_height, pad_width, pad_height, input_depth,
               output_depth, input_height, input_width, filter_height,
               filter_width, output_height, output_width, num_elements / batches,
-              data.per_channel_output_shift, data.per_channel_output_multiplier,
-              scratch_buffer);
+              1 /* num_groups */, data.per_channel_output_shift,
+              data.per_channel_output_multiplier, scratch_buffer);
         }
 #else  // #if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
         reference_integer_ops::TransposeConv(
diff --git a/tensorflow/lite/micro/tools/make/ext_libs/xa_nnlib_hifi4.patch b/tensorflow/lite/micro/tools/make/ext_libs/xa_nnlib_hifi4.patch
diff --git a/tensorflow/lite/micro/tools/make/ext_libs/xtensa_download.sh b/tensorflow/lite/micro/tools/make/ext_libs/xtensa_download.sh
@@ -41,18 +41,17 @@ DOWNLOADS_DIR=${1}
 PATCH=""
 
 if [[ ${2} == "hifi3" ]]; then
-  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_hifi4_09_05_2023.zip"
+  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_hifi4_09_30_2025.zip"
   LIBRARY_DIRNAME="xa_nnlib_hifi4"
-  LIBRARY_MD5="2a54e056aef73a4fcffde4643998501a"
+  LIBRARY_MD5="a1cba57501a35d0bffdd39c536689366"
 elif [[ ${2} == "hifi4" ]]; then
-  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_hifi4_09_05_2023.zip"
+  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_hifi4_09_30_2025.zip"
   LIBRARY_DIRNAME="xa_nnlib_hifi4"
-  LIBRARY_MD5="2a54e056aef73a4fcffde4643998501a"
-  PATCH="../../ext_libs/xa_nnlib_hifi4.patch"
+  LIBRARY_MD5="a1cba57501a35d0bffdd39c536689366"
 elif [[ ${2} == "hifi5" ]]; then
-  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi5/raw/master/archive/xa_nnlib_hifi5_09_05_2023.zip"
+  LIBRARY_URL="http://github.com/foss-xtensa/nnlib-hifi5/raw/master/archive/xa_nnlib_hifi5_09_30_2025.zip"
   LIBRARY_DIRNAME="xa_nnlib_hifi5"
-  LIBRARY_MD5="1deb55ef200bf5dbedc70b99b02140c0"
+  LIBRARY_MD5="f3ac445ae47143b1fddf26b85d763697"
 elif [[ ${2} == "vision_p6" ]]; then
   LIBRARY_URL="https://github.com/foss-xtensa/tflmlib_vision/raw/main/archive/xi_tflmlib_vision_p6_22_06_29.zip"
   LIBRARY_DIRNAME="xi_tflmlib_vision_p6"
