Optimized Kelvin kernel for Reshape - Reshape kernel that uses a Kelvin-specific memcpy implementation - Bazel changes to select Kelvin kernels if the build configuration is correct Change-Id: I21914f18c87b505f9fee765bfa96dd340b0b8c96

diff --git a/tensorflow/extra_rules.bzl b/tensorflow/extra_rules.bzl
index 9e20eda..c0a39dc 100644
--- a/tensorflow/extra_rules.bzl
+++ b/tensorflow/extra_rules.bzl

@@ -28,3 +28,7 @@
 def xtensa_vision_p6_config():
     """Config setting for all Vision P6 based cores."""
     return "//tensorflow/lite/micro/kernels:xtensa_vision_p6_default"
+
+def kelvin_config():
+    """Config setting for Kelvin-based cores."""
+    return "//tensorflow/lite/micro/kernels:kelvin_default"

diff --git a/tensorflow/lite/micro/kernels/BUILD b/tensorflow/lite/micro/kernels/BUILD
index 11ece2d..94a6d5b 100644
--- a/tensorflow/lite/micro/kernels/BUILD
+++ b/tensorflow/lite/micro/kernels/BUILD

@@ -6,6 +6,7 @@
     "xtensa_hifi_3z_config",
     "xtensa_hifi_5_config",
     "xtensa_vision_p6_config",
+    "kelvin_config",
 )
 
 package(
@@ -31,6 +32,10 @@
     packages = tflm_kernel_friends(),
 )
 
+exports_files(
+    glob(["*_test.cc"])
+)
+
 ####################################
 # C++ libraries
 ####################################
@@ -175,6 +180,9 @@
     "-DVISION_P6=1",
 ]
 
+KELVIN_COPTS = [
+]
+
 tflm_kernel_cc_library(
     name = "micro_ops",
     srcs = [
@@ -310,6 +318,7 @@
         xtensa_hifi_3z_config(): glob(["xtensa/**/*.h"]),
         xtensa_hifi_5_config(): glob(["xtensa/**/*.h"]),
         xtensa_vision_p6_config(): glob(["xtensa/**/*.h"]),
+        kelvin_config(): glob(["kelvin/**/*.h"]),
         "//conditions:default": [],
     }),
     accelerated_srcs = {
@@ -317,12 +326,14 @@
         xtensa_hifi_3z_config(): glob(["xtensa/**/*.cc"]),
         xtensa_hifi_5_config(): glob(["xtensa/**/*.cc"]),
         xtensa_vision_p6_config(): glob(["xtensa/**/*.cc"]),
+        kelvin_config(): glob(["kelvin/**/*.cc"]),
     },
     copts = micro_copts() + select({
         xtensa_fusion_f1_config(): HIFI4_COPTS,
         xtensa_hifi_3z_config(): HIFI4_COPTS,
         xtensa_hifi_5_config(): HIFI5_COPTS,
         xtensa_vision_p6_config(): VP6_COPTS,
+        kelvin_config(): KELVIN_COPTS,
         "//conditions:default": [],
     }),
     visibility = [
@@ -358,6 +369,7 @@
         xtensa_hifi_3z_config(): ["//third_party/xtensa/nnlib_hifi4:nnlib_hifi4_lib"],
         xtensa_hifi_5_config(): ["//third_party/xtensa/nnlib_hifi5:nnlib_hifi5_lib"],
         xtensa_vision_p6_config(): ["//third_party/xtensa/xi_tflmlib_vision_p6:xi_tflmlib_vision_p6_lib"],
+        kelvin_config(): ["@kelvin_sw//tflm/opt:opt"],
         "//conditions:default": [],
     }),
 )
@@ -1453,3 +1465,10 @@
         "cpu": "P6_200528",
     },
 )
+
+config_setting(
+    name = "kelvin_default",
+    values = {
+        "platforms": "//platforms/riscv32:kelvin"
+    },
+)

diff --git a/tensorflow/lite/micro/kernels/kelvin/reshape.cc b/tensorflow/lite/micro/kernels/kelvin/reshape.cc
new file mode 100644
index 0000000..9c88385
--- /dev/null
+++ b/tensorflow/lite/micro/kernels/kelvin/reshape.cc

@@ -0,0 +1,100 @@
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/memory_helpers.h"
+#include "tensorflow/lite/micro/micro_utils.h"
+#include "tflm/opt/opt.h"
+
+namespace tflite {
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus ReshapeOutput(TfLiteContext* context, TfLiteNode* node) {
+  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);
+  // Tensorflow's Reshape allows one of the shape components to have the
+  // special -1 value, meaning it will be calculated automatically based on the
+  // input. Here we calculate what that dimension should be so that the number
+  // of output elements in the same as the number of input elements.
+  int num_input_elements = NumElements(input);
+  TfLiteIntArray* output_shape = output->dims;
+
+  if (NumInputs(node) == 1 &&  // Legacy scalar supported with params.
+      output_shape->size == 1 && output_shape->data[0] == 0) {
+    // Legacy tflite models use a shape parameter of [0] to indicate scalars,
+    // so adjust accordingly. TODO(b/111614235): Allow zero-sized buffers during
+    // toco conversion.
+    output_shape->size = 0;
+  }
+
+  int num_output_elements = 1;
+  int stretch_dim = -1;
+  for (int i = 0; i < output_shape->size; ++i) {
+    int value = output_shape->data[i];
+    if (value == -1) {
+      TF_LITE_ENSURE_EQ(context, stretch_dim, -1);
+      stretch_dim = i;
+    } else {
+      num_output_elements *= value;
+    }
+  }
+  if (stretch_dim != -1) {
+    TfLiteEvalTensor* output_eval =
+        tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+    TF_LITE_ENSURE_STATUS(tflite::micro::CreateWritableTensorDimsWithCopy(
+        context, output, output_eval));
+    output_shape = output->dims;  // output tensor dims were moved
+    output_shape->data[stretch_dim] = num_input_elements / num_output_elements;
+    num_output_elements *= output_shape->data[stretch_dim];
+  }
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements);
+
+  micro_context->DeallocateTempTfLiteTensor(input);
+  micro_context->DeallocateTempTfLiteTensor(output);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE(context, NumInputs(node) == 1 || NumInputs(node) == 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, ReshapeOutput(context, node), kTfLiteOk);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteEvalTensor* input =
+      tflite::micro::GetEvalInput(context, node, kInputTensor);
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+
+  // TODO(b/162522304): storing input bytes in OpData increases some models
+  // significantly, possibly due to alignment issues.
+  size_t input_bytes;
+  TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &input_bytes));
+  input_bytes *= ElementCount(*input->dims);
+
+  // Do nothing for in-place reshape.
+  if (input->data.raw != output->data.raw) {
+    // Otherwise perform reshape with copy.
+    kelvin::opt::memcpy(output->data.raw, input->data.raw, input_bytes);
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+TFLMRegistration Register_RESHAPE() {
+  return tflite::micro::RegisterOp(nullptr, Prepare, Eval);
+}
+
+}  // namespace tflite