From 1ce131355028e149444441ef54f8290bf239d1e4 Mon Sep 17 00:00:00 2001
From: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com>
Date: Tue, 15 Oct 2024 15:23:54 +0800
Subject: [PATCH] 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 <dalg24+github@gmail.com>

* 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 <y.asahi@nr.titech.ac.jp>
Co-authored-by: Thomas Padioleau <thomas.padioleau@cea.fr>
Co-authored-by: Damien L-G <dalg24+github@gmail.com>
---
 common/src/KokkosFFT_transpose.hpp | 15 ++++++---------
 common/src/KokkosFFT_utils.hpp     | 20 ++++++++++++++++++++
 common/unit_test/Test_Utils.cpp    |  9 +++++++++
 3 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/common/src/KokkosFFT_transpose.hpp b/common/src/KokkosFFT_transpose.hpp
index dd78e14..85b8d64 100644
--- 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
index 2747c36..93a2e59 100644
--- 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
index d104bd9..9df472f 100644
--- 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}));
+}