Add kernel for Kelvin - Add kernel for Kelvin, which uses optimized implementations for Int8, Int16, and Int32 (all non-broadcast). Change-Id: I1c0dfcb58c3f9fa9dfe213fff2ec3746ca348df4
diff --git a/tensorflow/lite/micro/kernels/add_test.cc b/tensorflow/lite/micro/kernels/add_test.cc index 6e8b40c..bdf0224 100644 --- a/tensorflow/lite/micro/kernels/add_test.cc +++ b/tensorflow/lite/micro/kernels/add_test.cc
@@ -256,16 +256,42 @@ TF_LITE_MICRO_TEST(QuantizedAddNoActivationInt8) { const float scales[] = {0.25, 0.5, 1.0}; const int zero_points[] = {-10, 4, 13}; - int inout_shape[] = {4, 1, 2, 2, 1}; - const float input1_values[] = {-2.01, -1.01, -0.01, 0.98}; - const float input2_values[] = {1.01, 1.99, 2.99, 4.02}; - const float golden_values[] = {-1, 1, 3, 5}; + int inout_shape[] = {4, 1, 7, 6, 1}; + // clang-format off + const float input1_values[] = { + -2.01, -1.01, -0.01, 0.98, -2.01, -1.01, + -0.01, 0.98, -2.01, -1.01, -0.01, 0.98, + -2.01, -1.01, -0.01, 0.98, -2.01, -1.01, + -0.01, 0.98, -2.01, -1.01, -0.01, 0.98, + -2.01, -1.01, -0.01, 0.98, -2.01, -1.01, + -0.01, 0.98, -2.01, -1.01, -0.01, 0.98, + -2.0, 0.2, 0.7, 0.8, 1.1, 2.0 + }; + const float input2_values[] = { + 1.01, 1.99, 2.99, 4.02, 1.01, 1.99, + 2.99, 4.02, 1.01, 1.99, 2.99, 4.02, + 1.01, 1.99, 2.99, 4.02, 1.01, 1.99, + 2.99, 4.02, 1.01, 1.99, 2.99, 4.02, + 1.01, 1.99, 2.99, 4.02, 1.01, 1.99, + 2.99, 4.02, 1.01, 1.99, 2.99, 4.02, + 0.1, 0.2, 0.3, 0.5, 1.1, 0.1 + }; + const float golden_values[] = { + -1, 1, 3, 5, -1, 1, + 3, 5, -1, 1, 3, 5, + -1, 1, 3, 5, -1, 1, + 3, 5, -1, 1, 3, 5, + -1, 1, 3, 5, -1, 1, + 3, 5, -1, 1, 3, 5, + -1.9, 0.4, 1.0, 1.3, 2.2, 2.1 + }; + // clang-format on - constexpr int kOutputDimsCount = 4; - int8_t input1_quantized[kOutputDimsCount]; - int8_t input2_quantized[kOutputDimsCount]; - int8_t golden_quantized[kOutputDimsCount]; - int8_t output[kOutputDimsCount]; + constexpr int kOutputDimsCount = 42; + int8_t input1_quantized[kOutputDimsCount] __attribute__((aligned(64))); + int8_t input2_quantized[kOutputDimsCount] __attribute__((aligned(64))); + int8_t golden_quantized[kOutputDimsCount] __attribute__((aligned(64))); + int8_t output[kOutputDimsCount] __attribute__((aligned(64))); tflite::testing::TestAddQuantized( inout_shape, input1_values, input1_quantized, scales[0], zero_points[0],
diff --git a/tensorflow/lite/micro/kernels/kelvin/add.cc b/tensorflow/lite/micro/kernels/kelvin/add.cc new file mode 100644 index 0000000..b64c71f --- /dev/null +++ b/tensorflow/lite/micro/kernels/kelvin/add.cc
@@ -0,0 +1,181 @@ +// Copyright 2023 Google LLC +// Licensed under the Apache License, Version 2.0, see LICENSE for details. +// SPDX-License-Identifier: Apache-2.0 + +#include "tensorflow/lite/kernels/internal/reference/add.h" + +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h" +#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h" +#include "tensorflow/lite/micro/kernels/add.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/micro_log.h" +#include "tflm/opt/opt.h" + +namespace tflite { +namespace { + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr); + return context->AllocatePersistentBuffer(context, sizeof(OpDataAdd)); +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + return AddPrepare(context, node); +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + TFLITE_DCHECK(node->user_data != nullptr); + const OpDataAdd* data = static_cast<const OpDataAdd*>(node->user_data); + + const TfLiteEvalTensor* input1 = + tflite::micro::GetEvalInput(context, node, kAddInputTensor1); + const TfLiteEvalTensor* input2 = + tflite::micro::GetEvalInput(context, node, kAddInputTensor2); + TfLiteEvalTensor* output = + tflite::micro::GetEvalOutput(context, node, kAddOutputTensor); + + if (output->type == kTfLiteFloat32) { + tflite::ArithmeticParams op_params; + SetActivationParams(data->output_activation_min_f32, + data->output_activation_max_f32, &op_params); + if (data->requires_broadcast) { + reference_ops::BroadcastAdd4DSlow( + op_params, tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorData<float>(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorData<float>(input2), + tflite::micro::GetTensorShape(output), + tflite::micro::GetTensorData<float>(output)); + } else { + reference_ops::Add(op_params, tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorData<float>(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorData<float>(input2), + tflite::micro::GetTensorShape(output), + tflite::micro::GetTensorData<float>(output)); + } + } else if (output->type == kTfLiteInt32) { + tflite::ArithmeticParams op_params; + SetActivationParams(std::numeric_limits<int32_t>::lowest(), + std::numeric_limits<int32_t>::max(), &op_params); + if (data->requires_broadcast) { + reference_ops::BroadcastAdd4DSlow( + op_params, tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorData<int32_t>(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorData<int32_t>(input2), + tflite::micro::GetTensorShape(output), + tflite::micro::GetTensorData<int32_t>(output)); + } else { + kelvin::opt::elementwise_add_s32( + tflite::micro::GetTensorData<int32_t>(input1), + tflite::micro::GetTensorData<int32_t>(input2), + tflite::micro::GetTensorData<int32_t>(output), + op_params.quantized_activation_min, + op_params.quantized_activation_max, + MatchingElementsSize(tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorShape(output))); + } + return kTfLiteOk; + } else if (output->type == kTfLiteInt16) { + tflite::ArithmeticParams op_params; + op_params.left_shift = data->left_shift; + op_params.input1_offset = data->input1_offset; + op_params.input1_multiplier = data->input1_multiplier; + op_params.input1_shift = data->input1_shift; + op_params.input2_offset = data->input2_offset; + op_params.input2_multiplier = data->input2_multiplier; + op_params.input2_shift = data->input2_shift; + op_params.output_offset = data->output_offset; + op_params.output_multiplier = data->output_multiplier; + op_params.output_shift = data->output_shift; + SetActivationParams(data->output_activation_min, + data->output_activation_max, &op_params); + + bool need_broadcast = reference_ops::ProcessBroadcastShapes( + tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorShape(input2), &op_params); + + if (need_broadcast) { + reference_ops::BroadcastAdd4DSlow( + op_params, tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorData<int16_t>(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorData<int16_t>(input2), + tflite::micro::GetTensorShape(output), + tflite::micro::GetTensorData<int16_t>(output)); + } else { + kelvin::opt::elementwise_add_s16( + tflite::micro::GetTensorData<int16_t>(input1), + tflite::micro::GetTensorData<int16_t>(input2), + op_params.input1_offset, op_params.input1_multiplier, + op_params.input1_shift, op_params.input2_offset, + op_params.input2_multiplier, op_params.input2_shift, + op_params.left_shift, tflite::micro::GetTensorData<int16_t>(output), + op_params.output_offset, op_params.output_multiplier, + op_params.output_shift, op_params.quantized_activation_min, + op_params.quantized_activation_max, + MatchingElementsSize(tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorShape(output))); + } + return kTfLiteOk; + } else if (output->type == kTfLiteInt8) { + tflite::ArithmeticParams op_params; + op_params.left_shift = data->left_shift; + op_params.input1_offset = data->input1_offset; + op_params.input1_multiplier = data->input1_multiplier; + op_params.input1_shift = data->input1_shift; + op_params.input2_offset = data->input2_offset; + op_params.input2_multiplier = data->input2_multiplier; + op_params.input2_shift = data->input2_shift; + op_params.output_offset = data->output_offset; + op_params.output_multiplier = data->output_multiplier; + op_params.output_shift = data->output_shift; + SetActivationParams(data->output_activation_min, + data->output_activation_max, &op_params); + + bool need_broadcast = reference_ops::ProcessBroadcastShapes( + tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorShape(input2), &op_params); + + if (need_broadcast) { + reference_integer_ops::BroadcastAdd4DSlow( + op_params, tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorData<int8_t>(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorData<int8_t>(input2), + tflite::micro::GetTensorShape(output), + tflite::micro::GetTensorData<int8_t>(output)); + } else { + kelvin::opt::elementwise_add_s8( + tflite::micro::GetTensorData<int8_t>(input1), + tflite::micro::GetTensorData<int8_t>(input2), op_params.input1_offset, + op_params.input1_multiplier, op_params.input1_shift, + op_params.input2_offset, op_params.input2_multiplier, + op_params.input2_shift, op_params.left_shift, + tflite::micro::GetTensorData<int8_t>(output), op_params.output_offset, + op_params.output_multiplier, op_params.output_shift, + op_params.quantized_activation_min, + op_params.quantized_activation_max, + MatchingElementsSize(tflite::micro::GetTensorShape(input1), + tflite::micro::GetTensorShape(input2), + tflite::micro::GetTensorShape(output))); + } + } else { + MicroPrintf("Unsupported output type: %s", TfLiteTypeGetName(output->type)); + return kTfLiteError; + } + + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_ADD() { + return tflite::micro::RegisterOp(Init, Prepare, Eval); +} + +} // namespace tflite