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
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/7558
Reviewed-by: Pablo Marquez Tello <pablo.tello@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
diff --git a/src/core/NEON/kernels/NEGatherKernel.cpp b/src/core/NEON/kernels/NEGatherKernel.cpp
index 8d86a22..55c4525 100644
--- a/src/core/NEON/kernels/NEGatherKernel.cpp
+++ b/src/core/NEON/kernels/NEGatherKernel.cpp
@@ -47,21 +47,16 @@
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 @@
}
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 @@
}
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 @@
}
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())
{
- 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())
- {
- 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 fc0e678..6f00ddb 100644
--- a/src/core/NEON/kernels/NEGatherKernel.h
+++ b/src/core/NEON/kernels/NEGatherKernel.h
@@ -85,23 +85,18 @@
*
* 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);