Method gather_common

src/core/NEON/kernels/NEGatherKernel.cpp:86–145

template <typename TIndex>
void NEGatherKernel::gather_common(const Window &window, const ThreadInfo &info)
{
    ARM_COMPUTE_UNUSED(info);

    auto dst_win = window;

    const auto src_info = _input->info();
    const auto idx_info = _indices->info();
    const auto dst_info = _output->info();

    const auto num_dims     = dst_info->num_dimensions();
    const auto chunk_stride = src_info->strides_in_bytes()[_axis];

    const auto window_start_x = window.x().start();
    const auto window_end_x   = window.x().end();
    auto       window_size_x  = src_info->element_size();

    const auto idx_limit = static_cast<TIndex>(src_info->tensor_shape()[_axis]);

    if (_axis != 0)
    {
        dst_win.set(0, Window::Dimension(window_start_x, window_start_x + 1, 1));
        window_size_x *= window_end_x - window_start_x;
    }

    // Compute source and index tensors window based on the output window.
    auto   src_win = dst_win;
    Window idx_win;

    for (size_t i = 0; i < idx_info->num_dimensions(); ++i)
    {
        src_win.set(_axis + i, Window::Dimension(0, 1, 1));
        idx_win.set(_axis + i, window[_axis + i]);
    }

    // Use the custom strides to access all three tensors using the same loop.
    Iterator src_it(num_dims, _src_it_strides, _input->buffer(), src_info->offset_first_element_in_bytes(), src_win);
    Iterator idx_it(num_dims, _idx_it_strides, _indices->buffer(), idx_info->offset_first_element_in_bytes(), idx_win);
    Iterator dst_it(num_dims, dst_info->strides_in_bytes(), _output->buffer(),
                    dst_info->offset_first_element_in_bytes(), dst_win);

    execute_window_loop(
        dst_win,
        [&](const Coordinates &)
        {
            const auto idx = *reinterpret_cast<const TIndex *>(idx_it.ptr());

            if (idx >= 0 && idx < idx_limit)
            {
                const auto src_ptr = src_it.ptr() + idx * chunk_stride;

                std::copy_n(src_ptr, window_size_x, dst_it.ptr());
            }
            else
            {
                std::fill_n(dst_it.ptr(), window_size_x, 0);
            }
        },