Stacker Op (#2118)

`port c++ stacker op to open source in tflm_signal`

-port stacker op and corresponding to new open source location for C++

BUG=[b/289298641](https://b.corp.google.com/issues/289298641)

python/tflite_micro/python_ops_resolver.cc

@@ -111,6 +111,7 @@ PythonOpsResolver::PythonOpsResolver() {
   AddSquaredDifference();
   AddSqueeze();
   AddStridedSlice();
+  AddStacker();
   AddSub();
   AddSum();
   AddSvdf();

signal/micro/kernels/BUILD

@@ -12,6 +12,7 @@ cc_library(
         "framer.cc",
         "overlap_add.cc",
         "rfft.cc",
+        "stacker.cc",
         "window.cc",
     ],
     hdrs = [
@@ -167,3 +168,29 @@ cc_test(
         "//tensorflow/lite/micro/testing:micro_test",
     ],
 )
+
+cc_library(
+    name = "stacker_flexbuffers_generated_data",
+    srcs = [
+        "stacker_flexbuffers_generated_data.cc",
+    ],
+    hdrs = [
+        "stacker_flexbuffers_generated_data.h",
+    ],
+)
+
+cc_test(
+    name = "stacker_test",
+    srcs = [
+        "stacker_test.cc",
+    ],
+    deps = [
+        ":register_signal_ops",
+        ":stacker_flexbuffers_generated_data",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/micro:op_resolvers",
+        "//tensorflow/lite/micro:test_helpers",
+        "//tensorflow/lite/micro/kernels:kernel_runner",
+        "//tensorflow/lite/micro/testing:micro_test",
+    ],
+)

signal/micro/kernels/stacker.cc

@@ -0,0 +1,176 @@
+/* Copyright 2019 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 <stdint.h>
+
+#include "signal/src/circular_buffer.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/flatbuffer_utils.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/memory_helpers.h"
+#include "tensorflow/lite/micro/micro_utils.h"
+
+namespace tflite {
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+constexpr int kOutputValidTensor = 1;
+
+// Indices into the init flexbuffer's vector.
+// The parameter's name is in the comment that follows.
+// Elements in the vectors are ordered alphabetically by parameter name.
+constexpr int kNumChannelsIndex = 0;          // 'num_channels'
+constexpr int kStackerLeftContextIndex = 1;   // 'stacker_left_context'
+constexpr int kStackerRightContextIndex = 2;  // 'stacker_right_context'
+constexpr int kStackerStepIndex = 3;          // 'stacker_step'
+
+struct TFLMSignalStackerParams {
+  int32_t num_channels;
+  int32_t stacker_left_context;
+  int32_t stacker_right_context;
+  int32_t stacker_step;
+
+  size_t buffer_size;
+  size_t step_size;
+  bool stacker_has_first_frame;
+
+  int8_t* state;
+  tflm_signal::CircularBuffer* circular_buffer;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
+
+  auto* params =
+      static_cast<TFLMSignalStackerParams*>(context->AllocatePersistentBuffer(
+          context, sizeof(TFLMSignalStackerParams)));
+  if (params == nullptr) {
+    return nullptr;
+  }
+
+  tflite::FlexbufferWrapper fbw(buffer_t, length);
+  params->num_channels = fbw.ElementAsInt32(kNumChannelsIndex);
+  params->stacker_left_context = fbw.ElementAsInt32(kStackerLeftContextIndex);
+  params->stacker_right_context = fbw.ElementAsInt32(kStackerRightContextIndex);
+  params->stacker_step = fbw.ElementAsInt32(kStackerStepIndex);
+
+  params->buffer_size =
+      params->num_channels *
+      (params->stacker_left_context + params->stacker_right_context + 1);
+  params->step_size = params->num_channels * params->stacker_step;
+  params->stacker_has_first_frame = false;
+
+  size_t state_size =
+      tflm_signal::CircularBufferGetNeededMemory(params->buffer_size);
+  params->state = static_cast<int8_t*>(
+      context->AllocatePersistentBuffer(context, sizeof(int8_t) * state_size));
+
+  if (params->state == nullptr) {
+    return nullptr;
+  }
+
+  params->circular_buffer = tflm_signal::CircularBufferInit(
+      params->buffer_size, params->state, state_size);
+  return params;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);
+
+  MicroContext* micro_context = GetMicroContext(context);
+
+  TfLiteTensor* input =
+      micro_context->AllocateTempInputTensor(node, kInputTensor);
+  TF_LITE_ENSURE(context, input != nullptr);
+  TfLiteTensor* output =
+      micro_context->AllocateTempOutputTensor(node, kOutputTensor);
+  TF_LITE_ENSURE(context, output != nullptr);
+  TfLiteTensor* output_valid =
+      micro_context->AllocateTempOutputTensor(node, kOutputValidTensor);
+  TF_LITE_ENSURE(context, output_valid != nullptr);
+
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(output), 1);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(output_valid), 0);
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt16);
+  TF_LITE_ENSURE_TYPES_EQ(context, output_valid->type, kTfLiteBool);
+
+  micro_context->DeallocateTempTfLiteTensor(input);
+  micro_context->DeallocateTempTfLiteTensor(output);
+  micro_context->DeallocateTempTfLiteTensor(output_valid);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params = reinterpret_cast<TFLMSignalStackerParams*>(node->user_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  const TfLiteEvalTensor* input =
+      tflite::micro::GetEvalInput(context, node, kInputTensor);
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+  TfLiteEvalTensor* output_valid =
+      tflite::micro::GetEvalOutput(context, node, kOutputValidTensor);
+
+  const int16_t* input_data = tflite::micro::GetTensorData<int16_t>(input);
+
+  tflm_signal::CircularBufferWrite(params->circular_buffer, input_data,
+                                   params->num_channels);
+
+  // The first frame is replicated an extra left_context times to pad.
+  if (params->stacker_has_first_frame == false) {
+    tflm_signal::CircularBufferExtend(params->circular_buffer,
+                                      params->num_channels,
+                                      params->stacker_left_context);
+    params->stacker_has_first_frame = true;
+  }
+
+  int16_t* output_data = tflite::micro::GetTensorData<int16_t>(output);
+  bool* output_valid_data = tflite::micro::GetTensorData<bool>(output_valid);
+  if (tflm_signal::CircularBufferAvailable(params->circular_buffer) >=
+      params->buffer_size) {
+    tflm_signal::CircularBufferGet(params->circular_buffer, params->buffer_size,
+                                   output_data);
+    tflm_signal::CircularBufferDiscard(params->circular_buffer,
+                                       params->step_size);
+    *output_valid_data = true;
+  } else {
+    *output_valid_data = false;
+  }
+  return kTfLiteOk;
+}
+
+void Reset(TfLiteContext* context, void* buffer) {
+  auto* params = static_cast<TFLMSignalStackerParams*>(buffer);
+  tflm_signal::CircularBufferReset(params->circular_buffer);
+  params->stacker_has_first_frame = false;
+}
+
+}  // namespace
+
+namespace tflm_signal {
+TFLMRegistration* Register_STACKER() {
+  static TFLMRegistration r =
+      tflite::micro::RegisterOp(Init, Prepare, Eval, /*Free*/ nullptr, Reset);
+  return &r;
+}
+}  // namespace tflm_signal
+
+}  // namespace tflite

signal/micro/kernels/stacker_flexbuffers_generated_data.cc

@@ -0,0 +1,41 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+// This file is generated. See:
+// tensorflow/lite/micro/kernels/test_data_generation/README.md
+
+#include "signal/micro/kernels/stacker_flexbuffers_generated_data.h"
+
+const int g_gen_data_size_stacker_3_channels_step_1 = 88;
+const unsigned char g_gen_data_stacker_3_channels_step_1[] = {
+    0x6e, 0x75, 0x6d, 0x5f, 0x63, 0x68, 0x61, 0x6e, 0x6e, 0x65, 0x6c,
+    0x73, 0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x6c,
+    0x65, 0x66, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x78, 0x74,
+    0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x72, 0x69,
+    0x67, 0x68, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x78, 0x74,
+    0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x73, 0x74,
+    0x65, 0x70, 0x00, 0x04, 0x46, 0x3a, 0x26, 0x11, 0x04, 0x01, 0x04,
+    0x03, 0x01, 0x00, 0x01, 0x04, 0x04, 0x04, 0x04, 0x08, 0x24, 0x01,
+};
+const int g_gen_data_size_stacker_10_channels_step_2 = 88;
+const unsigned char g_gen_data_stacker_10_channels_step_2[] = {
+    0x6e, 0x75, 0x6d, 0x5f, 0x63, 0x68, 0x61, 0x6e, 0x6e, 0x65, 0x6c,
+    0x73, 0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x6c,
+    0x65, 0x66, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x78, 0x74,
+    0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x72, 0x69,
+    0x67, 0x68, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x78, 0x74,
+    0x00, 0x73, 0x74, 0x61, 0x63, 0x6b, 0x65, 0x72, 0x5f, 0x73, 0x74,
+    0x65, 0x70, 0x00, 0x04, 0x46, 0x3a, 0x26, 0x11, 0x04, 0x01, 0x04,
+    0x0a, 0x01, 0x00, 0x02, 0x04, 0x04, 0x04, 0x04, 0x08, 0x24, 0x01,
+};

signal/micro/kernels/stacker_flexbuffers_generated_data.h

@@ -0,0 +1,25 @@
+/* Copyright 2020 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 SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_STACKER_FLEXBUFFERS_DATA_H_
+#define SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_STACKER_FLEXBUFFERS_DATA_H_
+
+extern const int g_gen_data_size_stacker_3_channels_step_1;
+extern const unsigned char g_gen_data_stacker_3_channels_step_1[];
+
+extern const int g_gen_data_size_stacker_10_channels_step_2;
+extern const unsigned char g_gen_data_stacker_10_channels_step_2[];
+
+#endif  // SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_STACKER_FLEXBUFFERS_DATA_H_