Add Xtensa reference/trivial data-movement kernels

tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc

@@ -0,0 +1,138 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/lite/micro/kernels/circular_buffer.h"
+
+#include "tensorflow/lite/c/builtin_op_data.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/micro/flatbuffer_utils.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/kernels/xtensa/xtensa.h"
+
+/*
+ * The circular buffer custom operator is used to implement strided streaming
+ * convolutions on TFLite Micro.  Each time this operator is invoked, it checks
+ * whether or not to run, based on a predetermined stride in time.  If the op
+ * runs, it inserts the input into the end of the output buffer and shifts the
+ * output values towards the start of the buffer.  It discards the oldest value
+ * in the output buffer.
+ *
+ * Input: [<input N+1]
+ * Before shifting:
+ * Output: [<input 1>, <input 2>, <input ...>, <input N>]
+ *
+ * After shifting:
+ * Output: [<input 2>, <input 3>, <input ...>, <input N+1>]
+ *
+ * We make some assumptions in this custom operator:
+ * - Input shape must be [1, 1, 1, depth]
+ * - Output shape must be [1, num_slots, 1, depth]
+ * - Input and output types must match.
+ * - Input and output quantization params must be identical.
+ */
+namespace tflite {
+
+void* CircularBufferInit(TfLiteContext* context, const char* buffer,
+                         size_t length) {
+  TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
+  OpDataCircularBuffer* op_data = static_cast<OpDataCircularBuffer*>(
+      context->AllocatePersistentBuffer(context, sizeof(OpDataCircularBuffer)));
+
+  if (buffer != nullptr && length > 0) {
+    const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
+    tflite::FlexbufferWrapper wrapper(buffer_t, length);
+    op_data->cycles_max = wrapper.ElementAsInt32(kCircularBufferCyclesMaxIndex);
+  } else {
+    op_data->cycles_max = 0;
+  }
+
+  return op_data;
+}
+
+// Shifts buffer over by the output depth, and write new input to end of buffer.
+// num_slots is the number of samples stored in the output buffer.
+// depth is the size of each sample.
+void EvalInt8(const int8_t* input, int num_slots, int depth, int8_t* output) {
+  memmove(output, &output[depth], (num_slots - 1) * depth);
+  memcpy(&output[(num_slots - 1) * depth], input, depth);
+}
+
+TfLiteStatus CircularBufferEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteEvalTensor* input =
+      tflite::micro::GetEvalInput(context, node, kCircularBufferInputTensor);
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kCircularBufferOutputTensor);
+
+  TFLITE_DCHECK(node->user_data != nullptr);
+  OpDataCircularBuffer* data =
+      reinterpret_cast<OpDataCircularBuffer*>(node->user_data);
+
+  int num_slots = output->dims->data[1];
+  int depth = output->dims->data[2] * output->dims->data[3];
+
+  if (input->type == kTfLiteInt8) {
+#if defined(HIFI5) || defined(HIFI4) || defined(HIFI_IQ)
+    const int8_t* xa_input;
+    int8_t* xa_output;
+    int err;
+    xa_input = tflite::micro::GetTensorData<int8_t>(input);
+    xa_output =tflite::micro::GetTensorData<int8_t>(output);
+    err = xa_nn_memmove_8_8(xa_output, &xa_output[depth], (num_slots-1)*depth);
+    TF_LITE_ENSURE(context, (err==0) );
+    memcpy(&xa_output[(num_slots - 1) * depth], xa_input, depth);
+#else
+    EvalInt8(tflite::micro::GetTensorData<int8_t>(input), num_slots, depth,
+             tflite::micro::GetTensorData<int8_t>(output));
+#endif // defined(HIFI5) || defined(HIFI4)
+  } else {
+    TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.",
+                       TfLiteTypeGetName(input->type), input->type);
+    return kTfLiteError;
+  }
+
+  if (--data->cycles_until_run != 0) {
+    // Signal the interpreter to end current run if the delay before op invoke
+    // has not been reached.
+    // TODO(b/149795762): Add kTfLiteAbort to TfLiteStatus enum.
+    return static_cast<TfLiteStatus>(kTfLiteAbort);
+  }
+
+  data->cycles_until_run = data->cycles_max;
+
+  return kTfLiteOk;
+}
+
+// Restores period counter (cycles_until_run) on reset. Buffer memory cleanup
+// is not needed here: because the output tensor is a variable tensor.
+// ResetVariableTensors() automatically zero-points its memory upon reset.
+void CircularBufferReset(TfLiteContext* context, void* buffer) {
+  TFLITE_DCHECK(buffer != nullptr);
+  OpDataCircularBuffer* data = static_cast<OpDataCircularBuffer*>(buffer);
+  data->cycles_until_run = data->cycles_max;
+}
+
+TFLMRegistration* Register_CIRCULAR_BUFFER() {
+  static TFLMRegistration r = tflite::micro::RegisterOp(
+      CircularBufferInit, CircularBufferPrepare, CircularBufferEval,
+      /*free=*/nullptr, CircularBufferReset);
+  return &r;
+}
+
+}  // namespace tflite

tensorflow/lite/micro/kernels/xtensa/expand_dims.cc

@@ -0,0 +1,160 @@
+/* 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.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/micro_log.h"
+#include "tensorflow/lite/micro/micro_utils.h"
+
+namespace tflite {
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kAxisTensor = 1;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus GetAxisValueFromTensor(TfLiteContext* context,
+                                    const TfLiteTensor* axis,
+                                    int32_t* axis_value) {
+  const int axis_dims = (tflite::GetTensorShape(axis)).DimensionsCount();
+  if (axis_dims > 1) {
+    MicroPrintf("Axis has only one element for Expand_Dims.", axis_dims);
+    return kTfLiteError;
+  }
+
+  if (kTfLiteInt32 == (axis->type)) {
+    const int32_t* axis_ptr = tflite::GetTensorData<int32_t>(axis);
+    *axis_value = axis_ptr[0];
+    return kTfLiteOk;
+  } else {
+    MicroPrintf("Axis type %s (%d) not supported by Expand_Dims.",
+                TfLiteTypeGetName(axis->type), axis->type);
+    return kTfLiteError;
+  }
+}
+
+// Verifies that the output tensor's dimension shape is equivalent to inserting
+// a dimension of length 1 at the dimension index axis of input's shape as
+// defined in https://www.tensorflow.org/api_docs/python/tf/expand_dims.
+TfLiteStatus VerifyTensorDim(TfLiteContext* context, const TfLiteTensor* input,
+                             const TfLiteTensor* axis_tensor,
+                             const TfLiteTensor* output) {
+  int32_t axis_value = 0;
+  TF_LITE_ENSURE_OK(context,
+                    GetAxisValueFromTensor(context, axis_tensor, &axis_value));
+
+  tflite::RuntimeShape input_shape = tflite::GetTensorShape(input);
+  if (axis_value < 0) {
+    axis_value = input_shape.DimensionsCount() + 1 + axis_value;
+  }
+  TF_LITE_ENSURE(context, axis_value <= input_shape.DimensionsCount());
+
+  // TFLM only supports fixed dimension tensor and assumes that the output shape
+  // is fully specified in the model. As such, TFLM directly use the pointer to
+  // the dimension array in the model buffer.
+  tflite::RuntimeShape output_shape = tflite::GetTensorShape(output);
+
+  TF_LITE_ENSURE(context, output_shape.DimensionsCount() ==
+                              input_shape.DimensionsCount() + 1);
+  for (int i = 0; i < output_shape.DimensionsCount(); ++i) {
+    if (i < axis_value) {
+      TF_LITE_ENSURE(context, output_shape.Dims(i) == input_shape.Dims(i));
+    } else if (i == axis_value) {
+      TF_LITE_ENSURE(context, output_shape.Dims(i) == 1);
+    } else {
+      TF_LITE_ENSURE(context, output_shape.Dims(i) == input_shape.Dims(i - 1));
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ExpandDimsPrepare(TfLiteContext* context, TfLiteNode* node) {
+  MicroContext* micro_context = GetMicroContext(context);
+
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  TfLiteTensor* input =
+      micro_context->AllocateTempInputTensor(node, kInputTensor);
+  TF_LITE_ENSURE(context, input != nullptr);
+  TfLiteTensor* axis =
+      micro_context->AllocateTempInputTensor(node, kAxisTensor);
+  TF_LITE_ENSURE(context, axis != nullptr);
+  TfLiteTensor* output =
+      micro_context->AllocateTempOutputTensor(node, kOutputTensor);
+  TF_LITE_ENSURE(context, output != nullptr);
+  output->type = input->type;
+  if (IsDynamicTensor(axis)) {
+    MicroPrintf("DynamicTensor is not yet supported by Expand_Dims.");
+    return kTfLiteError;
+  }
+  TF_LITE_ENSURE_OK(context, VerifyTensorDim(context, input, axis, output));
+
+  micro_context->DeallocateTempTfLiteTensor(input);
+  micro_context->DeallocateTempTfLiteTensor(axis);
+  micro_context->DeallocateTempTfLiteTensor(output);
+  return kTfLiteOk;
+}
+
+template <typename T>
+void memCopyN(T* out, const T* in, const int num_elements) {
+#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
+  memcpy(out, in, num_elements * sizeof(T));
+#else // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
+  for (int i = 0; i < num_elements; ++i) {
+    out[i] = in[i];
+  }
+#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
+}
+
+TfLiteStatus ExpandDimsEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteEvalTensor* input =
+      tflite::micro::GetEvalInput(context, node, kInputTensor);
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+  const int flat_size = ElementCount(*input->dims);
+
+  switch (input->type) {
+    case kTfLiteFloat32: {
+      memCopyN(tflite::micro::GetTensorData<float>(output),
+               tflite::micro::GetTensorData<float>(input), flat_size);
+    } break;
+    case kTfLiteInt16: {
+      memCopyN(tflite::micro::GetTensorData<int16_t>(output),
+               tflite::micro::GetTensorData<int16_t>(input), flat_size);
+    } break;
+    case kTfLiteInt8: {
+      memCopyN(tflite::micro::GetTensorData<int8_t>(output),
+               tflite::micro::GetTensorData<int8_t>(input), flat_size);
+    } break;
+    default:
+      MicroPrintf(
+          "Expand_Dims only currently supports int8, int16 and float32, got "
+          "%d.",
+          input->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+}  // namespace
+
+TFLMRegistration Register_EXPAND_DIMS() {
+  return tflite::micro::RegisterOp(nullptr, ExpandDimsPrepare, ExpandDimsEval);
+}
+
+}  // namespace tflite