Add support for 2d and 3d indices for axis 1

* Resolves COMPMID-5055

Change-Id: I2d14de29d3ec913d20c971bc8bbc9ad71e2d998f
Signed-off-by: Pablo Marquez Tello <pablo.tello@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/7547
Reviewed-by: SiCong Li <sicong.li@arm.com>
Reviewed-by: Michalis Spyrou <michalis.spyrou@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>

3 files changed, 65 insertions, 17 deletions

diff --git a/tests/datasets/GatherDataset.h b/tests/datasets/GatherDataset.h
index 444b62c052..8fec5441b1 100644
--- a/tests/datasets/GatherDataset.h
+++ b/tests/datasets/GatherDataset.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022 Arm Limited.
+ * Copyright (c) 2018-2019, 2022 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -106,6 +106,19 @@ private:
     std::vector<int>         _axis{};
 };
 
+class SmallGatherMultiDimIndicesDataset final : public GatherDataset
+{
+public:
+    SmallGatherMultiDimIndicesDataset()
+    {
+        add_config(TensorShape(2U, 6U), TensorShape(4U, 9U), 1);
+        add_config(TensorShape(15U, 15U), TensorShape(3U, 2U, 2U), 1);
+        add_config(TensorShape(15U, 15U), TensorShape(2U, 11U), 1);
+        add_config(TensorShape(5U, 3U, 4U), TensorShape(2U, 7U), 1);
+        add_config(TensorShape(1U, 5U, 3U), TensorShape(1U, 7U, 3U), 1);
+    }
+};
+
 class SmallGatherDataset final : public GatherDataset
 {
 public:
diff --git a/tests/validation/NEON/Gather.cpp b/tests/validation/NEON/Gather.cpp
index 71f98ea627..0aea19939e 100644
--- a/tests/validation/NEON/Gather.cpp
+++ b/tests/validation/NEON/Gather.cpp
@@ -100,12 +100,14 @@ 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(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::F16)))
+                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::F16)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -125,7 +127,7 @@ TEST_SUITE(FP32)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<float>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::F32)))
+                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::F32)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -146,7 +148,7 @@ TEST_SUITE(U8)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<uint8_t>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::U8)))
+                       combine(gather_small_shapes, framework::dataset::make("DataType", DataType::U8)))
 {
     // Validate output
     validate(Accessor(_target), _reference);
@@ -166,7 +168,7 @@ TEST_SUITE(U16)
 FIXTURE_DATA_TEST_CASE(RunSmall,
                        NEGatherFixture<uint16_t>,
                        framework::DatasetMode::PRECOMMIT,
-                       combine(datasets::SmallGatherDataset(), framework::dataset::make("DataType", DataType::U16)))
+                       combine(gather_small_shapes, 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 c26438857f..8de1a473eb 100644
--- a/tests/validation/reference/Gather.cpp
+++ b/tests/validation/reference/Gather.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022 Arm Limited.
+ * Copyright (c) 2018-2019, 2022 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -45,22 +45,55 @@ SimpleTensor<T> gather(const SimpleTensor<T> &src, const SimpleTensor<uint32_t>
 
     Window win;
     win.use_tensor_dimensions(dst_shape);
-    execute_window_loop(win, [&](const Coordinates & id)
+    if(indices.shape().num_dimensions() == 1u)
     {
-        Coordinates offset;
-        for(unsigned int dim = 0; dim < id.num_dimensions(); ++dim)
+        execute_window_loop(win, [&](const Coordinates & id)
         {
-            if(dim == actual_axis)
+            Coordinates offset;
+            for(unsigned int dim = 0; dim < id.num_dimensions(); ++dim)
             {
-                offset.set(dim, indices_ptr[id[dim]]);
+                if(dim == actual_axis)
+                {
+                    offset.set(dim, indices_ptr[id[dim]]);
+                }
+                else
+                {
+                    offset.set(dim, id[dim]);
+                }
             }
-            else
+            *reinterpret_cast<T *>(dst(id)) = *reinterpret_cast<const T *>(src(offset));
+        });
+    }
+    else
+    {
+        if(actual_axis == 1)
+        {
+            win.set(Window::DimX, Window::Dimension(0, 1, 1));
+            execute_window_loop(win, [&](const Coordinates & id)
             {
-                offset.set(dim, id[dim]);
-            }
+                auto       *dst_ptr = dst(id);
+                Coordinates index_offset;
+                for(uint32_t k = 0; k < indices.shape().num_dimensions(); ++k)
+                {
+                    index_offset.set(k, id[k + 1]);
+                }
+                const uint32_t row = *reinterpret_cast<const uint32_t *>(indices(index_offset));
+                Coordinates    src_offset;
+                src_offset.set(0, 0);
+                src_offset.set(1, row);
+                for(uint32_t j = 2; j < src.shape().num_dimensions(); ++j)
+                {
+                    src_offset.set(j, id[1 + indices.shape().num_dimensions() + (j - 2)]);
+                }
+                const auto in_ptr_row = src(src_offset);
+                memcpy(dst_ptr, in_ptr_row, src.shape()[0] * src.element_size());
+            });
+        }
+        else
+        {
+            ARM_COMPUTE_ERROR("Not implemented.");
         }
-        *reinterpret_cast<T *>(dst(id)) = *reinterpret_cast<const T *>(src(offset));
-    });
+    }
 
     return dst;
 }
@@ -72,4 +105,4 @@ template SimpleTensor<uint8_t> gather(const SimpleTensor<uint8_t> &src, const Si
 } // namespace reference
 } // namespace validation
 } // namespace test
-} // namespace arm_compute
\ No newline at end of file
+} // namespace arm_compute