COMPMID-1527 - Implementing ReorgLayer on OpenCL

Also extended tests on NEON

Change-Id: Icb0eced534e904ef807972dd3a31988f501bb02e
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/147095
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>

4 files changed, 258 insertions, 103 deletions

diff --git a/tests/validation/CL/ReorgLayer.cpp b/tests/validation/CL/ReorgLayer.cpp
new file mode 100644
index 0000000000..e4caa04c3a
--- /dev/null
+++ b/tests/validation/CL/ReorgLayer.cpp
@@ -0,0 +1,168 @@
+/*
+ * Copyright (c) 2018 ARM Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#include "arm_compute/core/Types.h"
+#include "arm_compute/runtime/CL/functions/CLReorgLayer.h"
+#include "tests/CL/CLAccessor.h"
+#include "tests/CL/Helper.h"
+#include "tests/PaddingCalculator.h"
+#include "tests/datasets/ReorgLayerDataset.h"
+#include "tests/framework/Asserts.h"
+#include "tests/framework/Macros.h"
+#include "tests/framework/datasets/Datasets.h"
+#include "tests/validation/Validation.h"
+#include "tests/validation/fixtures/ReorgLayerFixture.h"
+
+namespace arm_compute
+{
+namespace test
+{
+namespace validation
+{
+TEST_SUITE(CL)
+TEST_SUITE(ReorgLayer)
+
+// *INDENT-OFF*
+// clang-format off
+DATA_TEST_CASE(Validate, framework::DatasetMode::ALL, zip(zip(zip(
+               framework::dataset::make("InputInfo", { TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::S64),    // Wrong output tensor
+                                                       TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::F32),    // Wrong output tensor
+                                                       TensorInfo(TensorShape(3U, 12U, 4U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(3U, 12U, 4U, 2U), 1, DataType::F32),     // Wrong data type
+                                                     }),
+               framework::dataset::make("OutputInfo",{ TensorInfo(TensorShape(3U, 4U, 10U, 2U), 1, DataType::S64),
+                                                       TensorInfo(TensorShape(5U, 6U, 4U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(5U, 6U, 2, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(1U, 4U, 36U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(1U, 4U, 36U, 2U), 1, DataType::F16),
+                                                     })),
+               framework::dataset::make("Stride", { 2, 2, 4, 3 })),
+               framework::dataset::make("Expected", { false, true, false, true, false })),
+               input_info, output_info, stride, expected)
+{
+    bool status = bool(CLReorgLayer::validate(&input_info, &output_info, stride));
+    ARM_COMPUTE_EXPECT(status == expected, framework::LogLevel::ERRORS);
+}
+// clang-format on
+// *INDENT-ON*
+
+DATA_TEST_CASE(Configuration, framework::DatasetMode::ALL, combine(combine(framework::dataset::concat(datasets::SmallReorgLayerDataset(), datasets::LargeReorgLayerDataset()),
+                                                                           framework::dataset::make("DataType", { DataType::F32, DataType::F16, DataType::QASYMM8 })),
+                                                                   framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })),
+               shape, stride, data_type, data_layout)
+{
+    // Permute the tensor shape in case of NHWC data layout
+    TensorShape shape_to_use = shape;
+    if(data_layout == DataLayout::NHWC)
+    {
+        permute(shape_to_use, PermutationVector(2U, 0U, 1U));
+    }
+
+    // Create tensors
+    CLTensor src = create_tensor<CLTensor>(shape_to_use, data_type, 1, QuantizationInfo(), data_layout);
+    CLTensor dst;
+
+    ARM_COMPUTE_EXPECT(src.info()->is_resizable(), framework::LogLevel::ERRORS);
+
+    // Create and configure function
+    CLReorgLayer reorg_layer;
+
+    // Auto-initialize the output within the function
+    reorg_layer.configure(&src, &dst, stride);
+
+    // Validate valid region
+    const ValidRegion src_valid_region = shape_to_valid_region(shape_to_use);
+    const ValidRegion dst_valid_region = shape_to_valid_region(dst.info()->tensor_shape());
+    validate(src.info()->valid_region(), src_valid_region);
+    validate(dst.info()->valid_region(), dst_valid_region);
+
+    // Validate padding
+    const int         step        = 1;
+    const PaddingSize src_padding = PaddingCalculator(shape_to_use.x(), step).required_padding();
+    const PaddingSize dst_padding = PaddingCalculator(dst.info()->tensor_shape().x(), step).required_padding();
+    validate(src.info()->padding(), src_padding);
+    validate(dst.info()->padding(), dst_padding);
+}
+
+template <typename T>
+using CLReorgLayerFixture = ReorgLayerValidationFixture<CLTensor, CLAccessor, CLReorgLayer, T>;
+
+TEST_SUITE(S32)
+FIXTURE_DATA_TEST_CASE(RunSmall, CLReorgLayerFixture<int32_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                  DataType::S32)),
+                                                                                                          framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+
+FIXTURE_DATA_TEST_CASE(RunLarge, CLReorgLayerFixture<int32_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                DataType::S32)),
+                                                                                                        framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+TEST_SUITE_END() // S32
+
+TEST_SUITE(S16)
+FIXTURE_DATA_TEST_CASE(RunSmall, CLReorgLayerFixture<int16_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                  DataType::S16)),
+                                                                                                          framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+
+FIXTURE_DATA_TEST_CASE(RunLarge, CLReorgLayerFixture<int16_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                DataType::S16)),
+                                                                                                        framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+TEST_SUITE_END() // S16
+
+TEST_SUITE(S8)
+FIXTURE_DATA_TEST_CASE(RunSmall, CLReorgLayerFixture<int8_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                 DataType::S8)),
+                                                                                                         framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+
+FIXTURE_DATA_TEST_CASE(RunLarge, CLReorgLayerFixture<int8_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType", DataType::S8)),
+                                                                                                       framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+{
+    // Validate output
+    validate(CLAccessor(_target), _reference);
+}
+TEST_SUITE_END() // S8
+
+TEST_SUITE_END() // ReorgLayer
+TEST_SUITE_END() // CL
+} // namespace validation
+} // namespace test
+} // namespace arm_compute
diff --git a/tests/validation/NEON/ReorgLayer.cpp b/tests/validation/NEON/ReorgLayer.cpp
index 6489b6529f..5a76315d77 100644
--- a/tests/validation/NEON/ReorgLayer.cpp
+++ b/tests/validation/NEON/ReorgLayer.cpp
@@ -28,7 +28,7 @@
 
 #include "tests/NEON/Accessor.h"
 #include "tests/PaddingCalculator.h"
-#include "tests/datasets/ShapeDatasets.h"
+#include "tests/datasets/ReorgLayerDataset.h"
 #include "tests/framework/Asserts.h"
 #include "tests/framework/Macros.h"
 #include "tests/framework/datasets/Datasets.h"
@@ -44,124 +44,124 @@ namespace validation
 TEST_SUITE(NEON)
 TEST_SUITE(ReorgLayer)
 
-DATA_TEST_CASE(Configuration,
-               framework::DatasetMode::ALL,
-               combine(combine(combine(concat(datasets::SmallShapes(), datasets::LargeShapes()),
-                                       framework::dataset::make("Stride", { 2, 3 })),
-                               framework::dataset::make("DataType", { DataType::QASYMM8, DataType::F16, DataType::F32 })),
-                       framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })),
+// *INDENT-OFF*
+// clang-format off
+DATA_TEST_CASE(Validate, framework::DatasetMode::ALL, zip(zip(zip(
+               framework::dataset::make("InputInfo", { TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::S64),    // Wrong output tensor
+                                                       TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(10U, 12U, 1U, 2U), 1, DataType::F32),    // Wrong output tensor
+                                                       TensorInfo(TensorShape(3U, 12U, 4U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(3U, 12U, 4U, 2U), 1, DataType::F32),     // Wrong data type
+                                                     }),
+               framework::dataset::make("OutputInfo",{ TensorInfo(TensorShape(3U, 4U, 10U, 2U), 1, DataType::S64),
+                                                       TensorInfo(TensorShape(5U, 6U, 4U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(5U, 6U, 2, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(1U, 4U, 36U, 2U), 1, DataType::F32),
+                                                       TensorInfo(TensorShape(1U, 4U, 36U, 2U), 1, DataType::F16),
+                                                     })),
+               framework::dataset::make("Stride", { 2, 2, 4, 3 })),
+               framework::dataset::make("Expected", { false, true, false, true, false })),
+               input_info, output_info, stride, expected)
+{
+    bool status = bool(NEReorgLayer::validate(&input_info, &output_info, stride));
+    ARM_COMPUTE_EXPECT(status == expected, framework::LogLevel::ERRORS);
+}
+// clang-format on
+// *INDENT-ON*
+
+DATA_TEST_CASE(Configuration, framework::DatasetMode::ALL, combine(combine(framework::dataset::concat(datasets::SmallReorgLayerDataset(), datasets::LargeReorgLayerDataset()),
+                                                                           framework::dataset::make("DataType", { DataType::F32, DataType::F16, DataType::QASYMM8 })),
+                                                                   framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })),
                shape, stride, data_type, data_layout)
 {
+    // Permute the tensor shape in case of NHWC data layout
+    TensorShape shape_to_use = shape;
+    if(data_layout == DataLayout::NHWC)
+    {
+        permute(shape_to_use, PermutationVector(2U, 0U, 1U));
+    }
+
     // Create tensors
-    Tensor ref_src = create_tensor<Tensor>(shape, data_type, 1, QuantizationInfo(), data_layout);
+    Tensor src = create_tensor<Tensor>(shape_to_use, data_type, 1, QuantizationInfo(), data_layout);
     Tensor dst;
 
-    // Create and Configure function
-    NEReorgLayer reorg_func;
-    reorg_func.configure(&ref_src, &dst, stride);
+    ARM_COMPUTE_EXPECT(src.info()->is_resizable(), framework::LogLevel::ERRORS);
+
+    // Create and configure function
+    NEReorgLayer reorg_layer;
+
+    // Auto-initialize the output within the function
+    reorg_layer.configure(&src, &dst, stride);
 
     // Validate valid region
-    const ValidRegion valid_region = shape_to_valid_region(dst.info()->tensor_shape());
-    validate(dst.info()->valid_region(), valid_region);
-}
+    const ValidRegion src_valid_region = shape_to_valid_region(shape_to_use);
+    const ValidRegion dst_valid_region = shape_to_valid_region(dst.info()->tensor_shape());
+    validate(src.info()->valid_region(), src_valid_region);
+    validate(dst.info()->valid_region(), dst_valid_region);
 
-// *INDENT-OFF*
-// clang-format off
-DATA_TEST_CASE(Validate, framework::DatasetMode::ALL, zip(zip(zip(
-        framework::dataset::make("InputInfo",{
-                TensorInfo(TensorShape(8U, 8U, 5U, 3U), 1, DataType::U16),     // Invalid stride
-                TensorInfo(TensorShape(8U, 8U, 5U, 3U), 1, DataType::U16),     // Invalid output shape
-                TensorInfo(TensorShape(8U, 8U, 5U, 3U), 1, DataType::U16),     // valid
-        }),
-        framework::dataset::make("OutputInfo", {
-                TensorInfo(TensorShape(4U, 4U, 20U, 3U), 1, DataType::U16),
-                TensorInfo(TensorShape(4U, 4U, 10U, 3U), 1, DataType::U16),
-                TensorInfo(TensorShape(4U, 4U, 20U, 3U), 1, DataType::U16),
-        })),
-        framework::dataset::make("Stride", { -1, 2, 2 })),
-        framework::dataset::make("Expected", { false, false, true })),
-        input_info, output_info, stride, expected)
-{
-    Status status = NEReorgLayer::validate(&input_info.clone()->set_is_resizable(false), &output_info.clone()->set_is_resizable(false), stride);
-    ARM_COMPUTE_EXPECT(bool(status) == expected, framework::LogLevel::ERRORS);
+    // Validate padding
+    const int         step        = 1;
+    const PaddingSize src_padding = PaddingCalculator(shape_to_use.x(), step).required_padding();
+    const PaddingSize dst_padding = PaddingCalculator(dst.info()->tensor_shape().x(), step).required_padding();
+    validate(src.info()->padding(), src_padding);
+    validate(dst.info()->padding(), dst_padding);
 }
-// clang-format on
-// *INDENT-ON*
 
 template <typename T>
 using NEReorgLayerFixture = ReorgLayerValidationFixture<Tensor, Accessor, NEReorgLayer, T>;
 
-TEST_SUITE(U8)
-FIXTURE_DATA_TEST_CASE(RunSmall,
-                       NEReorgLayerFixture<uint8_t>,
-                       framework::DatasetMode::PRECOMMIT,
-                       combine(combine(combine(datasets::SmallShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U8)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+TEST_SUITE(S32)
+FIXTURE_DATA_TEST_CASE(RunSmall, NEReorgLayerFixture<int32_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                  DataType::S32)),
+                                                                                                          framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
 
-FIXTURE_DATA_TEST_CASE(RunLarge,
-                       NEReorgLayerFixture<uint8_t>,
-                       framework::DatasetMode::NIGHTLY,
-                       combine(combine(combine(datasets::LargeShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U8)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+FIXTURE_DATA_TEST_CASE(RunLarge, NEReorgLayerFixture<int32_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                DataType::S32)),
+                                                                                                        framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
-TEST_SUITE_END() // U8
-
-TEST_SUITE(U16)
-FIXTURE_DATA_TEST_CASE(RunSmall,
-                       NEReorgLayerFixture<uint16_t>,
-                       framework::DatasetMode::PRECOMMIT,
-                       combine(combine(combine(datasets::SmallShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U16)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+TEST_SUITE_END() // S32
+
+TEST_SUITE(S16)
+FIXTURE_DATA_TEST_CASE(RunSmall, NEReorgLayerFixture<int16_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                  DataType::S16)),
+                                                                                                          framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
 
-FIXTURE_DATA_TEST_CASE(RunLarge,
-                       NEReorgLayerFixture<uint16_t>,
-                       framework::DatasetMode::NIGHTLY,
-                       combine(combine(combine(datasets::LargeShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U16)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+FIXTURE_DATA_TEST_CASE(RunLarge, NEReorgLayerFixture<int16_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                DataType::S16)),
+                                                                                                        framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
-TEST_SUITE_END() // U16
-
-TEST_SUITE(U32)
-FIXTURE_DATA_TEST_CASE(RunSmall,
-                       NEReorgLayerFixture<uint32_t>,
-                       framework::DatasetMode::PRECOMMIT,
-                       combine(combine(combine(datasets::SmallShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U32)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+TEST_SUITE_END() // S16
+
+TEST_SUITE(S8)
+FIXTURE_DATA_TEST_CASE(RunSmall, NEReorgLayerFixture<int8_t>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallReorgLayerDataset(), framework::dataset::make("DataType",
+                                                                                                                 DataType::S8)),
+                                                                                                         framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
 
-FIXTURE_DATA_TEST_CASE(RunLarge,
-                       NEReorgLayerFixture<uint32_t>,
-                       framework::DatasetMode::NIGHTLY,
-                       combine(combine(combine(datasets::LargeShapes(), framework::dataset::make("Stride", { 2, 3 })),
-                                       framework::dataset::make("DataType", DataType::U32)),
-                               framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
+FIXTURE_DATA_TEST_CASE(RunLarge, NEReorgLayerFixture<int8_t>, framework::DatasetMode::NIGHTLY, combine(combine(datasets::LargeReorgLayerDataset(), framework::dataset::make("DataType", DataType::S8)),
+                                                                                                       framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })))
 {
     // Validate output
     validate(Accessor(_target), _reference);
 }
-TEST_SUITE_END() // U32
+TEST_SUITE_END() // S8
 
 TEST_SUITE_END() // ReorgLayer
 TEST_SUITE_END() // NEON
diff --git a/tests/validation/fixtures/ReorgLayerFixture.h b/tests/validation/fixtures/ReorgLayerFixture.h
index 2bc5a6fb8c..3300e0dd7c 100644
--- a/tests/validation/fixtures/ReorgLayerFixture.h
+++ b/tests/validation/fixtures/ReorgLayerFixture.h
@@ -59,9 +59,7 @@ protected:
 
     TensorType compute_target(TensorShape input_shape, int32_t stride, DataType data_type, DataLayout data_layout)
     {
-        // Check if indeed the input shape can be reshape to the output one
-        ARM_COMPUTE_EXPECT(stride >= 0, framework::LogLevel::ERRORS);
-
+        // Note: The input shape passed to the function is always in NCHW
         if(data_layout == DataLayout::NHWC)
         {
             permute(input_shape, PermutationVector(2U, 0U, 1U));
diff --git a/tests/validation/reference/ReorgLayer.cpp b/tests/validation/reference/ReorgLayer.cpp
index cb13a737e0..2eb5d01926 100644
--- a/tests/validation/reference/ReorgLayer.cpp
+++ b/tests/validation/reference/ReorgLayer.cpp
@@ -24,6 +24,7 @@
 #include "ReorgLayer.h"
 
 #include "arm_compute/core/Types.h"
+#include "arm_compute/core/utils/misc/ShapeCalculator.h"
 
 namespace arm_compute
 {
@@ -33,29 +34,17 @@ namespace validation
 {
 namespace reference
 {
-namespace
-{
-TensorShape compute_reorg_shape(const TensorShape &src_shape, int32_t stride)
-{
-    ARM_COMPUTE_ERROR_ON(stride <= 0);
-
-    TensorShape dst_shape = src_shape;
-    dst_shape.set(0, src_shape.x() / stride);
-    dst_shape.set(1, src_shape.y() / stride);
-    dst_shape.set(2, src_shape.z() * stride * stride);
-
-    return dst_shape;
-}
-} // namespace
-
 template <typename T>
 SimpleTensor<T> reorg_layer(const SimpleTensor<T> &src, int32_t stride)
 {
-    // Calculate output shape
-    const TensorShape dst_shape = compute_reorg_shape(src.shape(), stride);
+    ARM_COMPUTE_ERROR_ON(src.shape().num_dimensions() > 4);
+    ARM_COMPUTE_ERROR_ON(src.data_layout() != DataLayout::NCHW);
+
+    TensorInfo        input_info(src.shape(), 1, src.data_type());
+    const TensorShape output_shape = misc::shape_calculator::compute_reorg_output_shape(input_info, stride);
 
     // Create destination tensor
-    SimpleTensor<T> dst{ dst_shape, src.data_type() };
+    SimpleTensor<T> dst{ output_shape, src.data_type() };
 
     const unsigned int W          = dst.shape().x();
     const unsigned int H          = dst.shape().y();
@@ -88,9 +77,9 @@ SimpleTensor<T> reorg_layer(const SimpleTensor<T> &src, int32_t stride)
     return dst;
 }
 
-template SimpleTensor<uint8_t> reorg_layer(const SimpleTensor<uint8_t> &src, int32_t stride);
-template SimpleTensor<uint16_t> reorg_layer(const SimpleTensor<uint16_t> &src, int32_t stride);
-template SimpleTensor<uint32_t> reorg_layer(const SimpleTensor<uint32_t> &src, int32_t stride);
+template SimpleTensor<int32_t> reorg_layer(const SimpleTensor<int32_t> &src, int32_t stride);
+template SimpleTensor<int16_t> reorg_layer(const SimpleTensor<int16_t> &src, int32_t stride);
+template SimpleTensor<int8_t> reorg_layer(const SimpleTensor<int8_t> &src, int32_t stride);
 } // namespace reference
 } // namespace validation
 } // namespace test