Revert "Add support for 2d and 3d indices for axis 0"

This reverts commit 0db8b8bbd941b3dab4238c03e734e7ac43c662ed.

Relates to [COMPMID-5055]

Signed-off-by: Mohammed Suhail Munshi <MohammedSuhail.Munshi@arm.com>
Change-Id: I143e7965e21b956abb05ba5c41e12c5b73b7345a

6 files changed, 52 insertions, 176 deletions

diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index aa51ad209a..df907c106e 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -1496,24 +1496,13 @@ inline TensorShape compute_pool3d_shape(const TensorShape &src, Pooling3dLayerIn
 
 inline TensorShape compute_gather_shape(const TensorShape &input_shape, const TensorShape &indices_shape, uint32_t actual_axis)
 {
+    ARM_COMPUTE_ERROR_ON(indices_shape.num_dimensions() > 1);
     ARM_COMPUTE_ERROR_ON(input_shape.num_dimensions() > 4);
     ARM_COMPUTE_ERROR_ON(actual_axis >= input_shape.num_dimensions());
 
-    TensorShape output_shape = input_shape;
-    if(indices_shape.num_dimensions() == 1u)
-    {
-        output_shape[actual_axis] = indices_shape[0];
-    }
-    else
-    {
-        const auto inddims{ indices_shape.num_dimensions() };
-        output_shape.shift_right(indices_shape.num_dimensions() - 1);
-        output_shape[0] = input_shape[0];
-        for(size_t idx(1); (idx - 1) < inddims; ++idx)
-        {
-            output_shape.set(actual_axis + idx, indices_shape[idx - 1], false);
-        }
-    }
+    TensorShape output_shape  = input_shape;
+    output_shape[actual_axis] = indices_shape[0];
+
     return output_shape;
 }
 } // namespace shape_calculator
diff --git a/src/core/NEON/kernels/NEGatherKernel.cpp b/src/core/NEON/kernels/NEGatherKernel.cpp
index 8d86a22b7e..7090da8015 100644
--- a/src/core/NEON/kernels/NEGatherKernel.cpp
+++ b/src/core/NEON/kernels/NEGatherKernel.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022 Arm Limited.
+ * Copyright (c) 2019-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -47,21 +47,16 @@ namespace
 template <typename U>
 void validate_indices(const ITensor *indices)
 {
-    auto *indices_ptr = (reinterpret_cast<const U *>(  indices->buffer() + indices->info()->offset_first_element_in_bytes() ));
-    for(size_t i = 0; i < indices->info()->total_size(); ++i)
+    for(size_t i = 0; i < indices->info()->tensor_shape()[0]; ++i)
     {
-        const U index_value = indices_ptr[i];
-        ARM_COMPUTE_UNUSED(index_value);
-        if(index_value < 0)
-        {
-            ARM_COMPUTE_ERROR_ON(index_value < 0);
-        }
+        ARM_COMPUTE_ERROR_ON(*(reinterpret_cast<U *>(indices->ptr_to_element(Coordinates(i)))) < 0);
     }
 }
 
 Status validate_arguments(const ITensorInfo *input, const ITensorInfo *indices, const ITensorInfo *output, int axis)
 {
     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, indices, output);
+    ARM_COMPUTE_RETURN_ERROR_ON(indices->num_dimensions() > 1);
     ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
 
     if(axis < 0)
@@ -70,7 +65,6 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *indices,
     }
 
     ARM_COMPUTE_RETURN_ERROR_ON(0 > axis || axis >= static_cast<int32_t>(input->num_dimensions()));
-    ARM_COMPUTE_RETURN_ERROR_ON(axis != 0 && indices->num_dimensions() > 1);
     ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
 
     if(output->total_size() != 0)
@@ -93,40 +87,6 @@ NEGatherKernel::NEGatherKernel()
 }
 
 template <typename U>
-inline void NEGatherKernel::gather_dims_0_axis(const Window &window, const ThreadInfo &info)
-{
-    ARM_COMPUTE_UNUSED(info);
-    ARM_COMPUTE_ERROR_ON(_indices->info()->num_dimensions() < 2);
-    validate_indices<U>(_indices);
-
-    Window output_window{ window };
-    output_window.set(Window::DimX, Window::Dimension(0, 1, 1));
-    Iterator output_it(_output, output_window);
-
-    const uint8_t *const in_ptr_start    = _input->buffer() + _input->info()->offset_first_element_in_bytes();
-    const uint32_t       input_stride_y  = _input->info()->strides_in_bytes()[1];
-    const uint32_t       output_stride_y = _output->info()->strides_in_bytes()[1];
-
-    const U *const dex_ptr_start = reinterpret_cast<const U *const>(_indices->buffer() + _indices->info()->offset_first_element_in_bytes());
-    execute_window_loop(output_window, [&](const Coordinates & id)
-    {
-        const auto        new_index = *(dex_ptr_start + id.y() + id.z() * _output->info()->tensor_shape()[1] + id[3] * _indices->info()->tensor_shape()[1] * _indices->info()->tensor_shape()[0]);
-        U                *out_ptr   = reinterpret_cast<U *>(output_it.ptr());
-        const char *const in_ptr    = reinterpret_cast<const char *const>(in_ptr_start + new_index * input_stride_y);
-        memcpy(out_ptr, in_ptr, output_stride_y);
-    },
-    output_it);
-}
-
-template <typename U>
-inline void NEGatherKernel::gather_dims_n_axis(const Window &window, const ThreadInfo &info)
-{
-    ARM_COMPUTE_UNUSED(info);
-    ARM_COMPUTE_UNUSED(window);
-    ARM_COMPUTE_ERROR("NOT_SUPPORTED!");
-}
-
-template <typename U>
 inline void NEGatherKernel::gather_0_axis(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_UNUSED(info);
@@ -187,75 +147,38 @@ void NEGatherKernel::configure(const ITensor *input, const ITensor *indices, ITe
     }
     ARM_COMPUTE_ERROR_ON(0 > _axis || _axis >= static_cast<int32_t>(input->info()->num_dimensions()));
 
-    if(indices->info()->num_dimensions() == 1u)
+    if(0 == _axis)
     {
-        if(0 == _axis)
-        {
-            switch(_indices->info()->data_type())
-            {
-                case DataType::U32:
-                    _func = &NEGatherKernel::gather_0_axis<uint32_t>;
-                    break;
-                case DataType::S32:
-                    _func = &NEGatherKernel::gather_0_axis<int32_t>;
-                    break;
-                default:
-                    ARM_COMPUTE_ERROR("Not supported");
-                    break;
-            }
-        }
-        else
+        switch(_indices->info()->data_type())
         {
-            switch(_indices->info()->data_type())
-            {
-                case DataType::U32:
-                    _func = &NEGatherKernel::gather_n_axis<uint32_t>;
-                    break;
-                case DataType::S32:
-                    _func = &NEGatherKernel::gather_n_axis<int32_t>;
-                    break;
-                default:
-                    ARM_COMPUTE_ERROR("Not supported");
-                    break;
-            }
+            case DataType::U32:
+                _func = &NEGatherKernel::gather_0_axis<uint32_t>;
+                break;
+            case DataType::S32:
+                _func = &NEGatherKernel::gather_0_axis<int32_t>;
+                break;
+            default:
+                ARM_COMPUTE_ERROR("Not supported");
+                break;
         }
     }
     else
     {
-        if(0 == _axis)
+        switch(_indices->info()->data_type())
         {
-            switch(_indices->info()->data_type())
-            {
-                case DataType::U32:
-                    _func = &NEGatherKernel::gather_dims_0_axis<uint32_t>;
-                    break;
-                case DataType::S32:
-                    _func = &NEGatherKernel::gather_dims_0_axis<int32_t>;
-                    break;
-                default:
-                    ARM_COMPUTE_ERROR("Not supported");
-                    break;
-            }
-        }
-        else
-        {
-            switch(_indices->info()->data_type())
-            {
-                case DataType::U32:
-                    _func = &NEGatherKernel::gather_dims_n_axis<uint32_t>;
-                    break;
-                case DataType::S32:
-                    _func = &NEGatherKernel::gather_dims_n_axis<int32_t>;
-                    break;
-                default:
-                    ARM_COMPUTE_ERROR("Not supported");
-                    break;
-            }
+            case DataType::U32:
+                _func = &NEGatherKernel::gather_n_axis<uint32_t>;
+                break;
+            case DataType::S32:
+                _func = &NEGatherKernel::gather_n_axis<int32_t>;
+                break;
+            default:
+                ARM_COMPUTE_ERROR("Not supported");
+                break;
         }
     }
-
     // Output auto initialization if not yet initialized
-    const TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->info()->tensor_shape(), indices->info()->tensor_shape(), _axis);
+    TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->info()->tensor_shape(), indices->info()->tensor_shape(), _axis);
     auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape));
 
     // Create window
diff --git a/src/core/NEON/kernels/NEGatherKernel.h b/src/core/NEON/kernels/NEGatherKernel.h
index fc0e67854b..0711f8190b 100644
--- a/src/core/NEON/kernels/NEGatherKernel.h
+++ b/src/core/NEON/kernels/NEGatherKernel.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022 Arm Limited.
+ * Copyright (c) 2019-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -85,23 +85,18 @@ private:
      *
      * For gather on the 0 axis an element by element copy is performed.
      *
-     * @param[in] window Region on which to run the kernel. (Must be a region of the window returned by window())
-     * @param[in] info   Info about running thread and CPU.
+     * @param[in] window Region on which to execute the kernel. (Must be a region of the window returned by window())
+     * @param[in] info   Info about executing thread and CPU.
      */
     template <typename U>
     void gather_0_axis(const Window &window, const ThreadInfo &info);
 
-    template <typename U>
-    void gather_dims_0_axis(const Window &window, const ThreadInfo &info);
-
-    template <typename U>
-    void gather_dims_n_axis(const Window &window, const ThreadInfo &info);
     /** Implementation of the gather operation.
      *
      * For 1<=axis a row-wise copy is taking place.
      *
-     * @param[in] window Region on which to run the kernel. (Must be a region of the window returned by window())
-     * @param[in] info   Info about running thread and CPU.
+     * @param[in] window Region on which to execute the kernel. (Must be a region of the window returned by window())
+     * @param[in] info   Info about executing thread and CPU.
      */
     template <typename U>
     void gather_n_axis(const Window &window, const ThreadInfo &info);
diff --git a/tests/datasets/GatherDataset.h b/tests/datasets/GatherDataset.h
index f4ad7c9497..29a99d5239 100644
--- a/tests/datasets/GatherDataset.h
+++ b/tests/datasets/GatherDataset.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2019, 2022 Arm Limited.
+ * Copyright (c) 2018-2019 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -106,18 +106,6 @@ private:
     std::vector<int>         _axis{};
 };
 
-
-class SmallGatherMultiDimIndicesDataset final : public GatherDataset
-{
-public:
-    SmallGatherMultiDimIndicesDataset()
-    {
-        add_config(TensorShape(15U, 15U), TensorShape(4U, 13U, 2U), 0);
-        add_config(TensorShape(15U, 15U), TensorShape(2U, 12U), 0);
-    }
-};
-
-
 class SmallGatherDataset final : public GatherDataset
 {
 public:
diff --git a/tests/validation/NEON/Gather.cpp b/tests/validation/NEON/Gather.cpp
index 1ab2668aad..ca1e166bd1 100644
--- a/tests/validation/NEON/Gather.cpp
+++ b/tests/validation/NEON/Gather.cpp
@@ -100,14 +100,12 @@ DATA_TEST_CASE(Validate, framework::DatasetMode::ALL, zip(zip(zip(zip(
 template <typename T>
 using NEGatherFixture = GatherFixture<Tensor, Accessor, NEGather, T>;
 
-const auto gather_small_shapes = arm_compute::test::framework::dataset::concat(datasets::SmallGatherDataset(),datasets::SmallGatherMultiDimIndicesDataset());
-
 TEST_SUITE(Float)
 TEST_SUITE(FP16)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<half>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine( gather_small_shapes, framework::dataset::make("DataType", DataType::F16)))
+                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::F16)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -127,7 +125,7 @@ TEST_SUITE(FP32)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<float>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::F32)))
+                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::F32)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -148,7 +146,7 @@ TEST_SUITE(U8)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<uint8_t>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::U8)))
+                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::U8)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -168,7 +166,7 @@ TEST_SUITE(U16)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<uint16_t>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::U16)))
+                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::U16)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
diff --git a/tests/validation/reference/Gather.cpp b/tests/validation/reference/Gather.cpp
index 02292fb74e..93ac09cf95 100644
--- a/tests/validation/reference/Gather.cpp
+++ b/tests/validation/reference/Gather.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2019, 2022 Arm Limited.
+ * Copyright (c) 2018-2019 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -45,39 +45,22 @@ SimpleTensor<T> gather(const SimpleTensor<T> &src, const SimpleTensor<uint32_t>
 
     Window win;
     win.use_tensor_dimensions(dst_shape);
-    if(indices.shape().num_dimensions() == 1u)
+    execute_window_loop(win, [&](const Coordinates & id)
     {
-        execute_window_loop(win, [&](const Coordinates & id)
+        Coordinates offset;
+        for(unsigned int dim = 0; dim < id.num_dimensions(); ++dim)
         {
-            Coordinates offset;
-            for(unsigned int dim = 0; dim < id.num_dimensions(); ++dim)
+            if(dim == actual_axis)
             {
-                if(dim == actual_axis)
-                {
-                    offset.set(dim, indices_ptr[id[dim]]);
-                }
-                else
-                {
-                    offset.set(dim, id[dim]);
-                }
+                offset.set(dim, indices_ptr[id[dim]]);
             }
-            *reinterpret_cast<T *>(dst(id)) = *reinterpret_cast<const T *>(src(offset));
-        });
-    }
-    else
-    {
-        if(actual_axis == 0)
-        {
-            win.set(Window::DimX, Window::Dimension(0, 1, 1));
-            uint32_t index = 0;
-            execute_window_loop(win, [&](const Coordinates & id)
+            else
             {
-                auto     *dst_ptr     = reinterpret_cast<T *>(dst(id));
-                const int row_to_copy = indices[index++];
-                std::copy_n(src.data() + row_to_copy * src.shape()[0], src.shape()[0], dst_ptr);
-            });
+                offset.set(dim, id[dim]);
+            }
         }
-    }
+        *reinterpret_cast<T *>(dst(id)) = *reinterpret_cast<const T *>(src(offset));
+    });
 
     return dst;
 }
@@ -89,4 +72,4 @@ template SimpleTensor<uint8_t> gather(const SimpleTensor<uint8_t> &src, const Si
 } // namespace reference
 } // namespace validation
 } // namespace test
-} // namespace arm_compute
+} // namespace arm_compute
\ No newline at end of file