diff --git a/src/KokkosComm/impl/contiguous.hpp b/src/KokkosComm/impl/contiguous.hpp index c4374b9b..cf81f266 100644 --- a/src/KokkosComm/impl/contiguous.hpp +++ b/src/KokkosComm/impl/contiguous.hpp @@ -3,7 +3,9 @@ #pragma once +#include #include +#include #include @@ -21,31 +23,37 @@ struct contiguous_view { template using contiguous_view_t = contiguous_view::type; -template -auto allocate_contiguous_for(const Space &space, const std::string &label, View &v) { - using non_const_packed_view_type = contiguous_view_t; - - if constexpr (KokkosComm::rank() == 1) { - return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0)); - } else if constexpr (KokkosComm::rank() == 2) { - return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0), - v.extent(1)); - } else { - static_assert(std::is_void_v, "allocate_contiguous_for for views > rank 2 not implemented"); +/// @brief Allocate a contiguous View suitable for packing a non-contiguous View. +/// @tparam Exec A Kokkos Execution Space type. +/// @tparam View A Kokkos View type. +/// @param exec The execution space instance in which to perform the view allocation. +/// @param v The View to make a suitable contiguous allocation for. +/// @param label The label to give to the allocated contiguous View. Defaults to "contiguous_view". +template +auto allocate_contiguous_for(const Exec& exec, const std::string& label, const View& v) -> contiguous_view_t { + // Unpack `v` extents into the `ContigView` constructor + return [&label, &exec, &v ](std::index_sequence) { + return contiguous_view_t(Kokkos::view_alloc(exec, Kokkos::WithoutInitializing, label), v.extent(Is)...); } + (std::make_index_sequence()>{}); } -template -auto resize_contiguous_for(const Space &space, DstView &out, const SrcView &in) { - static_assert(DstView::rank == SrcView::rank, ""); - - if constexpr (KokkosComm::rank() == 1) { - Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0)); - } else if constexpr (KokkosComm::rank() == 2) { - Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0), in.extent(1)); - } else { - static_assert(std::is_void_v, "realloc_contiguous_for for views > rank 2 not implemented"); +/// @brief Resize a View into a contiguous one suitable for packing a non-contiguous View. +/// @tparam Exec A Kokkos Execution Space type. +/// @tparam DstV The Kokkos View type of the destination View to resize. +/// @tparam DstV The Kokkos View type of the source View to resize for. +/// @param exec The execution space instance in which to perform the view reallocation. +/// @param dst The View to resize. +/// @param src The View to make a suitable contiguous resize for. +template +auto resize_contiguous_for(const Exec& exec, const DstV& dst, const SrcV& src) -> void { + static_assert(rank() == rank(), + "KokkosComm::Impl::resize_contiguous_for: source and destination Views must have the same rank"); + // Unpack `src` extents into `realloc` call + [&exec, &dst, &src ](std::index_sequence) { + Kokkos::realloc(Kokkos::view_alloc(exec, Kokkos::WithoutInitializing), dst, src.extent(Is)...); } + (std::make_index_sequence()>{}); } } // namespace KokkosComm::Impl