diff options
Diffstat (limited to 'libs/support/tflite/src/kernels/TensorFlowMax.cpp')
-rw-r--r-- | libs/support/tflite/src/kernels/TensorFlowMax.cpp | 390 |
1 files changed, 390 insertions, 0 deletions
diff --git a/libs/support/tflite/src/kernels/TensorFlowMax.cpp b/libs/support/tflite/src/kernels/TensorFlowMax.cpp new file mode 100644 index 000000000..abc6fda4e --- /dev/null +++ b/libs/support/tflite/src/kernels/TensorFlowMax.cpp @@ -0,0 +1,390 @@ +/* + * Copyright (c) 2018 Samsung Electronics Co., Ltd. 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 "support/tflite/kernels/TensorFlowMax.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" + +#include <iostream> + +namespace tflite +{ +namespace ops +{ +namespace custom +{ +namespace nnfw +{ +namespace TensorFlowMax +{ + +struct TensorFlowMaxOp +{ + TensorFlowMaxOp(TfLiteContext *context, TfLiteNode *node) + { + input = tflite::GetInput(context, node, 0); + axis = tflite::GetInput(context, node, 1); + output = tflite::GetOutput(context, node, 0); + } + const TfLiteTensor *input; + const TfLiteTensor *axis; + TfLiteTensor *output; +}; + +void *InitTensorFlowMax(TfLiteContext *context, const char *buffer, size_t length) +{ + // Creates two temp tensors to store index and axis for internal + // implementation only. + auto *scratch_tensor_index = new int; + context->AddTensors(context, 2, scratch_tensor_index); + return scratch_tensor_index; +} + +void FreeTensorFlowMax(TfLiteContext *context, void *buffer) +{ + delete static_cast<TensorFlowMaxOp *>(buffer); +} + +// Resizes the temp tensor that stores resolved axis. +TfLiteStatus ResizeTempAxis(TfLiteContext *context, TensorFlowMaxOp *op_context, + TfLiteTensor *resolved_axis) +{ + TfLiteIntArray *axis_size = TfLiteIntArrayCreate(1); + axis_size->data[0] = static_cast<int>(tflite::NumElements(op_context->axis)); + return context->ResizeTensor(context, resolved_axis, axis_size); +} + +// Resizes output array based on the input size and resolved axis. +TfLiteStatus ResizeOutputTensor(TfLiteContext *context, TensorFlowMaxOp *op_context) +{ + size_t num_axis = tflite::NumElements(op_context->axis); + const TfLiteIntArray *input_dims = op_context->input->dims; + int input_num_dims = tflite::NumDimensions(op_context->input); + const int *axis = op_context->axis->data.i32; + + { + // Calculates size of reducing axis. + int num_reduce_axis = num_axis; + for (int i = 0; i < num_axis; ++i) + { + int current = axis[i]; + if (current < 0) + { + current += input_num_dims; + } + TF_LITE_ENSURE(context, current >= 0 && current < input_num_dims); + for (int j = 0; j < i; ++j) + { + int previous = axis[j]; + if (previous < 0) + { + previous += input_num_dims; + } + if (current == previous) + { + --num_reduce_axis; + break; + } + } + } + // Determines output dimensions. + TfLiteIntArray *output_dims = TfLiteIntArrayCreate(input_num_dims - num_reduce_axis); + int num_skip_axis = 0; + for (int idx = 0; idx < input_num_dims; ++idx) + { + bool is_axis = false; + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) + { + if (axis[axis_idx] == idx || axis[axis_idx] + input_num_dims == idx) + { + ++num_skip_axis; + is_axis = true; + break; + } + } + if (!is_axis) + { + output_dims->data[idx - num_skip_axis] = input_dims->data[idx]; + } + } + return context->ResizeTensor(context, op_context->output, output_dims); + } +} + +// Initializes temp tensors to store index and resolved axis. +TfLiteStatus InitializeTemporaries(TfLiteContext *context, TfLiteNode *node, + TensorFlowMaxOp *op_context) +{ + // Creates a temp index to iterate through input data. + int *scratch_tensor_index = reinterpret_cast<int *>(node->user_data); + TfLiteIntArrayFree(node->temporaries); + node->temporaries = TfLiteIntArrayCreate(2); + node->temporaries->data[0] = *scratch_tensor_index; + TfLiteTensor *scratch_tensor = &context->tensors[node->temporaries->data[0]]; + scratch_tensor->type = kTfLiteInt32; + scratch_tensor->allocation_type = kTfLiteArenaRw; + TfLiteIntArray *index_size = TfLiteIntArrayCreate(1); + index_size->data[0] = tflite::NumDimensions(op_context->input); + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_tensor, index_size)); + + // Creates a temp tensor to store resolved axis given input data. + node->temporaries->data[1] = *scratch_tensor_index + 1; + TfLiteTensor *resolved_axis = &context->tensors[node->temporaries->data[1]]; + resolved_axis->type = kTfLiteInt32; + return kTfLiteOk; +} + +TfLiteStatus PrepareTensorFlowMax(TfLiteContext *context, TfLiteNode *node) +{ + TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1); + + TensorFlowMaxOp op_context(context, node); + TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context)); + + TfLiteTensor *resolved_axis = &context->tensors[node->temporaries->data[1]]; + // Leaves work to Eval if axis is not constant; else resizes output. + if (!tflite::IsConstantTensor(op_context.axis)) + { + tflite::SetTensorToDynamic(op_context.output); + tflite::SetTensorToDynamic(resolved_axis); + return kTfLiteOk; + } + resolved_axis->allocation_type = kTfLiteArenaRw; + TF_LITE_ENSURE_OK(context, ResizeTempAxis(context, &op_context, resolved_axis)); + return ResizeOutputTensor(context, &op_context); +} + +// Gets offset of index if expanded on axis. When expanded, the flattened offset +// will not change, if the output index changes on the given axis. For example, +// if you have a 2D tensor and you are expanding to 3D on axis 0, +// then index (0, 1, 2) and index (1, 1, 2) will map from the same flattened +// offset. +inline size_t ExpandedInputOffset(const int num_dims, const int *dims, const int *index, + const int num_axis, const int *axis) +{ + size_t offset = 0; + int out_idx = 0; + for (int in_idx = 0; in_idx < num_dims; ++in_idx) + { + // if we need to expand this axis + bool is_axis = false; + if (axis != nullptr) + { + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) + { + if (in_idx == axis[axis_idx]) + { + is_axis = true; + break; + } + } + } + if (!is_axis) + { + offset = offset * static_cast<size_t>(dims[in_idx]) + static_cast<size_t>(index[out_idx]); + out_idx++; + } + else + { + offset = offset * static_cast<size_t>(dims[in_idx]); + } + } + return offset; +} + +// Gets offset of index if reducing on axis. When reducing, the flattened offset +// will not change, if the input index changes on the given axis. For example, +// if you have a 3D tensor and you are reducing to 2D by eliminating axis 0, +// then index (0, 1, 2) and index (1, 1, 2) will map to the same flattened +// offset. +// TODO(kanlig): uses Dims to represent dimensions. +inline size_t ReducedOutputOffset(const int num_dims, const int *dims, const int *index, + const int num_axis, const int *axis) +{ + size_t offset = 0; + for (int idx = 0; idx < num_dims; ++idx) + { + // if we need to skip this axis + bool is_axis = false; + if (axis != nullptr) + { + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) + { + if (idx == axis[axis_idx]) + { + is_axis = true; + break; + } + } + } + if (!is_axis) + { + offset = offset * static_cast<size_t>(dims[idx]) + static_cast<size_t>(index[idx]); + } + } + return offset; +} + +// Gets next index to iterate through a multidimensional array. +inline bool NextIndex(TfLiteContext *context, const int num_dims, const int *dims, int *current) +{ + int carry = 1; + for (int idx = num_dims - 1; idx >= 0; --idx) + { + int current_val = current[idx] + carry; + TF_LITE_ENSURE(context, (dims[idx] >= current_val)); + if (dims[idx] == current_val) + { + current[idx] = 0; + } + else + { + current[idx] = current_val; + carry = 0; + break; + } + } + return (carry == 0); +} + +template <typename T> +inline TfLiteStatus +CustomMax(TfLiteContext *context, T *input_data, const int *input_dims, const int input_num_dims, + T *output_data, const int *output_dims, const int output_num_dims, const int *axis, + const int num_axis_dimensions, bool keep_dims, int *temp_index, int *resolved_axis) +{ + // resolves axis. + int num_resolved_axis = 0; + for (int idx = 0; idx < num_axis_dimensions; ++idx) + { + int current = axis[idx]; + TF_LITE_ENSURE(context, (current < input_num_dims && current + input_num_dims >= 0)); + if (current < 0) + { + current += input_num_dims; + } + bool is_dup = false; + for (int j = 0; j < num_resolved_axis; ++j) + { + if (resolved_axis[j] == current) + { + is_dup = true; + break; + } + } + if (!is_dup) + { + resolved_axis[num_resolved_axis++] = current; + } + } + + TF_LITE_ENSURE(context, (input_num_dims > 0)); + TF_LITE_ENSURE(context, (input_dims != nullptr)); + TF_LITE_ENSURE(context, (temp_index != nullptr)); + + // resets output data. + for (int idx = 0; idx < output_num_dims; ++idx) + { + temp_index[idx] = 0; + } + for (bool has_next = true; has_next; + has_next = NextIndex(context, output_num_dims, output_dims, temp_index)) + { + size_t output_offset = + ReducedOutputOffset(output_num_dims, output_dims, temp_index, 0, nullptr); + size_t input_offset = ExpandedInputOffset(input_num_dims, input_dims, temp_index, + num_resolved_axis, resolved_axis); + output_data[output_offset] = input_data[input_offset]; + } + + // resets temp index. + for (int idx = 0; idx < input_num_dims; ++idx) + { + temp_index[idx] = 0; + } + + // iterates through input_data. + for (bool has_next = true; has_next; + has_next = NextIndex(context, input_num_dims, input_dims, temp_index)) + { + size_t input_offset = ReducedOutputOffset(input_num_dims, input_dims, temp_index, 0, nullptr); + size_t output_offset = ReducedOutputOffset(input_num_dims, input_dims, temp_index, + num_resolved_axis, resolved_axis); + if (output_data[output_offset] < input_data[input_offset]) + { + output_data[output_offset] = input_data[input_offset]; + } + } + + return kTfLiteOk; +} + +TfLiteStatus EvalTensorFlowMax(TfLiteContext *context, TfLiteNode *node) +{ + + TensorFlowMaxOp op_context(context, node); + int num_axis = static_cast<int>(tflite::NumElements(op_context.axis)); + TfLiteTensor *temp_index = &context->tensors[node->temporaries->data[0]]; + TfLiteTensor *resolved_axis = &context->tensors[node->temporaries->data[1]]; + // Resize the output tensor if the output tensor is dynamic. + if (tflite::IsDynamicTensor(op_context.output)) + { + TF_LITE_ENSURE_OK(context, ResizeTempAxis(context, &op_context, resolved_axis)); + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + + TfLiteStatus returnStatus = kTfLiteOk; + switch (op_context.input->type) + { + case kTfLiteFloat32: + returnStatus = CustomMax<float>( + context, op_context.input->data.f, op_context.input->dims->data, + op_context.input->dims->size, op_context.output->data.f, op_context.output->dims->data, + op_context.output->dims->size, op_context.axis->data.i32, num_axis, false, + temp_index->data.i32, resolved_axis->data.i32); + break; + case kTfLiteInt32: + returnStatus = CustomMax<int>(context, op_context.input->data.i32, + op_context.input->dims->data, op_context.input->dims->size, + op_context.output->data.i32, op_context.output->dims->data, + op_context.output->dims->size, op_context.axis->data.i32, + num_axis, false, temp_index->data.i32, resolved_axis->data.i32); + break; + case kTfLiteUInt8: + returnStatus = CustomMax<uint8_t>( + context, op_context.input->data.uint8, op_context.input->dims->data, + op_context.input->dims->size, op_context.output->data.uint8, + op_context.output->dims->data, op_context.output->dims->size, op_context.axis->data.i32, + num_axis, false, temp_index->data.i32, resolved_axis->data.i32); + break; + case kTfLiteInt64: + returnStatus = CustomMax<int64_t>( + context, op_context.input->data.i64, op_context.input->dims->data, + op_context.input->dims->size, op_context.output->data.i64, op_context.output->dims->data, + op_context.output->dims->size, op_context.axis->data.i32, num_axis, false, + temp_index->data.i32, resolved_axis->data.i32); + break; + default: + returnStatus = kTfLiteError; + } + + return returnStatus; +} +} // namespace TensorFlowMax +} // namespace nnfw +} // namespace custom +} // namespace ops +} // namespace tflite |