| // 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/leaky_relu.h" |
| |
| #include "tensorflow/lite/c/common.h" |
| #include "tensorflow/lite/micro/kernels/kernel_util.h" |
| #include "tensorflow/lite/micro/kernels/leaky_relu.h" |
| #include "tensorflow/lite/micro/micro_log.h" |
| #include "tflm/opt/opt.h" |
| |
| namespace tflite { |
| |
| namespace { |
| void* LeakyReluInit(TfLiteContext* context, const char* buffer, size_t length) { |
| TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr); |
| return context->AllocatePersistentBuffer(context, sizeof(LeakyReluOpData)); |
| } |
| |
| TfLiteStatus LeakyReluEval(TfLiteContext* context, TfLiteNode* node) { |
| const TfLiteEvalTensor* input = |
| tflite::micro::GetEvalInput(context, node, kInputTensor); |
| TfLiteEvalTensor* output = |
| tflite::micro::GetEvalOutput(context, node, kOutputTensor); |
| const LeakyReluOpData& data = *static_cast<LeakyReluOpData*>(node->user_data); |
| |
| // Kelvin's vector ISA is used to implement Int8 and Int16. |
| // Float32 uses the reference op. |
| switch (input->type) { |
| case kTfLiteFloat32: { |
| LeakyReluParams op_params = {}; |
| const auto* params = |
| static_cast<TfLiteLeakyReluParams*>(node->builtin_data); |
| |
| op_params.alpha = params->alpha; |
| reference_ops::LeakyRelu(op_params, tflite::micro::GetTensorShape(input), |
| tflite::micro::GetTensorData<float>(input), |
| tflite::micro::GetTensorShape(output), |
| tflite::micro::GetTensorData<float>(output)); |
| return kTfLiteOk; |
| } break; |
| case kTfLiteInt8: { |
| kelvin::opt::leaky_relu_s8( |
| tflite::micro::GetTensorData<int8_t>(input), |
| tflite::micro::GetTensorData<int8_t>(output), |
| MatchingFlatSize(tflite::micro::GetTensorShape(input), |
| tflite::micro::GetTensorShape(output)), |
| data.input_zero_point, data.output_zero_point, |
| data.output_multiplier_alpha, data.output_shift_alpha, |
| data.output_multiplier_identity, data.output_shift_identity); |
| return kTfLiteOk; |
| } break; |
| case kTfLiteInt16: { |
| kelvin::opt::leaky_relu_s16( |
| tflite::micro::GetTensorData<int16_t>(input), |
| tflite::micro::GetTensorData<int16_t>(output), |
| MatchingFlatSize(tflite::micro::GetTensorShape(input), |
| tflite::micro::GetTensorShape(output)), |
| data.input_zero_point, data.output_zero_point, |
| data.output_multiplier_alpha, data.output_shift_alpha, |
| data.output_multiplier_identity, data.output_shift_identity); |
| return kTfLiteOk; |
| } break; |
| default: |
| MicroPrintf( |
| "Only float32, int8, int16 are supported by LEAKY_RELU, got %s.", |
| TfLiteTypeGetName(input->type)); |
| return kTfLiteError; |
| } |
| |
| return kTfLiteError; |
| } |
| |
| } // namespace |
| |
| TFLMRegistration Register_LEAKY_RELU() { |
| return tflite::micro::RegisterOp(LeakyReluInit, LeakyReluPrepare, |
| LeakyReluEval); |
| } |
| |
| } // namespace tflite |