271 lines
12 KiB
C++
271 lines
12 KiB
C++
/* Copyright 2022 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/c/common.h"
|
|
#include "tensorflow/lite/kernels/internal/quantization_util.h"
|
|
#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
|
|
#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
|
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
|
#include "tensorflow/lite/micro/kernels/kernel_util.h"
|
|
#include "tensorflow/lite/micro/micro_context.h"
|
|
#include "tensorflow/lite/micro/micro_log.h"
|
|
|
|
namespace tflite {
|
|
namespace {
|
|
constexpr int kInputTensor1 = 0;
|
|
constexpr int kInputTensor2 = 1;
|
|
constexpr int kOutputTensor = 0;
|
|
|
|
struct OpData {
|
|
bool requires_broadcast;
|
|
ArithmeticParams arithmetic_params;
|
|
};
|
|
|
|
template <typename T>
|
|
T SquaredDifference(T input1, T input2) {
|
|
const T difference = input1 - input2;
|
|
return difference * difference;
|
|
}
|
|
|
|
void* SquaredDifferenceInit(TfLiteContext* context, const char* buffer,
|
|
size_t length) {
|
|
TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
|
|
return context->AllocatePersistentBuffer(context, sizeof(OpData));
|
|
}
|
|
|
|
void PrepareQuantized(
|
|
const TfLiteQuantizationParams& input1_quantization_params,
|
|
const TfLiteQuantizationParams& input2_quantization_params,
|
|
const TfLiteQuantizationParams& output_quantization_params,
|
|
const int left_shift, const int32_t quantized_activation_min,
|
|
const int32_t quantized_activation_max, OpData* data) {
|
|
data->arithmetic_params.input1_offset =
|
|
-input1_quantization_params.zero_point;
|
|
data->arithmetic_params.input2_offset =
|
|
-input2_quantization_params.zero_point;
|
|
data->arithmetic_params.output_offset = output_quantization_params.zero_point;
|
|
data->arithmetic_params.left_shift = left_shift;
|
|
const double twice_max_input_scale =
|
|
2.0 * static_cast<double>(std::max(input1_quantization_params.scale,
|
|
input2_quantization_params.scale));
|
|
const double real_input1_multiplier =
|
|
static_cast<double>(input1_quantization_params.scale) /
|
|
twice_max_input_scale;
|
|
double real_input2_multiplier =
|
|
static_cast<double>(input2_quantization_params.scale) /
|
|
twice_max_input_scale;
|
|
const double real_output_multiplier =
|
|
(twice_max_input_scale * twice_max_input_scale) /
|
|
static_cast<double>((1 << data->arithmetic_params.left_shift * 2) *
|
|
output_quantization_params.scale);
|
|
QuantizeMultiplierSmallerThanOneExp(
|
|
real_input1_multiplier, &data->arithmetic_params.input1_multiplier,
|
|
&data->arithmetic_params.input1_shift);
|
|
QuantizeMultiplierSmallerThanOneExp(
|
|
real_input2_multiplier, &data->arithmetic_params.input2_multiplier,
|
|
&data->arithmetic_params.input2_shift);
|
|
QuantizeMultiplier(real_output_multiplier,
|
|
&data->arithmetic_params.output_multiplier,
|
|
&data->arithmetic_params.output_shift);
|
|
data->arithmetic_params.quantized_activation_min = quantized_activation_min;
|
|
data->arithmetic_params.quantized_activation_max = quantized_activation_max;
|
|
}
|
|
|
|
TfLiteStatus SquaredDifferencePrepare(TfLiteContext* context,
|
|
TfLiteNode* node) {
|
|
TFLITE_DCHECK(node->user_data != nullptr);
|
|
OpData* data = reinterpret_cast<OpData*>(node->user_data);
|
|
data->requires_broadcast = false;
|
|
|
|
TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
|
|
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
|
|
|
MicroContext* micro_context = GetMicroContext(context);
|
|
|
|
TfLiteTensor* input1 =
|
|
micro_context->AllocateTempInputTensor(node, kInputTensor1);
|
|
TF_LITE_ENSURE(context, input1 != nullptr);
|
|
TfLiteTensor* input2 =
|
|
micro_context->AllocateTempInputTensor(node, kInputTensor2);
|
|
TF_LITE_ENSURE(context, input2 != nullptr);
|
|
TfLiteTensor* output =
|
|
micro_context->AllocateTempOutputTensor(node, kOutputTensor);
|
|
TF_LITE_ENSURE(context, output != nullptr);
|
|
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
|
|
output->type = input2->type;
|
|
|
|
const TfLiteQuantizationParams& input1_quantization_params = input1->params;
|
|
const TfLiteQuantizationParams& input2_quantization_params = input2->params;
|
|
const TfLiteQuantizationParams& output_quantization_params = output->params;
|
|
if (input1->type == kTfLiteInt8) {
|
|
const int32_t integer_type_min = std::numeric_limits<int8_t>::min();
|
|
const int32_t integer_type_max = std::numeric_limits<int8_t>::max();
|
|
TF_LITE_ENSURE(context,
|
|
input1_quantization_params.zero_point >= integer_type_min);
|
|
TF_LITE_ENSURE(context,
|
|
input1_quantization_params.zero_point <= integer_type_max);
|
|
TF_LITE_ENSURE(context,
|
|
input2_quantization_params.zero_point >= integer_type_min);
|
|
TF_LITE_ENSURE(context,
|
|
input2_quantization_params.zero_point <= integer_type_max);
|
|
TF_LITE_ENSURE(context,
|
|
output_quantization_params.zero_point >= integer_type_min);
|
|
TF_LITE_ENSURE(context,
|
|
output_quantization_params.zero_point <= integer_type_max);
|
|
// leftshift = 7 is selected so that maximum shifted result 255^2 * (1 << (7
|
|
// * 2 )) does not overflow signed 32-bit integer
|
|
PrepareQuantized(input1_quantization_params, input2_quantization_params,
|
|
output_quantization_params, /*left_shift=*/7,
|
|
/*quantized_activation_min*/ integer_type_min,
|
|
/*quantized_activation_max*/ integer_type_max, data);
|
|
} else if (input1->type == kTfLiteInt16) {
|
|
const int32_t integer_type_min = std::numeric_limits<int16_t>::min();
|
|
const int32_t integer_type_max = std::numeric_limits<int16_t>::max();
|
|
TF_LITE_ENSURE(context, input1_quantization_params.zero_point == 0);
|
|
TF_LITE_ENSURE(context, input2_quantization_params.zero_point == 0);
|
|
TF_LITE_ENSURE(context, output_quantization_params.zero_point == 0);
|
|
|
|
// leftshift = 0 as number is already 16-bit. so that maximum shifted result
|
|
// 32767^2 * (1 << (0 * 2 ))
|
|
PrepareQuantized(input1_quantization_params, input2_quantization_params,
|
|
output_quantization_params, /*left_shift=*/0,
|
|
/*quantized_activation_min*/ integer_type_min,
|
|
/*quantized_activation_max*/ integer_type_max, data);
|
|
}
|
|
|
|
data->requires_broadcast = !HaveSameShapes(input1, input2);
|
|
|
|
micro_context->DeallocateTempTfLiteTensor(input1);
|
|
micro_context->DeallocateTempTfLiteTensor(input2);
|
|
micro_context->DeallocateTempTfLiteTensor(output);
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
template <typename T>
|
|
T SquaredDifference(T x, T y, const ArithmeticParams& params) {
|
|
const int32_t input1_val = params.input1_offset + x;
|
|
const int32_t input2_val = params.input2_offset + y;
|
|
const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
|
|
const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
|
|
const int32_t scaled_input1_val =
|
|
MultiplyByQuantizedMultiplierSmallerThanOneExp(
|
|
shifted_input1_val, params.input1_multiplier, params.input1_shift);
|
|
const int32_t scaled_input2_val =
|
|
MultiplyByQuantizedMultiplierSmallerThanOneExp(
|
|
shifted_input2_val, params.input2_multiplier, params.input2_shift);
|
|
const int32_t raw_diff = scaled_input1_val - scaled_input2_val;
|
|
|
|
// Max of this is 32767^2 * (1 << 0), so won't overflow 32 bits.
|
|
const int32_t squared_raw_diff = raw_diff * raw_diff;
|
|
const int32_t raw_output =
|
|
MultiplyByQuantizedMultiplier(squared_raw_diff, params.output_multiplier,
|
|
params.output_shift) +
|
|
params.output_offset;
|
|
const int32_t clamped_output =
|
|
std::min(params.quantized_activation_max,
|
|
std::max(params.quantized_activation_min, raw_output));
|
|
return static_cast<T>(clamped_output);
|
|
}
|
|
|
|
template <typename T>
|
|
void EvalQuantizedSquaredDifference(TfLiteContext* context, TfLiteNode* node,
|
|
const OpData* data,
|
|
const TfLiteEvalTensor* input1,
|
|
const TfLiteEvalTensor* input2,
|
|
TfLiteEvalTensor* output) {
|
|
const auto* op_data = static_cast<const OpData*>(node->user_data);
|
|
if (data->requires_broadcast) {
|
|
reference_integer_ops::BroadcastBinaryFunction4DSlow(
|
|
op_data->arithmetic_params, tflite::micro::GetTensorShape(input1),
|
|
tflite::micro::GetTensorData<T>(input1),
|
|
tflite::micro::GetTensorShape(input2),
|
|
tflite::micro::GetTensorData<T>(input2),
|
|
tflite::micro::GetTensorShape(output),
|
|
tflite::micro::GetTensorData<T>(output),
|
|
reference_integer_ops::CheckArithmeticParams, SquaredDifference);
|
|
} else {
|
|
const int flat_size = tflite::micro::GetTensorShape(input1).FlatSize();
|
|
reference_integer_ops::ElementWise(
|
|
flat_size, op_data->arithmetic_params,
|
|
tflite::micro::GetTensorData<T>(input1),
|
|
tflite::micro::GetTensorData<T>(input2),
|
|
tflite::micro::GetTensorData<T>(output),
|
|
reference_integer_ops::CheckArithmeticParams, SquaredDifference);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void EvalSquaredDifference(TfLiteContext* context, TfLiteNode* node,
|
|
const OpData* data, const TfLiteEvalTensor* input1,
|
|
const TfLiteEvalTensor* input2,
|
|
TfLiteEvalTensor* output) {
|
|
if (data->requires_broadcast) {
|
|
reference_ops::BroadcastBinaryFunction4DSlow<T, T, T>(
|
|
tflite::micro::GetTensorShape(input1),
|
|
tflite::micro::GetTensorData<T>(input1),
|
|
tflite::micro::GetTensorShape(input2),
|
|
tflite::micro::GetTensorData<T>(input2),
|
|
tflite::micro::GetTensorShape(output),
|
|
tflite::micro::GetTensorData<T>(output), SquaredDifference<T>);
|
|
} else {
|
|
reference_ops::BinaryFunction<T, T, T>(
|
|
tflite::micro::GetTensorShape(input1),
|
|
tflite::micro::GetTensorData<T>(input1),
|
|
tflite::micro::GetTensorShape(input2),
|
|
tflite::micro::GetTensorData<T>(input2),
|
|
tflite::micro::GetTensorShape(output),
|
|
tflite::micro::GetTensorData<T>(output), SquaredDifference<T>);
|
|
}
|
|
}
|
|
|
|
TfLiteStatus SquaredDifferenceEval(TfLiteContext* context, TfLiteNode* node) {
|
|
OpData* data = reinterpret_cast<OpData*>(node->user_data);
|
|
|
|
const TfLiteEvalTensor* input1 =
|
|
tflite::micro::GetEvalInput(context, node, kInputTensor1);
|
|
const TfLiteEvalTensor* input2 =
|
|
tflite::micro::GetEvalInput(context, node, kInputTensor2);
|
|
TfLiteEvalTensor* output =
|
|
tflite::micro::GetEvalOutput(context, node, kOutputTensor);
|
|
|
|
if (output->type == kTfLiteFloat32) {
|
|
EvalSquaredDifference<float>(context, node, data, input1, input2, output);
|
|
} else if (output->type == kTfLiteInt32) {
|
|
EvalSquaredDifference<int32_t>(context, node, data, input1, input2, output);
|
|
} else if (output->type == kTfLiteInt8) {
|
|
EvalQuantizedSquaredDifference<int8_t>(context, node, data, input1, input2,
|
|
output);
|
|
} else if (output->type == kTfLiteInt16) {
|
|
EvalQuantizedSquaredDifference<int16_t>(context, node, data, input1, input2,
|
|
output);
|
|
} else {
|
|
MicroPrintf(
|
|
"SquaredDifference only supports FLOAT32, INT32 , INT16 and INT8 now, "
|
|
"got %d.",
|
|
output->type);
|
|
return kTfLiteError;
|
|
}
|
|
|
|
return kTfLiteOk;
|
|
}
|
|
} // namespace
|
|
|
|
TFLMRegistration Register_SQUARED_DIFFERENCE() {
|
|
return tflite::micro::RegisterOp(
|
|
SquaredDifferenceInit, SquaredDifferencePrepare, SquaredDifferenceEval);
|
|
}
|
|
|
|
} // namespace tflite
|