From e3454ae0f1d56e66f4575fd5ce863245fca6775e Mon Sep 17 00:00:00 2001
From: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com>
Date: Fri, 2 Aug 2024 10:47:16 +0200
Subject: [PATCH] Improve check functions (#125)

* Add a trait to extract the base value type from a container

* extend check functions to work on std::array or std::vector

* Avoid throw from a check function

* remove uunnecessary check from are_valid_axes and add overlaped axes tests

* rename: key to value for find function

* fix: remove a variable rank which is no longer used

---------

Co-authored-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>
---
 common/src/KokkosFFT_traits.hpp  |  25 ++++
 common/src/KokkosFFT_utils.hpp   |  72 +++++++++---
 common/unit_test/Test_Traits.cpp |  49 ++++++++
 common/unit_test/Test_Utils.cpp  | 190 +++++++++++++++++++++++++++----
 4 files changed, 300 insertions(+), 36 deletions(-)

diff --git a/common/src/KokkosFFT_traits.hpp b/common/src/KokkosFFT_traits.hpp
index 312797bd..284e65f8 100644
--- a/common/src/KokkosFFT_traits.hpp
+++ b/common/src/KokkosFFT_traits.hpp
@@ -215,6 +215,31 @@ inline constexpr bool are_operatable_views_v =
 
 // Other traits
 
+template <typename ContainerType>
+struct base_container_value;
+
+template <template <typename, typename...> class ContainerType,
+          typename ValueType, typename... Args>
+struct base_container_value<ContainerType<ValueType, Args...>> {
+  using value_type = ValueType;
+};
+
+// Specialization for std::array
+template <typename ValueType, std::size_t N>
+struct base_container_value<std::array<ValueType, N>> {
+  using value_type = ValueType;
+};
+
+// Specialization for Kokkos::Array
+template <typename ValueType, std::size_t N>
+struct base_container_value<Kokkos::Array<ValueType, N>> {
+  using value_type = ValueType;
+};
+
+/// \brief Helper to extract the base value type from a container
+template <typename T>
+using base_container_value_type = typename base_container_value<T>::value_type;
+
 /// \brief Helper to define a managable View type from the original view type
 template <typename T>
 struct managable_view_type {
diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp
index 4159991f..fd820f8a 100644
--- a/common/src/KokkosFFT_utils.hpp
+++ b/common/src/KokkosFFT_utils.hpp
@@ -11,6 +11,7 @@
 #include <algorithm>
 #include <numeric>
 #include "KokkosFFT_traits.hpp"
+#include "KokkosFFT_common_types.hpp"
 
 #if defined(KOKKOS_ENABLE_CXX17)
 #include <cstdlib>
@@ -85,21 +86,28 @@ auto convert_negative_shift(const ViewType& view, int _shift, int _axis) {
   return std::tuple<int, int, int>({shift0, shift1, shift2});
 }
 
-template <typename T>
-bool is_found(std::vector<T>& values, const T& key) {
-  return std::find(values.begin(), values.end(), key) != values.end();
+template <typename ContainerType, typename ValueType>
+bool is_found(ContainerType& values, const ValueType& value) {
+  using value_type = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  static_assert(std::is_same_v<value_type, ValueType>,
+                "Container value type must match ValueType");
+  return std::find(values.begin(), values.end(), value) != values.end();
 }
 
-template <typename T>
-bool has_duplicate_values(const std::vector<T>& values) {
-  std::set<T> set_values(values.begin(), values.end());
+template <typename ContainerType>
+bool has_duplicate_values(const ContainerType& values) {
+  using value_type = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  std::set<value_type> set_values(values.begin(), values.end());
   return set_values.size() < values.size();
 }
 
-template <typename IntType, std::enable_if_t<std::is_integral_v<IntType>,
-                                             std::nullptr_t> = nullptr>
-bool is_out_of_range_value_included(const std::vector<IntType>& values,
-                                    IntType max) {
+template <
+    typename ContainerType, typename IntType,
+    std::enable_if_t<std::is_integral_v<IntType>, std::nullptr_t> = nullptr>
+bool is_out_of_range_value_included(const ContainerType& values, IntType max) {
+  using value_type = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  static_assert(std::is_same_v<value_type, IntType>,
+                "Container value type must match IntType");
   bool is_included = false;
   for (auto value : values) {
     is_included = value >= max;
@@ -107,6 +115,39 @@ bool is_out_of_range_value_included(const std::vector<IntType>& values,
   return is_included;
 }
 
+template <
+    typename ViewType, template <typename, std::size_t> class ArrayType,
+    typename IntType, std::size_t DIM = 1,
+    std::enable_if_t<std::is_integral_v<IntType>, std::nullptr_t> = nullptr>
+bool are_valid_axes(const ViewType& view, const ArrayType<IntType, DIM>& axes) {
+  static_assert(
+      DIM >= 1 && DIM <= KokkosFFT::MAX_FFT_DIM,
+      "are_valid_axes: the Rank of FFT axes must be between 1 and MAX_FFT_DIM");
+  static_assert(ViewType::rank() >= DIM,
+                "are_valid_axes: View rank must be larger than or equal to the "
+                "Rank of FFT axes");
+
+  // Convert the input axes to be in the range of [0, rank-1]
+  // int type is choosen for consistency with the rest of the code
+  // the axes are defined with int type
+  std::array<int, DIM> non_negative_axes;
+
+  // In case axis is out of range, 'convert_negative_axis' will throw an
+  // runtime_error and we will return false. Without runtime_error, it is
+  // ensured that the 'non_negative_axes' are in the range of [0, rank-1]
+  try {
+    for (std::size_t i = 0; i < DIM; i++) {
+      int axis = KokkosFFT::Impl::convert_negative_axis(view, axes[i]);
+      non_negative_axes[i] = axis;
+    }
+  } catch (std::runtime_error& e) {
+    return false;
+  }
+
+  bool is_valid = !KokkosFFT::Impl::has_duplicate_values(non_negative_axes);
+  return is_valid;
+}
+
 template <std::size_t DIM = 1>
 bool is_transpose_needed(std::array<int, DIM> map) {
   std::array<int, DIM> contiguous_map;
@@ -114,14 +155,17 @@ bool is_transpose_needed(std::array<int, DIM> map) {
   return map != contiguous_map;
 }
 
-template <typename T>
-std::size_t get_index(std::vector<T>& values, const T& key) {
-  auto it           = find(values.begin(), values.end(), key);
+template <typename ContainerType, typename ValueType>
+std::size_t get_index(ContainerType& values, const ValueType& value) {
+  using value_type = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  static_assert(std::is_same_v<value_type, ValueType>,
+                "Container value type must match ValueType");
+  auto it           = std::find(values.begin(), values.end(), value);
   std::size_t index = 0;
   if (it != values.end()) {
     index = it - values.begin();
   } else {
-    throw std::runtime_error("key is not included in values");
+    throw std::runtime_error("value is not included in values");
   }
 
   return index;
diff --git a/common/unit_test/Test_Traits.cpp b/common/unit_test/Test_Traits.cpp
index 8d09753d..b17ac3a6 100644
--- a/common/unit_test/Test_Traits.cpp
+++ b/common/unit_test/Test_Traits.cpp
@@ -9,6 +9,9 @@
 // All the tests in this file are compile time tests, so we skip all the tests
 // by GTEST_SKIP(). gtest is used for type parameterization.
 
+// Int like types
+using base_int_types = ::testing::Types<int, std::size_t>;
+
 // Define the types to combine
 using base_real_types = std::tuple<float, double, long double>;
 
@@ -39,6 +42,19 @@ using paired_view_types =
     tuple_to_types_t<cartesian_product_t<base_real_types, base_layout_types,
                                          base_real_types, base_layout_types>>;
 
+template <typename T>
+struct ContainerTypes : public ::testing::Test {
+  static constexpr std::size_t rank = 3;
+  using value_type                  = T;
+  using vector_type                 = std::vector<T>;
+  using std_array_type              = std::array<T, rank>;
+  using Kokkos_array_type           = Kokkos::Array<T, rank>;
+
+  virtual void SetUp() {
+    GTEST_SKIP() << "Skipping all tests for this fixture";
+  }
+};
+
 template <typename T>
 struct RealAndComplexTypes : public ::testing::Test {
   using real_type    = T;
@@ -91,12 +107,45 @@ struct PairedViewTypes : public ::testing::Test {
   }
 };
 
+TYPED_TEST_SUITE(ContainerTypes, base_int_types);
 TYPED_TEST_SUITE(RealAndComplexTypes, real_types);
 TYPED_TEST_SUITE(RealAndComplexViewTypes, view_types);
 TYPED_TEST_SUITE(PairedValueTypes, paired_value_types);
 TYPED_TEST_SUITE(PairedLayoutTypes, paired_layout_types);
 TYPED_TEST_SUITE(PairedViewTypes, paired_view_types);
 
+// Tests for base value type deduction
+template <typename ValueType, typename ContainerType>
+void test_get_container_value_type() {
+  using value_type_ContainerType =
+      KokkosFFT::Impl::base_container_value_type<ContainerType>;
+
+  // base value type of ContainerType is ValueType
+  static_assert(std::is_same_v<value_type_ContainerType, ValueType>,
+                "Value type not deduced correctly from ContainerType");
+}
+
+TYPED_TEST(ContainerTypes, get_value_type_from_vector) {
+  using value_type     = typename TestFixture::value_type;
+  using container_type = typename TestFixture::vector_type;
+
+  test_get_container_value_type<value_type, container_type>();
+}
+
+TYPED_TEST(ContainerTypes, get_value_type_from_std_array) {
+  using value_type     = typename TestFixture::value_type;
+  using container_type = typename TestFixture::std_array_type;
+
+  test_get_container_value_type<value_type, container_type>();
+}
+
+TYPED_TEST(ContainerTypes, get_value_type_from_kokkos_array) {
+  using value_type     = typename TestFixture::value_type;
+  using container_type = typename TestFixture::Kokkos_array_type;
+
+  test_get_container_value_type<value_type, container_type>();
+}
+
 // Tests for real type deduction
 template <typename RealType, typename ComplexType>
 void test_get_real_type() {
diff --git a/common/unit_test/Test_Utils.cpp b/common/unit_test/Test_Utils.cpp
index b443a112..19698570 100644
--- a/common/unit_test/Test_Utils.cpp
+++ b/common/unit_test/Test_Utils.cpp
@@ -8,6 +8,9 @@
 
 using test_types = ::testing::Types<Kokkos::LayoutLeft, Kokkos::LayoutRight>;
 
+// Int like types
+using base_int_types = ::testing::Types<int, std::size_t>;
+
 // Basically the same fixtures, used for labeling tests
 template <typename T>
 struct ConvertNegativeAxis : public ::testing::Test {
@@ -19,8 +22,17 @@ struct ConvertNegativeShift : public ::testing::Test {
   using layout_type = T;
 };
 
+template <typename T>
+struct ContainerTypes : public ::testing::Test {
+  static constexpr std::size_t rank = 5;
+  using value_type                  = T;
+  using vector_type                 = std::vector<T>;
+  using array_type                  = std::array<T, rank>;
+};
+
 TYPED_TEST_SUITE(ConvertNegativeAxis, test_types);
 TYPED_TEST_SUITE(ConvertNegativeShift, test_types);
+TYPED_TEST_SUITE(ContainerTypes, base_int_types);
 
 // Tests for convert_negative_axes over ND views
 template <typename LayoutType>
@@ -267,25 +279,49 @@ TEST(IsTransposeNeeded, 1Dto3D) {
   EXPECT_TRUE(KokkosFFT::Impl::is_transpose_needed(map3D_210));
 }
 
-TEST(GetIndex, Vectors) {
-  std::vector<int> v = {0, 1, 4, 2, 3};
-
-  EXPECT_EQ(KokkosFFT::Impl::get_index(v, 0), 0);
-  EXPECT_EQ(KokkosFFT::Impl::get_index(v, 1), 1);
-  EXPECT_EQ(KokkosFFT::Impl::get_index(v, 2), 3);
-  EXPECT_EQ(KokkosFFT::Impl::get_index(v, 3), 4);
-  EXPECT_EQ(KokkosFFT::Impl::get_index(v, 4), 2);
-
-  EXPECT_THROW(KokkosFFT::Impl::get_index(v, -1), std::runtime_error);
+template <typename ContainerType>
+void test_is_found() {
+  using IntType   = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  ContainerType v = {0, 1, 4, 2, 3};
+
+  EXPECT_TRUE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(0)));
+  EXPECT_TRUE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(1)));
+  EXPECT_TRUE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(2)));
+  EXPECT_TRUE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(3)));
+  EXPECT_TRUE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(4)));
+
+  if constexpr (std::is_signed_v<IntType>) {
+    EXPECT_FALSE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(-1)));
+  }
+  EXPECT_FALSE(KokkosFFT::Impl::is_found(v, static_cast<IntType>(5)));
+}
 
-  EXPECT_THROW(KokkosFFT::Impl::get_index(v, 5), std::runtime_error);
+template <typename ContainerType>
+void test_get_index() {
+  using IntType   = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  ContainerType v = {0, 1, 4, 2, 3};
+
+  EXPECT_EQ(KokkosFFT::Impl::get_index(v, static_cast<IntType>(0)), 0);
+  EXPECT_EQ(KokkosFFT::Impl::get_index(v, static_cast<IntType>(1)), 1);
+  EXPECT_EQ(KokkosFFT::Impl::get_index(v, static_cast<IntType>(2)), 3);
+  EXPECT_EQ(KokkosFFT::Impl::get_index(v, static_cast<IntType>(3)), 4);
+  EXPECT_EQ(KokkosFFT::Impl::get_index(v, static_cast<IntType>(4)), 2);
+
+  if constexpr (std::is_signed_v<IntType>) {
+    EXPECT_THROW(KokkosFFT::Impl::get_index(v, static_cast<IntType>(-1)),
+                 std::runtime_error);
+  }
+  EXPECT_THROW(KokkosFFT::Impl::get_index(v, static_cast<IntType>(5)),
+               std::runtime_error);
 }
 
-TEST(HasDuplicateValues, Array) {
-  std::vector<int> v0 = {0, 1, 1};
-  std::vector<int> v1 = {0, 1, 1, 1};
-  std::vector<int> v2 = {0, 1, 2, 3};
-  std::vector<int> v3 = {0};
+template <typename ContainerType0, typename ContainerType1,
+          typename ContainerType2>
+void test_has_duplicate_values() {
+  ContainerType0 v0 = {0, 1, 1};
+  ContainerType1 v1 = {0, 1, 1, 1};
+  ContainerType1 v2 = {0, 1, 2, 3};
+  ContainerType2 v3 = {0};
 
   EXPECT_TRUE(KokkosFFT::Impl::has_duplicate_values(v0));
   EXPECT_TRUE(KokkosFFT::Impl::has_duplicate_values(v1));
@@ -293,13 +329,123 @@ TEST(HasDuplicateValues, Array) {
   EXPECT_FALSE(KokkosFFT::Impl::has_duplicate_values(v3));
 }
 
-TEST(IsOutOfRangeValueIncluded, Array) {
-  std::vector<int> v = {0, 1, 2, 3};
+template <typename ContainerType>
+void test_is_out_of_range_value_included() {
+  using IntType   = KokkosFFT::Impl::base_container_value_type<ContainerType>;
+  ContainerType v = {0, 1, 2, 3, 4};
+
+  EXPECT_TRUE(KokkosFFT::Impl::is_out_of_range_value_included(
+      v, static_cast<IntType>(2)));
+  EXPECT_TRUE(KokkosFFT::Impl::is_out_of_range_value_included(
+      v, static_cast<IntType>(3)));
+  EXPECT_FALSE(KokkosFFT::Impl::is_out_of_range_value_included(
+      v, static_cast<IntType>(5)));
+  EXPECT_FALSE(KokkosFFT::Impl::is_out_of_range_value_included(
+      v, static_cast<IntType>(6)));
+}
+
+template <typename IntType>
+void test_are_valid_axes() {
+  using real_type  = double;
+  using View1DType = Kokkos::View<real_type*>;
+  using View2DType = Kokkos::View<real_type**>;
+  using View3DType = Kokkos::View<real_type***>;
+  using View4DType = Kokkos::View<real_type****>;
+
+  std::array<IntType, 1> axes0 = {0};
+  std::array<IntType, 2> axes1 = {0, 1};
+  std::array<IntType, 3> axes2 = {0, 1, 2};
+  std::array<IntType, 3> axes3 = {0, 1, 1};
+  std::array<IntType, 3> axes4 = {0, 1, 3};
+
+  View1DType view1;
+  View2DType view2;
+  View3DType view3;
+  View4DType view4;
+
+  // 1D axes on 1D+ Views
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view1, axes0));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view2, axes0));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view3, axes0));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view4, axes0));
+
+  // 2D axes on 2D+ Views
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view2, axes1));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view3, axes1));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view4, axes1));
+
+  // 3D axes on 3D+ Views
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view3, axes2));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view4, axes2));
+  EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view4, axes4));
+
+  // 3D axes on 3D Views with out of range -> should fail
+  EXPECT_FALSE(KokkosFFT::Impl::are_valid_axes(view3, axes4));
+
+  // axes include overlap -> should fail
+  EXPECT_FALSE(KokkosFFT::Impl::are_valid_axes(view3, axes3));
+  EXPECT_FALSE(KokkosFFT::Impl::are_valid_axes(view4, axes3));
+
+  if constexpr (std::is_signed_v<IntType>) {
+    // {0, 1, -2} is converted to {0, 1, 1} for 3D View and {0, 1, 2}
+    // for 4D View. Invalid for 3D View but valid for 4D View
+    std::array<IntType, 3> axes5 = {0, 1, -2};
+
+    // axes include overlap -> should fail
+    EXPECT_FALSE(KokkosFFT::Impl::are_valid_axes(view3, axes5));
+
+    // axes do not include overlap -> OK
+    EXPECT_TRUE(KokkosFFT::Impl::are_valid_axes(view4, axes5));
+  }
+}
+
+TYPED_TEST(ContainerTypes, is_found_from_vector) {
+  using container_type = typename TestFixture::vector_type;
+  test_is_found<container_type>();
+}
+
+TYPED_TEST(ContainerTypes, is_found_from_array) {
+  using container_type = typename TestFixture::array_type;
+  test_is_found<container_type>();
+}
+
+TYPED_TEST(ContainerTypes, get_index_from_vector) {
+  using container_type = typename TestFixture::vector_type;
+  test_get_index<container_type>();
+}
+
+TYPED_TEST(ContainerTypes, get_index_from_array) {
+  using container_type = typename TestFixture::array_type;
+  test_get_index<container_type>();
+}
+
+TYPED_TEST(ContainerTypes, has_duplicate_values_in_vector) {
+  using container_type = typename TestFixture::vector_type;
+  test_has_duplicate_values<container_type, container_type, container_type>();
+}
+
+TYPED_TEST(ContainerTypes, has_duplicate_values_in_array) {
+  using value_type      = typename TestFixture::value_type;
+  using container_type0 = std::array<value_type, 3>;
+  using container_type1 = std::array<value_type, 4>;
+  using container_type2 = std::array<value_type, 1>;
+  test_has_duplicate_values<container_type0, container_type1,
+                            container_type2>();
+}
+
+TYPED_TEST(ContainerTypes, is_out_of_range_value_included_in_vector) {
+  using container_type = typename TestFixture::vector_type;
+  test_is_out_of_range_value_included<container_type>();
+}
+
+TYPED_TEST(ContainerTypes, is_out_of_range_value_included_in_array) {
+  using container_type = typename TestFixture::array_type;
+  test_is_out_of_range_value_included<container_type>();
+}
 
-  EXPECT_TRUE(KokkosFFT::Impl::is_out_of_range_value_included(v, 2));
-  EXPECT_TRUE(KokkosFFT::Impl::is_out_of_range_value_included(v, 3));
-  EXPECT_FALSE(KokkosFFT::Impl::is_out_of_range_value_included(v, 4));
-  EXPECT_FALSE(KokkosFFT::Impl::is_out_of_range_value_included(v, 5));
+TYPED_TEST(ContainerTypes, are_valid_axes) {
+  using value_type = typename TestFixture::value_type;
+  test_are_valid_axes<value_type>();
 }
 
 TEST(ExtractExtents, 1Dto8D) {