Apply to_array helper (#171)

* Apply to_array helper * fix: ToArray tests for Kokkos v.4.2.0 * fix: to_array should not modify std::array * Update common/src/KokkosFFT_utils.hpp Remove unnecessary variadic integers Co-authored-by: Damien L-G <[email protected]> * use std::size_t in place of size_t * Make ToArray tests Kokkos 4.2 compatible * Align to_array implementation with Kokkos implementation * Update ToArray test with Kokkos 4.4 * Make rvalue test a run time test * make tests uniform based on a review --------- Co-authored-by: Yuuichi Asahi <[email protected]> Co-authored-by: Thomas Padioleau <[email protected]> Co-authored-by: Damien L-G <[email protected]>
kokkos · Oct 15, 2024 · 1ce1313 · 1ce1313
1 parent da02175
commit 1ce1313
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 9 deletions.
diff --git a/common/src/KokkosFFT_transpose.hpp b/common/src/KokkosFFT_transpose.hpp
@@ -132,7 +132,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
       tile_type{{4, 4, 4}}  // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, 3> map = {_map[0], _map[1], _map[2]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range, KOKKOS_LAMBDA(int i0, int i1, int i2) {
         int _indices[rank] = {i0, i1, i2};
@@ -164,7 +164,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
                    // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, rank> map = {_map[0], _map[1], _map[2], _map[3]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range,
       KOKKOS_LAMBDA(int i0, int i1, int i2, int i3) {
@@ -199,7 +199,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
                    // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, rank> map = {_map[0], _map[1], _map[2], _map[3], _map[4]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range,
       KOKKOS_LAMBDA(int i0, int i1, int i2, int i3, int i4) {
@@ -236,8 +236,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
                    // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, rank> map = {_map[0], _map[1], _map[2],
-                                  _map[3], _map[4], _map[5]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range,
       KOKKOS_LAMBDA(int i0, int i1, int i2, int i3, int i4, int i5) {
@@ -275,8 +274,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
                    // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, rank> map = {_map[0], _map[1], _map[2], _map[3],
-                                  _map[4], _map[5], _map[6]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range,
       KOKKOS_LAMBDA(int i0, int i1, int i2, int i3, int i4, int i5) {
@@ -320,8 +318,7 @@ void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in,
                    // [TO DO] Choose optimal tile sizes for each device
   );
 
-  Kokkos::Array<int, rank> map = {_map[0], _map[1], _map[2], _map[3],
-                                  _map[4], _map[5], _map[6], _map[7]};
+  Kokkos::Array<int, rank> map = to_array(_map);
   Kokkos::parallel_for(
       "KokkosFFT::transpose", range,
       KOKKOS_LAMBDA(int i0, int i1, int i2, int i3, int i4, int i5) {

diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp
@@ -215,6 +215,26 @@ Layout create_layout(const std::array<int, N>& extents) {
   return layout;
 }
 
+template <typename T, std::size_t N, std::size_t... Is>
+constexpr Kokkos::Array<std::remove_cv_t<T>, N> to_array_lvalue(
+    std::array<T, N>& a, std::index_sequence<Is...>) {
+  return {{a[Is]...}};
+}
+template <typename T, std::size_t N, std::size_t... Is>
+constexpr Kokkos::Array<std::remove_cv_t<T>, N> to_array_rvalue(
+    std::array<T, N>&& a, std::index_sequence<Is...>) {
+  return {{std::move(a[Is])...}};
+}
+
+template <typename T, std::size_t N>
+constexpr Kokkos::Array<T, N> to_array(std::array<T, N>& a) {
+  return to_array_lvalue(a, std::make_index_sequence<N>());
+}
+template <typename T, std::size_t N>
+constexpr Kokkos::Array<T, N> to_array(std::array<T, N>&& a) {
+  return to_array_rvalue(std::move(a), std::make_index_sequence<N>());
+}
+
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/common/unit_test/Test_Utils.cpp b/common/unit_test/Test_Utils.cpp
@@ -541,3 +541,12 @@ TEST(IndexSequence, 3Dto5D) {
   EXPECT_EQ(default_axes1, ref_axes1);
   EXPECT_EQ(default_axes2, ref_axes2);
 }
+
+TEST(ToArray, lvalue) {
+  std::array arr{1, 2, 3};
+  ASSERT_EQ(KokkosFFT::Impl::to_array(arr), (Kokkos::Array{1, 2, 3}));
+}
+
+TEST(ToArray, rvalue) {
+  ASSERT_EQ(KokkosFFT::Impl::to_array(std::array{1, 2}), (Kokkos::Array{1, 2}));
+}