mirror of
https://github.com/vee1e/tflite-micro.git
synced 2026-09-03 18:58:00 +00:00
2b7a4ad621
@tensorflow/micro Eval method now uses TfLiteEvalTensor. The `has_low_rank_input_condition` variable has been removed from OpData as SELECT_V2 does not use this variable. All code outside of the registration has been placed into the anonymous namespace. bug=fixes #2014
196 lines
7.2 KiB
C++
196 lines
7.2 KiB
C++
/* Copyright 2023 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/kernels/internal/reference/select.h"
|
|
|
|
#include <stddef.h>
|
|
#include <stdint.h>
|
|
|
|
#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"
|
|
|
|
namespace tflite {
|
|
namespace {
|
|
|
|
constexpr int kInputTensorCondition = 0;
|
|
constexpr int kInputTensorX = 1;
|
|
constexpr int kInputTensorY = 2;
|
|
constexpr int kOutputTensor = 0;
|
|
|
|
struct OpData {
|
|
bool requires_broadcast;
|
|
};
|
|
|
|
void* SelectInit(TfLiteContext* context, const char* buffer, size_t length) {
|
|
TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
|
|
auto* data = static_cast<OpData*>(
|
|
context->AllocatePersistentBuffer(context, sizeof(OpData)));
|
|
data->requires_broadcast = false;
|
|
return data;
|
|
}
|
|
|
|
TfLiteStatus CheckBroadcastShape(TfLiteContext* context,
|
|
const TfLiteTensor* input1,
|
|
const TfLiteTensor* input2,
|
|
const TfLiteTensor* input3,
|
|
const TfLiteIntArray* output_shape) {
|
|
const int dims1 = NumDimensions(input1);
|
|
const int dims2 = NumDimensions(input2);
|
|
const int dims3 = NumDimensions(input3);
|
|
const int out_dims = std::max(std::max(dims1, dims2), dims3);
|
|
TF_LITE_ENSURE_EQ(context, out_dims, output_shape->size);
|
|
|
|
for (int i = 0; i < out_dims; ++i) {
|
|
const int d1 = i >= dims1 ? 1 : SizeOfDimension(input1, dims1 - i - 1);
|
|
const int d2 = i >= dims2 ? 1 : SizeOfDimension(input2, dims2 - i - 1);
|
|
const int d3 = i >= dims3 ? 1 : SizeOfDimension(input3, dims3 - i - 1);
|
|
const int min_value = std::min(std::min(d1, d2), d3);
|
|
int max_value = std::max(std::max(d1, d2), d3);
|
|
// If one dimension is 0, others must be 0 or 1.
|
|
if (min_value == 0) max_value = 0;
|
|
if (!(d1 == 1 || d1 == max_value) || !(d2 == 1 || d2 == max_value) ||
|
|
!(d3 == 1 || d3 == max_value)) {
|
|
MicroPrintf("Given shapes are not broadcastable.");
|
|
return kTfLiteError;
|
|
}
|
|
TF_LITE_ENSURE_EQ(context, output_shape->data[out_dims - i - 1], max_value);
|
|
}
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) {
|
|
OpData* data = reinterpret_cast<OpData*>(node->user_data);
|
|
|
|
TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
|
|
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
|
|
|
MicroContext* micro_context = GetMicroContext(context);
|
|
TfLiteTensor* input_condition =
|
|
micro_context->AllocateTempInputTensor(node, kInputTensorCondition);
|
|
|
|
TfLiteTensor* input_x =
|
|
micro_context->AllocateTempInputTensor(node, kInputTensorX);
|
|
|
|
TfLiteTensor* input_y =
|
|
micro_context->AllocateTempInputTensor(node, kInputTensorY);
|
|
|
|
TfLiteTensor* output =
|
|
micro_context->AllocateTempOutputTensor(node, kOutputTensor);
|
|
|
|
// Input must be bool.
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input_condition->type, kTfLiteBool);
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input_x->type, input_y->type);
|
|
output->type = input_x->type;
|
|
|
|
// Respect the original output shape when there are mixed shapes to represent
|
|
// a scalar data.
|
|
bool possible_mixed_scaler =
|
|
GetTensorShape(input_condition).FlatSize() == 1 &&
|
|
GetTensorShape(input_x).FlatSize() == 1 &&
|
|
GetTensorShape(input_y).FlatSize() == 1 &&
|
|
GetTensorShape(output).FlatSize() == 1;
|
|
|
|
bool same_shape = HaveSameShapes(input_condition, input_x) &&
|
|
HaveSameShapes(input_x, input_y);
|
|
if (!same_shape && !possible_mixed_scaler) {
|
|
TF_LITE_ENSURE_OK(
|
|
context, CheckBroadcastShape(context, input_condition, input_x, input_y,
|
|
output->dims));
|
|
data->requires_broadcast = true;
|
|
}
|
|
|
|
micro_context->DeallocateTempTfLiteTensor(input_condition);
|
|
micro_context->DeallocateTempTfLiteTensor(input_x);
|
|
micro_context->DeallocateTempTfLiteTensor(input_y);
|
|
micro_context->DeallocateTempTfLiteTensor(output);
|
|
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
template <typename T>
|
|
void CallSelect(const TfLiteEvalTensor* input_condition,
|
|
const TfLiteEvalTensor* input_x,
|
|
const TfLiteEvalTensor* input_y, TfLiteEvalTensor* output,
|
|
bool need_broadcast) {
|
|
using Func = decltype(reference_ops::Select<bool, T>)*;
|
|
Func select_func;
|
|
if (need_broadcast) {
|
|
select_func = reference_ops::BroadcastSelect5DSlow<bool, T>;
|
|
} else {
|
|
select_func = reference_ops::Select<bool, T>;
|
|
}
|
|
|
|
select_func(tflite::micro::GetTensorShape(input_condition),
|
|
tflite::micro::GetTensorData<bool>(input_condition),
|
|
tflite::micro::GetTensorShape(input_x),
|
|
tflite::micro::GetTensorData<T>(input_x),
|
|
tflite::micro::GetTensorShape(input_y),
|
|
tflite::micro::GetTensorData<T>(input_y),
|
|
tflite::micro::GetTensorShape(output),
|
|
tflite::micro::GetTensorData<T>(output));
|
|
}
|
|
|
|
TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) {
|
|
OpData* data = static_cast<OpData*>(node->user_data);
|
|
|
|
const TfLiteEvalTensor* input_condition =
|
|
tflite::micro::GetEvalInput(context, node, kInputTensorCondition);
|
|
|
|
const TfLiteEvalTensor* input_x =
|
|
tflite::micro::GetEvalInput(context, node, kInputTensorX);
|
|
|
|
const TfLiteEvalTensor* input_y =
|
|
tflite::micro::GetEvalInput(context, node, kInputTensorY);
|
|
|
|
TfLiteEvalTensor* output =
|
|
tflite::micro::GetEvalOutput(context, node, kOutputTensor);
|
|
|
|
switch (input_x->type) {
|
|
case kTfLiteFloat32:
|
|
CallSelect<float>(input_condition, input_x, input_y, output,
|
|
data->requires_broadcast);
|
|
break;
|
|
case kTfLiteInt8:
|
|
CallSelect<int8_t>(input_condition, input_x, input_y, output,
|
|
data->requires_broadcast);
|
|
break;
|
|
case kTfLiteInt16:
|
|
CallSelect<int16_t>(input_condition, input_x, input_y, output,
|
|
data->requires_broadcast);
|
|
break;
|
|
default:
|
|
MicroPrintf("Does not support type other than %s, but got %s",
|
|
"int8|int16|float32", TfLiteTypeGetName(input_x->type));
|
|
return kTfLiteError;
|
|
}
|
|
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// SelectV2 op selects values of 'x' if the corresponding value of 'condition'
|
|
// is true or the value of 'y' if false. There are valid condition input sizes:
|
|
//
|
|
// 1. Either the same shape (in which case the select is elementwise), or
|
|
// 2. Broadcastable shapes between 'condition', 'x' and 'y'.
|
|
TFLMRegistration Register_SELECT_V2() {
|
|
return tflite::micro::RegisterOp(tflite::SelectInit, tflite::SelectPrepare,
|
|
tflite::SelectEval);
|
|
}
|
|
|
|
} // namespace tflite
|