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>

1 files changed, 45 insertions, 12 deletions

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