Method run

src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp:115–171

}

void NESpaceToDepthLayerKernel::run(const Window &window, const ThreadInfo &info)
{
    ARM_COMPUTE_UNUSED(info);
    ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
    ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(ICPPKernel::window(), window);

    const int channel_idx  = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::CHANNEL);
    const int element_size = _input->info()->element_size();

    const size_t channel_size = _input->info()->dimension(channel_idx);

    Window slice_out = window.first_slice_window_3D();

    int batch_id = 0;

    // Main loop for NCHW and NHWC
    if (_data_layout == DataLayout::NCHW)
    {
        do
        {
            Iterator out(_output, slice_out);
            execute_window_loop(
                slice_out,
                [&](const Coordinates &id)
                {
                    const size_t channel_id = id.z();
                    const size_t in_x       = id.x() * _block_shape + (channel_id / channel_size) % _block_shape;
                    const size_t in_y       = id.y() * _block_shape + (channel_id / channel_size) / _block_shape;
                    const int    z          = channel_id % channel_size;
                    Coordinates  input_coords{in_x, in_y, z, batch_id};
                    memcpy(out.ptr(), _input->ptr_to_element(input_coords), element_size);
                },
                out);
            ++batch_id;
        } while (window.slide_window_slice_3D(slice_out));
    }
    else
    {
        do
        {
            Iterator out(_output, slice_out);
            execute_window_loop(
                slice_out,
                [&](const Coordinates &id)
                {
                    const size_t channel_id = id.x();
                    const size_t in_x       = id.y() * _block_shape + (channel_id / channel_size) % _block_shape;
                    const size_t in_y       = id.z() * _block_shape + (channel_id / channel_size) / _block_shape;
                    const int    z          = channel_id % channel_size;
                    Coordinates  input_coords{z, in_x, in_y, batch_id};
                    memcpy(out.ptr(), _input->ptr_to_element(input_coords), element_size);
                },
                out);
            ++batch_id;
        } while (window.slide_window_slice_3D(slice_out));
    }
}
} // namespace arm_compute