feat(core): add contiguous packing support for Views of arbitrary rank

src/KokkosComm/impl/contiguous.hpp

@@ -3,7 +3,9 @@
 
 #pragma once
 
+#include <cstddef>
 #include <string>
+#include <utility>
 
 #include <Kokkos_Core.hpp>
 
@@ -21,31 +23,37 @@ struct contiguous_view {
 template <KokkosView View>
 using contiguous_view_t = contiguous_view<View>::type;
 
-template <KokkosView View, KokkosExecutionSpace Space>
-auto allocate_contiguous_for(const Space &space, const std::string &label, View &v) {
-  using non_const_packed_view_type = contiguous_view_t<View>;
-
-  if constexpr (KokkosComm::rank<View>() == 1) {
-    return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0));
-  } else if constexpr (KokkosComm::rank<View>() == 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<View>, "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 <KokkosExecutionSpace Exec, KokkosView View>
+auto allocate_contiguous_for(const Exec& exec, const std::string& label, const View& v) -> contiguous_view_t<View> {
+  // Unpack `v` extents into the `ContigView` constructor
+  return [&label, &exec, &v ]<size_t... Is>(std::index_sequence<Is...>) {
+    return contiguous_view_t<View>(Kokkos::view_alloc(exec, Kokkos::WithoutInitializing, label), v.extent(Is)...);
   }
+  (std::make_index_sequence<rank<View>()>{});
 }
 
-template <KokkosExecutionSpace Space, KokkosView DstView, KokkosView SrcView>
-auto resize_contiguous_for(const Space &space, DstView &out, const SrcView &in) {
-  static_assert(DstView::rank == SrcView::rank, "");
-
-  if constexpr (KokkosComm::rank<DstView>() == 1) {
-    Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0));
-  } else if constexpr (KokkosComm::rank<DstView>() == 2) {
-    Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0), in.extent(1));
-  } else {
-    static_assert(std::is_void_v<DstView>, "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 <KokkosExecutionSpace Exec, KokkosView DstV, KokkosView SrcV>
+auto resize_contiguous_for(const Exec& exec, const DstV& dst, const SrcV& src) -> void {
+  static_assert(rank<DstV>() == rank<SrcV>(),
+                "KokkosComm::Impl::resize_contiguous_for: source and destination Views must have the same rank");
+  // Unpack `src` extents into `realloc` call
+  [&exec, &dst, &src ]<size_t... Is>(std::index_sequence<Is...>) {
+    Kokkos::realloc(Kokkos::view_alloc(exec, Kokkos::WithoutInitializing), dst, src.extent(Is)...);
   }
+  (std::make_index_sequence<rank<DstV>()>{});
 }
 
 }  // namespace KokkosComm::Impl