wehub-resource-sync
375 lines
14 KiB
C++
375 lines
14 KiB
C++
/* Copyright 2017 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 <stdint.h>
|
|
|
|
#include <algorithm>
|
|
#include <limits>
|
|
#include <tuple>
|
|
#include <utility>
|
|
|
|
#include "tensorflow/lite/core/c/common.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/kernel_util.h"
|
|
#include "tensorflow/lite/string_util.h"
|
|
|
|
namespace tflite {
|
|
namespace ops {
|
|
namespace builtin {
|
|
namespace tile {
|
|
|
|
constexpr int kInputTensor = 0;
|
|
constexpr int kInputMultipliers = 1;
|
|
constexpr int kOutputTensor = 0;
|
|
|
|
namespace {
|
|
struct OpData {
|
|
// Indicates that 'Eval' is a noop as the output was written during 'Prepare'.
|
|
bool noop;
|
|
};
|
|
|
|
template <typename T>
|
|
TfLiteStatus MultiplyShapeDims(TfLiteContext* context,
|
|
const TfLiteIntArray& shape,
|
|
const TfLiteTensor* multipliers,
|
|
int num_dimensions,
|
|
TfLiteIntArray** output_shape) {
|
|
const T* multipliers_v = GetTensorData<T>(multipliers);
|
|
|
|
TfLiteIntArray* temp_shape = TfLiteIntArrayCreate(num_dimensions);
|
|
if (temp_shape == nullptr) {
|
|
TF_LITE_KERNEL_LOG(context, "Failed to allocate memory for output shape.");
|
|
return kTfLiteError;
|
|
}
|
|
for (int i = 0; i < num_dimensions; ++i) {
|
|
int64_t shape_data = static_cast<int64_t>(shape.data[i]);
|
|
int64_t multiplier = static_cast<int64_t>(multipliers_v[i]);
|
|
|
|
if (multiplier < 0) {
|
|
TfLiteIntArrayFree(temp_shape);
|
|
TF_LITE_KERNEL_LOG(context, "Multipliers must be non-negative.");
|
|
return kTfLiteError;
|
|
}
|
|
|
|
if (shape_data < 0 ||
|
|
(shape_data > 0 &&
|
|
multiplier > std::numeric_limits<int32_t>::max() / shape_data)) {
|
|
TfLiteIntArrayFree(temp_shape);
|
|
TF_LITE_KERNEL_LOG(context,
|
|
"Cannot multiply %lld and %lld. Output shape "
|
|
"dimensions must be in range [0, INT32_MAX].",
|
|
static_cast<long long>(shape_data),
|
|
static_cast<long long>(multiplier));
|
|
return kTfLiteError;
|
|
}
|
|
temp_shape->data[i] = static_cast<int32_t>(shape_data * multiplier);
|
|
}
|
|
*output_shape = temp_shape;
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
TfLiteStatus ResizeOutput(TfLiteContext* context, TfLiteNode* node) {
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
const TfLiteTensor* multipliers;
|
|
TF_LITE_ENSURE_OK(
|
|
context, GetInputSafe(context, node, kInputMultipliers, &multipliers));
|
|
|
|
const int num_dimensions = NumDimensions(input);
|
|
const int num_multipliers = NumElements(multipliers);
|
|
TF_LITE_ENSURE_EQ(context, num_dimensions, num_multipliers);
|
|
TfLiteIntArray* output_shape = nullptr;
|
|
switch (multipliers->type) {
|
|
case kTfLiteInt32:
|
|
TF_LITE_ENSURE_OK(context, MultiplyShapeDims<int32_t>(
|
|
context, *input->dims, multipliers,
|
|
num_dimensions, &output_shape));
|
|
break;
|
|
case kTfLiteInt64:
|
|
TF_LITE_ENSURE_OK(context, MultiplyShapeDims<int64_t>(
|
|
context, *input->dims, multipliers,
|
|
num_dimensions, &output_shape));
|
|
break;
|
|
default:
|
|
TF_LITE_KERNEL_LOG(context,
|
|
"Multipliers of type '%s' are not supported by tile.",
|
|
TfLiteTypeGetName(multipliers->type));
|
|
return kTfLiteError;
|
|
}
|
|
return context->ResizeTensor(context, output, output_shape);
|
|
}
|
|
|
|
template <typename T, typename M>
|
|
void CopyMultipleTimes(const T* in_data, int32_t in_size, M multiplier,
|
|
T* out_data) {
|
|
for (M i = 0; i < multiplier; ++i) {
|
|
const T* in_end = in_data + in_size;
|
|
T* new_out_data = std::copy(in_data, in_end, out_data);
|
|
in_data = out_data;
|
|
out_data = new_out_data;
|
|
}
|
|
}
|
|
|
|
template <typename M>
|
|
void CopyStringMultipleTimes(const TfLiteTensor* in_data, int in_data_index,
|
|
const int dimension_size, M multiplier,
|
|
DynamicBuffer* buffer) {
|
|
for (M i = 0; i < multiplier; ++i) {
|
|
for (int j = 0; j < dimension_size; ++j) {
|
|
const StringRef string_ref = GetString(in_data, in_data_index + j);
|
|
buffer->AddString(string_ref.str, string_ref.len);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename M>
|
|
std::pair<int, int> TileOneDimension(const TfLiteIntArray& in_dimensions,
|
|
const T* in_data, const M* multipliers,
|
|
T* out_data, int dimension) {
|
|
if (in_dimensions.size == 0) {
|
|
// If input tensor is a scalar, then just copy it to output (no need to
|
|
// multiply).
|
|
*out_data = *in_data;
|
|
return std::make_pair(0, 0);
|
|
}
|
|
|
|
const int dimension_size = in_dimensions.data[dimension];
|
|
if (dimension == in_dimensions.size - 1) {
|
|
CopyMultipleTimes(in_data, dimension_size, multipliers[dimension],
|
|
out_data);
|
|
return std::make_pair(
|
|
dimension_size,
|
|
dimension_size * static_cast<int>(multipliers[dimension]));
|
|
}
|
|
int total_stride_size = 0, total_tiled_stride_size = 0;
|
|
const T* copy_from_data = in_data;
|
|
T* copy_to_data = out_data;
|
|
for (int i = 0; i < dimension_size; ++i) {
|
|
int stride_size = 0, tiled_stride_size = 0;
|
|
std::tie(stride_size, tiled_stride_size) =
|
|
TileOneDimension(in_dimensions, copy_from_data, multipliers,
|
|
copy_to_data, dimension + 1);
|
|
copy_from_data += stride_size;
|
|
copy_to_data += tiled_stride_size;
|
|
total_stride_size += stride_size;
|
|
total_tiled_stride_size += tiled_stride_size;
|
|
}
|
|
CopyMultipleTimes(out_data, total_tiled_stride_size,
|
|
multipliers[dimension] - 1,
|
|
out_data + total_tiled_stride_size);
|
|
return std::make_pair(
|
|
total_stride_size,
|
|
total_tiled_stride_size * static_cast<int>(multipliers[dimension]));
|
|
}
|
|
|
|
template <typename M>
|
|
std::pair<int, int> TileStringOneDimension(
|
|
const TfLiteIntArray& in_dimensions, const TfLiteTensor* in_data,
|
|
int in_data_index, const M* multipliers, DynamicBuffer* buffer,
|
|
int buffer_index, int dimension, TfLiteTensor* out_data) {
|
|
if (in_dimensions.size == 0) {
|
|
const StringRef string_ref = GetString(in_data, in_data_index);
|
|
buffer->AddString(string_ref.str, string_ref.len);
|
|
buffer->WriteToTensor(out_data, /*new_shape=*/nullptr);
|
|
return {0, 0};
|
|
}
|
|
|
|
const int dimension_size = in_dimensions.data[dimension];
|
|
if (dimension == in_dimensions.size - 1) {
|
|
CopyStringMultipleTimes(in_data, in_data_index, dimension_size,
|
|
multipliers[dimension], buffer);
|
|
return {dimension_size,
|
|
dimension_size * static_cast<int>(multipliers[dimension])};
|
|
}
|
|
|
|
int total_stride_size = 0, total_tiled_stride_size = 0;
|
|
for (int i = 0; i < dimension_size; ++i) {
|
|
int stride_size, tiled_stride_size;
|
|
std::tie(stride_size, tiled_stride_size) = TileStringOneDimension(
|
|
in_dimensions, in_data, in_data_index + total_stride_size, multipliers,
|
|
buffer, buffer_index + total_tiled_stride_size, dimension + 1,
|
|
out_data);
|
|
total_stride_size += stride_size;
|
|
total_tiled_stride_size += tiled_stride_size;
|
|
}
|
|
|
|
buffer->WriteToTensor(out_data, /*new_shape=*/nullptr);
|
|
CopyStringMultipleTimes(out_data, buffer_index, total_tiled_stride_size,
|
|
multipliers[dimension] - 1, buffer);
|
|
|
|
return {total_stride_size,
|
|
total_tiled_stride_size * static_cast<int>(multipliers[dimension])};
|
|
}
|
|
|
|
template <typename T>
|
|
void Tile(const TfLiteIntArray& in_dimensions, const TfLiteTensor* in_data,
|
|
const TfLiteTensor* multipliers, TfLiteTensor* out_data) {
|
|
// Recursively tiles from the outermost to the innermost dimension.
|
|
switch (multipliers->type) {
|
|
case kTfLiteInt32:
|
|
TileOneDimension(in_dimensions, GetTensorData<T>(in_data),
|
|
GetTensorData<int32_t>(multipliers),
|
|
GetTensorData<T>(out_data), 0);
|
|
break;
|
|
case kTfLiteInt64:
|
|
TileOneDimension(in_dimensions, GetTensorData<T>(in_data),
|
|
GetTensorData<int64_t>(multipliers),
|
|
GetTensorData<T>(out_data), 0);
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
|
|
void TileString(const TfLiteIntArray& in_dimensions,
|
|
const TfLiteTensor* in_data, const TfLiteTensor* multipliers,
|
|
DynamicBuffer* buffer, TfLiteTensor* out_data) {
|
|
// Recursively tiles from the outermost to the innermost dimension.
|
|
switch (multipliers->type) {
|
|
case kTfLiteInt32:
|
|
TileStringOneDimension(in_dimensions, in_data, 0,
|
|
GetTensorData<int32_t>(multipliers), buffer, 0, 0,
|
|
out_data);
|
|
break;
|
|
case kTfLiteInt64:
|
|
TileStringOneDimension(in_dimensions, in_data, 0,
|
|
GetTensorData<int64_t>(multipliers), buffer, 0, 0,
|
|
out_data);
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
} // namespace
|
|
|
|
TfLiteStatus EvalImpl(TfLiteContext* context, const TfLiteTensor* input,
|
|
const TfLiteTensor* multipliers, TfLiteTensor* output) {
|
|
if (GetTensorShape(output).FlatSize() == 0) {
|
|
if (output->type == kTfLiteString) {
|
|
// For safety, ensure that we write to the output tensor.
|
|
DynamicBuffer buffer;
|
|
buffer.WriteToTensor(output, /*new_shape=*/nullptr);
|
|
}
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
if (output->type == kTfLiteString) {
|
|
DynamicBuffer buffer;
|
|
TileString(*(input->dims), input, multipliers, &buffer, output);
|
|
buffer.WriteToTensor(output, /*new_shape=*/nullptr);
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
switch (TfLiteTypeGetSizeBits(output->type)) {
|
|
case 8:
|
|
Tile<int8_t>(*(input->dims), input, multipliers, output);
|
|
break;
|
|
case 16:
|
|
Tile<int16_t>(*(input->dims), input, multipliers, output);
|
|
break;
|
|
case 32:
|
|
Tile<int32_t>(*(input->dims), input, multipliers, output);
|
|
break;
|
|
case 64:
|
|
Tile<int64_t>(*(input->dims), input, multipliers, output);
|
|
break;
|
|
default:
|
|
TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by tile.",
|
|
TfLiteTypeGetName(output->type));
|
|
return kTfLiteError;
|
|
}
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
|
TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
|
|
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
|
OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
|
|
op_data->noop = false;
|
|
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
|
|
|
|
const TfLiteTensor* multipliers;
|
|
TF_LITE_ENSURE_OK(
|
|
context, GetInputSafe(context, node, kInputMultipliers, &multipliers));
|
|
// Only int32 and int64 multiplier types are supported.
|
|
if (multipliers->type != kTfLiteInt32 && multipliers->type != kTfLiteInt64) {
|
|
TF_LITE_KERNEL_LOG(context,
|
|
"Multipliers of type '%s' are not supported by tile.",
|
|
TfLiteTypeGetName(multipliers->type));
|
|
return kTfLiteError;
|
|
}
|
|
|
|
if (IsConstantOrPersistentTensor(multipliers)) {
|
|
if (IsConstantOrPersistentTensor(input)) {
|
|
SetTensorToPersistentRo(output);
|
|
TF_LITE_ENSURE_OK(context, ResizeOutput(context, node));
|
|
op_data->noop = true;
|
|
return EvalImpl(context, input, multipliers, output);
|
|
}
|
|
TF_LITE_ENSURE_OK(context, ResizeOutput(context, node));
|
|
} else {
|
|
SetTensorToDynamic(output);
|
|
}
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
const TfLiteTensor* multipliers;
|
|
TF_LITE_ENSURE_OK(
|
|
context, GetInputSafe(context, node, kInputMultipliers, &multipliers));
|
|
|
|
if (reinterpret_cast<OpData*>(node->user_data)->noop) {
|
|
return kTfLiteOk;
|
|
}
|
|
if (IsDynamicTensor(output)) {
|
|
TF_LITE_ENSURE_OK(context, ResizeOutput(context, node));
|
|
}
|
|
return EvalImpl(context, input, multipliers, output);
|
|
}
|
|
|
|
void* Init(TfLiteContext* context, const char* buffer, size_t length) {
|
|
return new OpData();
|
|
}
|
|
|
|
void Free(TfLiteContext* context, void* buffer) {
|
|
delete reinterpret_cast<OpData*>(buffer);
|
|
}
|
|
|
|
} // namespace tile
|
|
TfLiteRegistration* Register_TILE() {
|
|
static TfLiteRegistration r = {tile::Init, tile::Free, tile::Prepare,
|
|
tile::Eval};
|
|
return &r;
|
|
}
|
|
} // namespace builtin
|
|
} // namespace ops
|
|
} // namespace tflite
|