1 files changed, 129 insertions, 0 deletions
diff --git a/runtimes/libs/ARMComputeEx/arm_compute/core/utils/misc/ShapeCalculatorEx.h b/runtimes/libs/ARMComputeEx/arm_compute/core/utils/misc/ShapeCalculatorEx.h
new file mode 100644
index 000000000..bacb1140c
--- /dev/null
+++ b/runtimes/libs/ARMComputeEx/arm_compute/core/utils/misc/ShapeCalculatorEx.h
@@ -0,0 +1,129 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ * Copyright (c) 2016-2018 ARM Limited.
+ *
+ * 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 __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__
+#define __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__
+
+#include "arm_compute/core/Helpers.h"
+#include "arm_compute/core/ITensorInfo.h"
+#include "arm_compute/core/Utils.h"
+
+#include "arm_compute/core/utils/helpers/tensor_transform.h"
+
+#include <cmath>
+
+namespace arm_compute
+{
+namespace misc
+{
+namespace shape_calculator
+{
+
+/** Calculate the upsampled output shape used for transpose convolution
+ *
+ * @param[in] input              Input tensor info
+ * @param[in] weights            Weights tensor shape
+ * @param[in] info               Padding and stride info
+ * @param[in] out_dims           Output shape dimensions
+ * @param[in] invalid_right      The number of zeros added to right edge of the output.
+ * @param[in] invalid_bottom     The number of zeros added to bottom edge of the output.
+ * @param[out] pad_left          Padding on left
+ * @param[out] pad_right         Padding on right
+ * @param[out] pad_top           Padding on top
+ * @param[out] pad_bottom        Padding on bottom
+ *
+ * @return the calculated shape
+ */
+inline TensorShape compute_transposeconv_upsampled_shape(
+    const ITensorInfo &input, const ITensorInfo &weights, const PadStrideInfo &info,
+    std::pair<unsigned int, unsigned int> &out_dims, unsigned int invalid_right,
+    unsigned int invalid_bottom, unsigned int &pad_left, unsigned int &pad_right,
+    unsigned int &pad_top, unsigned int &pad_bottom)
+{
+  unsigned int sx = info.stride().first;
+  unsigned int sy = info.stride().second;
+  const DataLayout data_layout = input.data_layout();
+  const size_t idx_w = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
+  const size_t idx_h = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+
+  // Find the upsampled dimensions
+  // transpose conv out:
+  //    tconv_out + pad = 1 + (in - 1) * stride + invalid
+  //    tconv_out = 1 + (in - 1) * stride + invalid - pad
+  // upsample out:
+  //    upsample_out = 1 + (in - 1) * stride
+  unsigned int out_x = (input.dimension(idx_w) - 1) * sx + 1;
+  unsigned int out_y = (input.dimension(idx_h) - 1) * sy + 1;
+
+  // Find the padding needed for the convolution with stride 1 in order to match output shape
+  // upsample+pad out:
+  //    upsample_out + pad = tconv_out + kernel - 1
+  //    pad = tconv_out + kernel - 1 - upsample_out
+  unsigned int padx = out_dims.first - (out_x - weights.dimension(idx_w) + 1);
+  unsigned int pady = out_dims.second - (out_y - weights.dimension(idx_h) + 1);
+  out_x += padx;
+  out_y += pady;
+
+  unsigned int padx_all_except_invallid = padx + info.pad_left() + info.pad_right() - invalid_right;
+  unsigned int pady_all_except_invallid =
+      pady + info.pad_top() + info.pad_bottom() - invalid_bottom;
+  pad_left = (padx_all_except_invallid + 1) / 2 - info.pad_left();
+  pad_right = pady_all_except_invallid / 2 - info.pad_right() + invalid_right;
+  pad_top = (padx_all_except_invallid + 1) / 2 - info.pad_top();
+  pad_bottom = pady_all_except_invallid / 2 - info.pad_bottom() + invalid_bottom;
+
+  TensorShape scale_out_shape(input.tensor_shape());
+  scale_out_shape.set(idx_w, out_x);
+  scale_out_shape.set(idx_h, out_y);
+
+  return scale_out_shape;
+}
+
+/** Calculate the output shape of the transpose convolution layer
+ *
+ * @param[in] out_dims Output x and y shape dimensions
+ * @param[in] input    Input tensor info
+ * @param[in] weights  Weights tensor shape
+ *
+ * @return the calculated shape
+ */
+inline TensorShape
+compute_transposeconv_output_shape(const std::pair<unsigned int, unsigned int> &out_dims,
+                                   const ITensorInfo &input, const ITensorInfo &weights)
+{
+  const TensorShape input_shape{input.tensor_shape()};
+  const TensorShape weights_shape{weights.tensor_shape()};
+
+  const DataLayout data_layout = input.data_layout();
+  const int width_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
+  const int height_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+  const int channel_idx =
+      get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
+  const int batch_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::BATCHES);
+
+  TensorShape out_shape{input_shape};
+  out_shape.set(width_idx, out_dims.first);
+  out_shape.set(height_idx, out_dims.second);
+  out_shape.set(channel_idx, weights_shape[batch_idx]);
+  return out_shape;
+}
+
+} // namespace shape_calculator
+} // namespace misc
+} // namespace arm_compute
+
+#endif // __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__