From 1e73bbda958564c09380113d17e3f36c563411ae Mon Sep 17 00:00:00 2001 From: Ryan Kuester Date: Mon, 28 Dec 2020 17:45:12 -0600 Subject: [PATCH] Extract reference for op SPACE_TO_DEPTH to standalone header Move the reference implementation to its own header so that micro can use it without including unrelated depedencies via reference_ops.h. --- tensorflow/lite/kernels/internal/BUILD | 2 + .../internal/reference/reference_ops.h | 53 +------------ .../internal/reference/space_to_depth.h | 78 +++++++++++++++++++ 3 files changed, 81 insertions(+), 52 deletions(-) create mode 100644 tensorflow/lite/kernels/internal/reference/space_to_depth.h diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD index c619564ee5b..f5d7ad0c225 100644 --- a/tensorflow/lite/kernels/internal/BUILD +++ b/tensorflow/lite/kernels/internal/BUILD @@ -493,6 +493,7 @@ cc_library( "reference/resize_nearest_neighbor.h", "reference/round.h", "reference/softmax.h", + "reference/space_to_depth.h", "reference/strided_slice.h", "reference/sub.h", "reference/svdf.h", @@ -587,6 +588,7 @@ cc_library( "reference/resize_nearest_neighbor.h", "reference/round.h", "reference/softmax.h", + "reference/space_to_depth.h", "reference/strided_slice.h", "reference/string_comparisons.h", "reference/sub.h", diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h index 35d811c534f..765ce4df063 100644 --- a/tensorflow/lite/kernels/internal/reference/reference_ops.h +++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h @@ -61,6 +61,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h" #include "tensorflow/lite/kernels/internal/reference/round.h" #include "tensorflow/lite/kernels/internal/reference/softmax.h" +#include "tensorflow/lite/kernels/internal/reference/space_to_depth.h" #include "tensorflow/lite/kernels/internal/reference/strided_slice.h" #include "tensorflow/lite/kernels/internal/reference/string_comparisons.h" #include "tensorflow/lite/kernels/internal/reference/sub.h" @@ -125,58 +126,6 @@ inline void DepthToSpace(const tflite::DepthToSpaceParams& op_params, } } -template -inline void SpaceToDepth(const tflite::SpaceToDepthParams& op_params, - const RuntimeShape& unextended_input_shape, - const T* input_data, - const RuntimeShape& unextended_output_shape, - T* output_data) { - TFLITE_DCHECK_LE(unextended_input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4); - const RuntimeShape input_shape = - RuntimeShape::ExtendedShape(4, unextended_input_shape); - const RuntimeShape output_shape = - RuntimeShape::ExtendedShape(4, unextended_output_shape); - - const int input_depth = input_shape.Dims(3); - const int input_width = input_shape.Dims(2); - const int input_height = input_shape.Dims(1); - const int input_batch = input_shape.Dims(0); - - const int output_depth = output_shape.Dims(3); - const int output_width = output_shape.Dims(2); - const int output_height = output_shape.Dims(1); - const int output_batch = output_shape.Dims(0); - - const int32 block_size = op_params.block_size; - - TFLITE_DCHECK_EQ(input_width, output_width * block_size); - TFLITE_DCHECK_EQ(input_height, output_height * block_size); - TFLITE_DCHECK_EQ(input_depth * block_size * block_size, output_depth); - TFLITE_DCHECK_EQ(input_batch, output_batch); - - for (int in_b = 0; in_b < input_batch; ++in_b) { - for (int in_h = 0; in_h < input_height; ++in_h) { - for (int in_w = 0; in_w < input_width; ++in_w) { - for (int in_d = 0; in_d < input_depth; ++in_d) { - const int out_d = - in_d + ((in_h % block_size) * block_size + in_w % block_size) * - input_depth; - const int out_w = in_w / block_size; - const int out_h = in_h / block_size; - const int out_b = in_b; - - const int input_index = Offset(input_shape, in_b, in_h, in_w, in_d); - const int output_index = - Offset(output_shape, out_b, out_h, out_w, out_d); - - output_data[output_index] = input_data[input_index]; - } - } - } - } -} - inline void Elu(const RuntimeShape& input_shape, const float* input_data, const RuntimeShape& output_shape, float* output_data) { const int flat_size = MatchingFlatSize(input_shape, output_shape); diff --git a/tensorflow/lite/kernels/internal/reference/space_to_depth.h b/tensorflow/lite/kernels/internal/reference/space_to_depth.h new file mode 100644 index 00000000000..61343ac6e35 --- /dev/null +++ b/tensorflow/lite/kernels/internal/reference/space_to_depth.h @@ -0,0 +1,78 @@ +/* Copyright 2020 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. +==============================================================================*/ +#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SPACE_TO_DEPTH_H_ +#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SPACE_TO_DEPTH_H_ + +#include "tensorflow/lite/kernels/internal/types.h" + +namespace tflite { +namespace reference_ops { + +template +inline void SpaceToDepth(const tflite::SpaceToDepthParams& op_params, + const RuntimeShape& unextended_input_shape, + const T* input_data, + const RuntimeShape& unextended_output_shape, + T* output_data) { + TFLITE_DCHECK_LE(unextended_input_shape.DimensionsCount(), 4); + TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4); + const RuntimeShape input_shape = + RuntimeShape::ExtendedShape(4, unextended_input_shape); + const RuntimeShape output_shape = + RuntimeShape::ExtendedShape(4, unextended_output_shape); + + const int input_depth = input_shape.Dims(3); + const int input_width = input_shape.Dims(2); + const int input_height = input_shape.Dims(1); + const int input_batch = input_shape.Dims(0); + + const int output_depth = output_shape.Dims(3); + const int output_width = output_shape.Dims(2); + const int output_height = output_shape.Dims(1); + const int output_batch = output_shape.Dims(0); + + const int32 block_size = op_params.block_size; + + TFLITE_DCHECK_EQ(input_width, output_width * block_size); + TFLITE_DCHECK_EQ(input_height, output_height * block_size); + TFLITE_DCHECK_EQ(input_depth * block_size * block_size, output_depth); + TFLITE_DCHECK_EQ(input_batch, output_batch); + + for (int in_b = 0; in_b < input_batch; ++in_b) { + for (int in_h = 0; in_h < input_height; ++in_h) { + for (int in_w = 0; in_w < input_width; ++in_w) { + for (int in_d = 0; in_d < input_depth; ++in_d) { + const int out_d = + in_d + ((in_h % block_size) * block_size + in_w % block_size) * + input_depth; + const int out_w = in_w / block_size; + const int out_h = in_h / block_size; + const int out_b = in_b; + + const int input_index = Offset(input_shape, in_b, in_h, in_w, in_d); + const int output_index = + Offset(output_shape, out_b, out_h, out_w, out_d); + + output_data[output_index] = input_data[input_index]; + } + } + } + } +} + +} // namespace reference_ops +} // namespace tflite + +#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SPACE_TO_DEPTH_H_