COMPMID-1528: Add ReorgLayer on NEON

Change-Id: I44369b4a716767163e2233b7d87bff300c523383
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/146314
Reviewed-by: Pablo Tello <pablo.tello@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>

1 files changed, 16 insertions, 0 deletions

diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index 1a86d27727..9c7cfecd4c 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -57,6 +57,22 @@ inline TensorShape compute_permutation_output_shape(const ITensorInfo &input, co
     permute(output_shape, perm);
     return output_shape;
 }
+inline TensorShape compute_reorg_output_shape(const ITensorInfo &input, int32_t stride)
+{
+    ARM_COMPUTE_ERROR_ON(stride <= 0);
+
+    const DataLayout   data_layout = input.data_layout();
+    const unsigned int width_idx   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
+    const unsigned int height_idx  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+    const unsigned int channel_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
+
+    TensorShape output_shape{ input.tensor_shape() };
+    output_shape.set(width_idx, input.tensor_shape()[width_idx] / stride);
+    output_shape.set(height_idx, input.tensor_shape()[height_idx] / stride);
+    output_shape.set(channel_idx, input.tensor_shape()[channel_idx] * stride * stride);
+
+    return output_shape;
+}
 inline TensorShape compute_weights_reshaped_shape(const ITensorInfo &weights, bool has_bias = false, unsigned int num_groups = 1)
 {
     // Number of groups greater than one are only supported for NCHW data layout, and the number of weights must be a multiple of it.