Adding streaming_conv reference code. (#2521)

* Adding streaming_conv reference code.

* Streaming conv support under xtensa

* Enabled HiFi5 support for streaming conv kernel.

* Add bazel build rules for streaming_conv

* Fix formatting

* Fix dims_shape array type

* Fix VP6 build

* Fix asan error

* Disable streaming conv test on HIFIMINI & VP6

---------

Co-authored-by: Nileshkumar Vaishnav <nileshv@cadence.com>
Co-authored-by: RJ Ascani <rjascani@google.com>

Author: 3 people

tensorflow/lite/micro/kernels/BUILD

@@ -282,6 +282,8 @@ tflm_kernel_cc_library(
         "split_v.cc",
         "squared_difference.cc",
         "squeeze.cc",
+        "streaming_conv.cc",
+        "streaming_conv_common.cc",
         "strided_slice.cc",
         "strided_slice_common.cc",
         "sub.cc",
@@ -321,6 +323,7 @@ tflm_kernel_cc_library(
         "reduce.h",
         "reshape.h",
         "softmax.h",
+        "streaming_conv.h",
         "strided_slice.h",
         "sub.h",
         "svdf.h",

tensorflow/lite/micro/kernels/conv.cc

@@ -147,9 +147,4 @@ TFLMRegistration Register_CONV_2D() {
   return tflite::micro::RegisterOp(ConvInit, ConvPrepare, ConvEval);
 }
 
-TFLMRegistration Register_STREAMING_CONV_2D() {
-  // TODO(rjascani): These should be replaced with Streaming wrapper functions.
-  return tflite::micro::RegisterOp(ConvInit, ConvPrepare, ConvEval);
-}
-
 }  // namespace tflite

tensorflow/lite/micro/kernels/conv_test.cc

@@ -56,6 +56,18 @@ static TfLiteConvParams common_conv_params = {
     kTfLiteNoType         // quantized_bias_type
 };
 
+#if not defined(HIFIMINI) && not defined(VISION_P6)
+static TfLiteConvParams common_streaming_conv_params = {
+    kTfLitePaddingValid,  // padding
+    1,                    // stride_width
+    1,                    // stride_height
+    kTfLiteActNone,       // activation
+    1,                    // dilation_width_factor
+    1,                    // dilation_height_factor
+    kTfLiteNoType         // quantized_bias_type
+};
+#endif
+
 }  // namespace
 }  // namespace testing
 }  // namespace tflite
@@ -216,6 +228,55 @@ TF_LITE_MICRO_TEST(HybridModeIsError) {
                                   tflite::Register_CONV_2D(), output_data));
 }
 
+#if not defined(HIFIMINI) && not defined(VISION_P6)
+TF_LITE_MICRO_TEST(SimpleTestQuantized16x8PerChannel64bBiasStreaming) {
+  constexpr int kInputElementsStreamConv = 20;
+  static int kInputShapeStreamConv[] = {4, 2, 5, 1, 2};
+  static const float kInputDataStreamConv[kInputElementsStreamConv] = {
+      1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5};
+
+  constexpr int kFilterElementsStreamConv = 8;
+  static int kFilterShapeStreamConv[] = {4, 1, 2, 2, 2};
+  static const float kFilterDataStreamConv[kFilterElementsStreamConv] = {
+      1, 1, 1, 1, 1, 1, 1, 1};
+
+  constexpr int kBiasElementsStreamConv = 1;
+  static int kBiasShapeStreamConv[] = {1, 1};
+  static const float kBiasDataStreamConv[kBiasElementsStreamConv] = {0};
+
+  constexpr int kOutputElementsStreamConv = 8;
+  static int kOutputShapeStreamConv[] = {4, 2, 4, 1, 1};
+  static const float kGoldenDataStreamConv[kOutputElementsStreamConv] = {
+      6, 10, 14, 18, 12, 20, 28, 36};
+
+  int16_t output_data[kOutputElementsStreamConv];
+
+  const float input_scale = 0.5f;
+  const float output_scale = 1.0f;
+  const int input_zero_point = 0;
+  const int output_zero_point = 0;
+
+  int16_t input_quantized[kInputElementsStreamConv];
+  int8_t filter_quantized[kFilterElementsStreamConv];
+  std::int64_t bias_quantized[kBiasElementsStreamConv];
+  int16_t golden_quantized[kOutputElementsStreamConv];
+  int zero_points[kBiasElementsStreamConv + 1];
+  float scales[kBiasElementsStreamConv + 1];
+
+  TF_LITE_MICRO_EXPECT_EQ(
+      kTfLiteOk,
+      tflite::testing::TestConvQuantizedPerChannel(
+          kInputShapeStreamConv, kInputDataStreamConv, input_quantized,
+          input_scale, input_zero_point, kFilterShapeStreamConv,
+          kFilterDataStreamConv, filter_quantized, kBiasShapeStreamConv,
+          kBiasDataStreamConv, bias_quantized, scales, zero_points,
+          kOutputShapeStreamConv, kGoldenDataStreamConv, golden_quantized,
+          output_scale, output_zero_point,
+          &tflite::testing::common_streaming_conv_params,
+          tflite::Register_STREAMING_CONV_2D(), output_data));
+}
+#endif
+
 TF_LITE_MICRO_TEST(SimpleTestQuantized16x8PerChannel64bBias) {
   const int output_dims_count = 12;
   int16_t output_data[output_dims_count];

tensorflow/lite/micro/kernels/streaming_conv.cc

@@ -0,0 +1,236 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/lite/micro/kernels/streaming_conv.h"
+
+#include "tensorflow/lite/c/builtin_op_data.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/reference/conv.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/conv.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/micro_log.h"
+
+namespace tflite {
+namespace {
+
+void updateStreamingConvBuffer(int16_t* pbuf, const int16_t* pinp, int ih,
+                               int ic, int kw) {
+  /* This will update the persistent mem. Input data is coming in as ihx1xic,
+   * buffer is ihxkwxic */
+  int i, j;
+
+  /* step 1 : Move the older columns in persistent buffer */
+  int striplength = (kw - 1) * ic;
+  int16_t* pbuf_dst = pbuf;
+  int16_t* pbuf_src = pbuf + ic;
+
+  for (i = 0; i < ih; i++) {
+    memmove(pbuf_dst, pbuf_src, striplength * sizeof(int16_t));
+    pbuf_dst += kw * ic;
+    pbuf_src += kw * ic;
+  }
+
+  /* step 2: Copy new input column to buffer */
+  int16_t* pbuf_base = pbuf + (ic * (kw - 1));
+  const int16_t* pinp_base = pinp;
+
+  for (i = 0; i < ih; i++) {
+    for (j = 0; j < ic; j++) {
+      pbuf_base[j] = pinp_base[j];
+    }
+    pbuf_base += (kw * ic);
+    pinp_base += ic;
+  }
+}
+
+template <typename AccumScalar>
+inline void StreamingConvPerChannel(
+    const ConvParams& params, const int32_t* output_multiplier,
+    const int32_t* output_shift, const RuntimeShape& input_shape,
+    const int16_t* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const AccumScalar* bias_data, const RuntimeShape& output_shape,
+    int16_t* output_data, int16_t* input_state) {
+  // Get parameters.
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+
+  // Set min and max value of the output.
+  const int32_t output_activation_min = params.quantized_activation_min;
+  const int32_t output_activation_max = params.quantized_activation_max;
+
+  // Consistency check.
+  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_depth = input_shape.Dims(3);
+  const int output_depth = MatchingDim(filter_shape, 0, output_shape, 3);
+  if (bias_data) {
+    TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
+  }
+
+  // Check dimensions of the tensors.
+  const int input_height = input_shape.Dims(1);
+  int input_width = input_shape.Dims(2);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int filter_input_depth = filter_shape.Dims(3);
+
+  /* Update streaming conv buffer with input data */
+  input_width = filter_width;
+  const int32_t dims_shape[4] = {1, input_height, filter_width, input_depth};
+  RuntimeShape input_state_shape(4, dims_shape);
+
+  const int groups = input_depth / filter_input_depth;
+  TFLITE_DCHECK_EQ(input_depth % filter_input_depth, 0);
+  const int filters_per_group = output_depth / groups;
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  for (int batch = 0; batch < batches; ++batch) {
+    updateStreamingConvBuffer(input_state,
+                              &input_data[Offset(input_shape, batch, 0, 0, 0)],
+                              input_height, input_depth, filter_width);
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      const int in_y_origin = (out_y * stride_height) - pad_height;
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        const int in_x_origin = (out_x * stride_width) - pad_width;
+        for (int out_channel = 0; out_channel < output_depth; ++out_channel) {
+          auto group = out_channel / filters_per_group;
+          AccumScalar acc = 0;
+          for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+            const int in_y = in_y_origin + dilation_height_factor * filter_y;
+            for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+              const int in_x = in_x_origin + dilation_width_factor * filter_x;
+
+              // Zero padding by omitting the areas outside the image.
+              const bool is_point_inside_image =
+                  (in_x >= 0) && (in_x < input_width) && (in_y >= 0) &&
+                  (in_y < input_height);
+
+              if (!is_point_inside_image) {
+                continue;
+              }
+
+              for (int in_channel = 0; in_channel < filter_input_depth;
+                   ++in_channel) {
+                int32_t input_val = input_state[Offset(
+                    input_state_shape, 0, in_y, in_x,
+                    in_channel + group * filter_input_depth)];
+                int32_t filter_val = filter_data[Offset(
+                    filter_shape, out_channel, filter_y, filter_x, in_channel)];
+                // Accumulate with 64 bits accumulator.
+                // int64_t += int8_t * int16_t so the highest value we can
+                // get from each accumulation is [-127, 127] * ([-32768,
+                // 32767] -
+                // [-32768, 32767]), which is [-8322945, 8322945].
+                // log2(8322945) = 22.99.
+                acc += filter_val * input_val;
+              }
+            }
+          }
+          if (bias_data) {
+            acc += bias_data[out_channel];
+          }
+          int32_t scaled_acc = MultiplyByQuantizedMultiplier(
+              acc, output_multiplier[out_channel], output_shift[out_channel]);
+          scaled_acc = std::max(scaled_acc, output_activation_min);
+          scaled_acc = std::min(scaled_acc, output_activation_max);
+          output_data[Offset(output_shape, batch, out_y, out_x, out_channel)] =
+              static_cast<int16_t>(scaled_acc);
+        }
+      }
+    }
+  }
+}
+
+TfLiteStatus StreamingConvEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteEvalTensor* input =
+      tflite::micro::GetEvalInput(context, node, kConvInputTensor);
+  const TfLiteEvalTensor* filter =
+      tflite::micro::GetEvalInput(context, node, kConvWeightsTensor);
+  const TfLiteEvalTensor* bias =
+      (NumInputs(node) == 3)
+          ? tflite::micro::GetEvalInput(context, node, kConvBiasTensor)
+          : nullptr;
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kConvOutputTensor);
+
+  TFLITE_DCHECK(node->builtin_data != nullptr);
+  const auto& params =
+      *(reinterpret_cast<TfLiteConvParams*>(node->builtin_data));
+  TFLITE_DCHECK(node->user_data != nullptr);
+  const auto& sdata =
+      *(static_cast<const OpDataStreamingConv*>(node->user_data));
+  const auto& data = sdata.op_data;
+
+  switch (input->type) {  // Already know in/out types are same.
+    case kTfLiteInt16: {
+      if (bias == nullptr || bias->type == kTfLiteInt32) {
+        StreamingConvPerChannel(
+            StreamingConvParamsQuantized(params, data),
+            data.per_channel_output_multiplier, data.per_channel_output_shift,
+            tflite::micro::GetTensorShape(input),
+            tflite::micro::GetTensorData<int16_t>(input),
+            tflite::micro::GetTensorShape(filter),
+            tflite::micro::GetTensorData<int8_t>(filter),
+            tflite::micro::GetTensorShape(bias),
+            tflite::micro::GetOptionalTensorData<std::int32_t>(bias),
+            tflite::micro::GetTensorShape(output),
+            tflite::micro::GetTensorData<int16_t>(output),
+            (int16_t*)sdata.input_state);
+      } else if (bias->type == kTfLiteInt64) {
+        StreamingConvPerChannel(
+            StreamingConvParamsQuantized(params, data),
+            data.per_channel_output_multiplier, data.per_channel_output_shift,
+            tflite::micro::GetTensorShape(input),
+            tflite::micro::GetTensorData<int16_t>(input),
+            tflite::micro::GetTensorShape(filter),
+            tflite::micro::GetTensorData<int8_t>(filter),
+            tflite::micro::GetTensorShape(bias),
+            tflite::micro::GetOptionalTensorData<std::int64_t>(bias),
+            tflite::micro::GetTensorShape(output),
+            tflite::micro::GetTensorData<int16_t>(output),
+            (int16_t*)sdata.input_state);
+      } else {
+        MicroPrintf("Bias type %s (%d) not supported.",
+                    TfLiteTypeGetName(bias->type), bias->type);
+        return kTfLiteError;
+      }
+      break;
+    }
+    default:
+      MicroPrintf("Type %s (%d) not supported.", TfLiteTypeGetName(input->type),
+                  input->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+TFLMRegistration Register_STREAMING_CONV_2D() {
+  return tflite::micro::RegisterOp(StreamingConvInit, StreamingConvPrepare,
+                                   StreamingConvEval);
+}
+
+}  // namespace tflite