wehub-resource-sync
337 lines
13 KiB
C++
337 lines
13 KiB
C++
/* Copyright 2018 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/slice.h"
|
|
|
|
#include <stdint.h>
|
|
|
|
#include <algorithm>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "Eigen/Core"
|
|
#include "tensorflow/lite/context_util.h"
|
|
#include "tensorflow/lite/core/c/common.h"
|
|
#include "tensorflow/lite/kernels/internal/compatibility.h"
|
|
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
|
|
#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
|
|
#include "tensorflow/lite/kernels/internal/tensor.h"
|
|
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
|
|
#include "tensorflow/lite/kernels/internal/types.h"
|
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
|
#include "tensorflow/lite/string_type.h"
|
|
|
|
namespace tflite {
|
|
namespace ops {
|
|
namespace builtin {
|
|
namespace slice {
|
|
|
|
enum KernelType {
|
|
kReference,
|
|
kGenericOptimized,
|
|
};
|
|
|
|
constexpr int kInputTensor = 0;
|
|
constexpr int kBeginTensor = 1;
|
|
constexpr int kSizeTensor = 2;
|
|
constexpr int kOutputTensor = 0;
|
|
|
|
// The optimized helper supports up to 5D. Higher-rank tensors use the
|
|
// runtime-rank reference path.
|
|
const int kMaxDim = 5;
|
|
|
|
template <typename T>
|
|
TfLiteStatus CalculateOutputShapeVector(TfLiteContext* context,
|
|
const TfLiteTensor* input,
|
|
const TfLiteTensor* begin,
|
|
const TfLiteTensor* size,
|
|
std::vector<int>* output_shape_vector) {
|
|
for (int idx = 0; idx < NumDimensions(input); ++idx) {
|
|
T size_value = GetTensorData<T>(size)[idx];
|
|
if (size_value < 0) {
|
|
if (size_value != -1) {
|
|
TF_LITE_KERNEL_LOG(context, "Invalid size.");
|
|
return kTfLiteError;
|
|
}
|
|
size_value = SizeOfDimension(input, idx) - GetTensorData<T>(begin)[idx];
|
|
} else {
|
|
if (SizeOfDimension(input, idx) <
|
|
GetTensorData<T>(begin)[idx] + size_value) {
|
|
TF_LITE_KERNEL_LOG(context, "Invalid begin and size.");
|
|
return kTfLiteError;
|
|
}
|
|
}
|
|
output_shape_vector->push_back(static_cast<int>(size_value));
|
|
}
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
template <typename T>
|
|
void GetBeginAndSizeVectors(int dimensions, const TfLiteTensor* begin,
|
|
const TfLiteTensor* size, std::vector<int>* begins,
|
|
std::vector<int>* sizes) {
|
|
for (int idx = 0; idx < dimensions; ++idx) {
|
|
begins->push_back(GetTensorData<T>(begin)[idx]);
|
|
sizes->push_back(GetTensorData<T>(size)[idx]);
|
|
}
|
|
}
|
|
|
|
TfLiteStatus ResizeOutputShape(TfLiteContext* context,
|
|
const TfLiteTensor* input,
|
|
const TfLiteTensor* begin,
|
|
const TfLiteTensor* size, TfLiteTensor* output) {
|
|
std::vector<int> output_shape_vector;
|
|
|
|
if (begin->type == kTfLiteInt32) {
|
|
TF_LITE_ENSURE_STATUS(CalculateOutputShapeVector<int32_t>(
|
|
context, input, begin, size, &output_shape_vector));
|
|
} else if (begin->type == kTfLiteInt64) {
|
|
TF_LITE_ENSURE_STATUS(CalculateOutputShapeVector<int64_t>(
|
|
context, input, begin, size, &output_shape_vector));
|
|
} else {
|
|
TF_LITE_KERNEL_LOG(context, "Type %d is currently not supported by Slice.",
|
|
begin->type);
|
|
return kTfLiteError;
|
|
}
|
|
|
|
TfLiteIntArray* output_shape =
|
|
TfLiteIntArrayCreate(output_shape_vector.size());
|
|
std::copy(output_shape_vector.begin(), output_shape_vector.end(),
|
|
output_shape->data);
|
|
return context->ResizeTensor(context, output, output_shape);
|
|
}
|
|
|
|
bool ShapeHasRank(const TfLiteIntArray* shape) {
|
|
// Note that we consider scalar as false here because there is
|
|
// no differentiation between scalar and dynamic properly supported.
|
|
if (shape == nullptr || shape->size == 0) return false;
|
|
return true;
|
|
}
|
|
|
|
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
|
TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
|
|
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
|
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
const TfLiteTensor* begin;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBeginTensor, &begin));
|
|
const TfLiteTensor* size;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSizeTensor, &size));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
|
|
// Ensure validity of input tensor and its dimension.
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
|
|
TF_LITE_ENSURE(context,
|
|
begin->type == kTfLiteInt32 || begin->type == kTfLiteInt64);
|
|
TF_LITE_ENSURE(context,
|
|
size->type == kTfLiteInt32 || size->type == kTfLiteInt64);
|
|
TF_LITE_ENSURE_EQ(context, NumDimensions(begin), 1);
|
|
TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1);
|
|
TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size));
|
|
// If the shape of output is fully specified then resize even if
|
|
// the input shape is not staticly defined.
|
|
if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims)) {
|
|
return kTfLiteOk;
|
|
}
|
|
// Postpone allocation of output if any of the indexing tensors is not
|
|
// constant, or the input tensor has dynamic dimension.
|
|
if (!(IsConstantOrPersistentTensor(begin) &&
|
|
IsConstantOrPersistentTensor(size)) ||
|
|
HasUnspecifiedDimension(input)) {
|
|
SetTensorToDynamic(output);
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
return ResizeOutputShape(context, input, begin, size, output);
|
|
}
|
|
|
|
template <KernelType kernel_type>
|
|
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
const TfLiteTensor* begin;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBeginTensor, &begin));
|
|
const TfLiteTensor* size;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSizeTensor, &size));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
|
|
if (IsDynamicTensor(output)) {
|
|
TF_LITE_ENSURE_OK(context,
|
|
ResizeOutputShape(context, input, begin, size, output));
|
|
}
|
|
|
|
const int input_dims = NumDimensions(input);
|
|
std::vector<int> begins;
|
|
begins.reserve(input_dims);
|
|
std::vector<int> sizes;
|
|
sizes.reserve(input_dims);
|
|
|
|
if (begin->type == kTfLiteInt32) {
|
|
GetBeginAndSizeVectors<int32_t>(NumDimensions(input), begin, size, &begins,
|
|
&sizes);
|
|
} else if (begin->type == kTfLiteInt64) {
|
|
GetBeginAndSizeVectors<int64_t>(NumDimensions(input), begin, size, &begins,
|
|
&sizes);
|
|
} else {
|
|
TF_LITE_KERNEL_LOG(context, "Type %d is currently not supported by Slice.",
|
|
begin->type);
|
|
return kTfLiteError;
|
|
}
|
|
|
|
if (input_dims > kMaxDim) {
|
|
if (input->type == kTfLiteString) {
|
|
reference_ops::Slice<string>(begins, GetTensorShape(input), input,
|
|
GetTensorShape(output), output);
|
|
return kTfLiteOk;
|
|
} else if (input->type == kTfLiteInt4) {
|
|
reference_ops::SliceInt4(begins, GetTensorShape(input), input,
|
|
GetTensorShape(output), output);
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
#define TF_LITE_SLICE_DYNAMIC(data_type) \
|
|
{ \
|
|
reference_ops::Slice<data_type>(begins, GetTensorShape(input), input, \
|
|
GetTensorShape(output), output); \
|
|
}
|
|
switch (TfLiteTypeGetSizeBits(output->type)) {
|
|
case 8:
|
|
TF_LITE_SLICE_DYNAMIC(int8_t);
|
|
break;
|
|
case 16:
|
|
TF_LITE_SLICE_DYNAMIC(int16_t);
|
|
break;
|
|
case 32:
|
|
TF_LITE_SLICE_DYNAMIC(int32_t);
|
|
break;
|
|
case 64:
|
|
TF_LITE_SLICE_DYNAMIC(int64_t);
|
|
break;
|
|
default:
|
|
TF_LITE_KERNEL_LOG(context,
|
|
"Type %d is currently not supported by Slice.",
|
|
input->type);
|
|
return kTfLiteError;
|
|
}
|
|
#undef TF_LITE_SLICE_DYNAMIC
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
std::vector<int> padded_begins;
|
|
padded_begins.reserve(kMaxDim);
|
|
std::vector<int> padded_sizes;
|
|
padded_sizes.reserve(kMaxDim);
|
|
for (int i = input_dims; i < kMaxDim; ++i) {
|
|
padded_begins.push_back(0);
|
|
padded_sizes.push_back(1);
|
|
}
|
|
padded_begins.insert(padded_begins.end(), begins.begin(), begins.end());
|
|
padded_sizes.insert(padded_sizes.end(), sizes.begin(), sizes.end());
|
|
|
|
#define TF_LITE_SLICE_INT4() \
|
|
{ \
|
|
TF_LITE_ENSURE_EQ(context, padded_begins.size(), kMaxDim); \
|
|
TF_LITE_ENSURE_EQ(context, padded_sizes.size(), kMaxDim); \
|
|
tflite::SliceParams op_params; \
|
|
op_params.begin_count = kMaxDim; \
|
|
op_params.size_count = kMaxDim; \
|
|
for (int i = 0; i < kMaxDim; ++i) { \
|
|
op_params.begin[i] = padded_begins[i]; \
|
|
op_params.size[i] = padded_sizes[i]; \
|
|
} \
|
|
\
|
|
if (kernel_type == kGenericOptimized) { \
|
|
optimized_ops::SliceInt4(op_params, GetTensorShape(input), input, \
|
|
GetTensorShape(output), output); \
|
|
} else { \
|
|
reference_ops::SliceInt4(op_params, GetTensorShape(input), input, \
|
|
GetTensorShape(output), output); \
|
|
} \
|
|
}
|
|
|
|
#define TF_LITE_SLICE(data_type) \
|
|
{ \
|
|
TF_LITE_ENSURE_EQ(context, padded_begins.size(), kMaxDim); \
|
|
TF_LITE_ENSURE_EQ(context, padded_sizes.size(), kMaxDim); \
|
|
tflite::SliceParams op_params; \
|
|
op_params.begin_count = kMaxDim; \
|
|
op_params.size_count = kMaxDim; \
|
|
for (int i = 0; i < kMaxDim; ++i) { \
|
|
op_params.begin[i] = padded_begins[i]; \
|
|
op_params.size[i] = padded_sizes[i]; \
|
|
} \
|
|
\
|
|
if (kernel_type == kGenericOptimized) { \
|
|
optimized_ops::Slice<data_type>(op_params, GetTensorShape(input), input, \
|
|
GetTensorShape(output), output); \
|
|
} else { \
|
|
reference_ops::Slice<data_type>(op_params, GetTensorShape(input), input, \
|
|
GetTensorShape(output), output); \
|
|
} \
|
|
}
|
|
|
|
if (input->type == kTfLiteString) {
|
|
TF_LITE_SLICE(string);
|
|
return kTfLiteOk;
|
|
} else if (input->type == kTfLiteInt4) {
|
|
TF_LITE_SLICE_INT4();
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
switch (TfLiteTypeGetSizeBits(output->type)) {
|
|
case 8:
|
|
TF_LITE_SLICE(int8_t);
|
|
break;
|
|
case 16:
|
|
TF_LITE_SLICE(int16_t);
|
|
break;
|
|
case 32:
|
|
TF_LITE_SLICE(int32_t);
|
|
break;
|
|
case 64:
|
|
TF_LITE_SLICE(int64_t);
|
|
break;
|
|
default:
|
|
TF_LITE_KERNEL_LOG(
|
|
context, "Type %d is currently not supported by Slice.", input->type);
|
|
return kTfLiteError;
|
|
}
|
|
#undef TF_LITE_SLICE
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
} // namespace slice
|
|
|
|
TfLiteRegistration* Register_SLICE_REF() {
|
|
static TfLiteRegistration r = {nullptr, nullptr, slice::Prepare,
|
|
slice::Eval<slice::kReference>};
|
|
return &r;
|
|
}
|
|
|
|
TfLiteRegistration* Register_SLICE() {
|
|
static TfLiteRegistration r = {nullptr, nullptr, slice::Prepare,
|
|
slice::Eval<slice::kGenericOptimized>};
|
|
return &r;
|
|
}
|
|
|
|
} // namespace builtin
|
|
} // namespace ops
|
|
} // namespace tflite
|