slice.hpp

Go to the documentation of this file.

#ifndef KOKKOS_UTILS_VIEW_SLICE_HPP
#define KOKKOS_UTILS_VIEW_SLICE_HPP
 
#include "kokkos-utils/concepts/View.hpp"

 
namespace Kokkos::utils::view
{
 
namespace impl
{
 
template <typename ViewType, typename... Indices, size_t... AllsIndices>
KOKKOS_FUNCTION
constexpr auto slice(ViewType&& view, Indices&&... indices, std::index_sequence<AllsIndices...>)
{
    return Kokkos::subview(
        std::forward<ViewType>(view),
        std::forward<Indices>(indices)...,
        std::get<AllsIndices>(std::array<Kokkos::ALL_t, sizeof...(AllsIndices)>{})...
    );
}
 
} // namespace impl

template <size_t Rank, typename ViewType, typename... Indices>
requires concepts::ViewOfRank<std::remove_reference_t<ViewType>, Rank>
KOKKOS_FUNCTION
constexpr auto slice(ViewType&& view, Indices&&... indices)
{
    constexpr auto size = Rank - sizeof...(Indices);
    return impl::slice<ViewType, Indices...>(
        std::forward<ViewType>(view),
        std::forward<Indices>(indices)...,
        std::make_index_sequence<size>{}
    );
}

template <typename ViewType, typename... Indices>
requires concepts::ViewOfRank<std::remove_reference_t<ViewType>, sizeof...(Indices)>
KOKKOS_FUNCTION
constexpr auto slice(ViewType&& view, Indices&&... indices)
{
    return slice<std::remove_reference_t<ViewType>::rank()>(
        std::forward<ViewType>(view),
        std::forward<Indices>(indices)...
    );
}
 
} // namespace Kokkos::utils::view
 
#endif // KOKKOS_UTILS_VIEW_SLICE_HPP