More traits for unary and binary operations (#119)

* More traits for unary and binary operations * Reduce combinations for same precision, layout and rank tests * Add mssing execution_space type check * fix: rename precision traits and improve comments --------- Co-authored-by: Yuuichi Asahi <[email protected]>
kokkos · Jul 23, 2024 · bab1620 · bab1620
1 parent aff3416
commit bab1620
Show file tree

Hide file tree

Showing 3 changed files with 694 additions and 22 deletions.
diff --git a/common/src/KokkosFFT_traits.hpp b/common/src/KokkosFFT_traits.hpp
@@ -9,6 +9,9 @@
 
 namespace KokkosFFT {
 namespace Impl {
+
+// Traits for unary operation
+
 template <typename T>
 struct base_floating_point {
   using value_type = T;
@@ -60,15 +63,15 @@ struct is_admissible_value_type<
 
 template <typename T>
 struct is_admissible_value_type<
-    T, std::enable_if_t<Kokkos::is_view<T>::value &&
+    T, std::enable_if_t<Kokkos::is_view_v<T> &&
                         (is_real_v<typename T::non_const_value_type> ||
                          is_complex_v<typename T::non_const_value_type>)>>
     : std::true_type {};
 
 /// \brief Helper to check if a type is an acceptable value type
 /// (float/double/Kokkos::complex<float>/Kokkos::complex<double>) for Kokkos-FFT
-///        When applied to Kokkos::View, then check if a value type is an
-///        acceptable real/complex type.
+/// When applied to Kokkos::View, then check if a value type is an acceptable
+/// real/complex type.
 template <typename T>
 inline constexpr bool is_admissible_value_type_v =
     is_admissible_value_type<T>::value;
@@ -80,7 +83,7 @@ template <typename ViewType>
 struct is_layout_left_or_right<
     ViewType,
     std::enable_if_t<
-        Kokkos::is_view<ViewType>::value &&
+        Kokkos::is_view_v<ViewType> &&
         (std::is_same_v<typename ViewType::array_layout, Kokkos::LayoutLeft> ||
          std::is_same_v<typename ViewType::array_layout, Kokkos::LayoutRight>)>>
     : std::true_type {};
@@ -96,7 +99,7 @@ struct is_admissible_view : std::false_type {};
 
 template <typename ViewType>
 struct is_admissible_view<
-    ViewType, std::enable_if_t<Kokkos::is_view<ViewType>::value &&
+    ViewType, std::enable_if_t<Kokkos::is_view_v<ViewType> &&
                                is_layout_left_or_right_v<ViewType> &&
                                is_admissible_value_type_v<ViewType>>>
     : std::true_type {};
@@ -107,6 +110,111 @@ template <typename ViewType>
 inline constexpr bool is_admissible_view_v =
     is_admissible_view<ViewType>::value;
 
+template <typename ExecutionSpace, typename ViewType, typename Enable = void>
+struct is_operatable_view : std::false_type {};
+
+template <typename ExecutionSpace, typename ViewType>
+struct is_operatable_view<
+    ExecutionSpace, ViewType,
+    std::enable_if_t<
+        Kokkos::is_execution_space_v<ExecutionSpace> &&
+        is_admissible_view_v<ViewType> &&
+        Kokkos::SpaceAccessibility<
+            ExecutionSpace, typename ViewType::memory_space>::accessible>>
+    : std::true_type {};
+
+/// \brief Helper to check if a View is an acceptable View for Kokkos-FFT and
+/// memory space is accessible from the ExecutionSpace
+template <typename ExecutionSpace, typename ViewType>
+inline constexpr bool is_operatable_view_v =
+    is_operatable_view<ExecutionSpace, ViewType>::value;
+
+// Traits for binary operations
+template <typename T1, typename T2, typename Enable = void>
+struct have_same_base_floating_point_type : std::false_type {};
+
+template <typename T1, typename T2>
+struct have_same_base_floating_point_type<
+    T1, T2,
+    std::enable_if_t<!Kokkos::is_view_v<T1> && !Kokkos::is_view_v<T2> &&
+                     std::is_same_v<base_floating_point_type<T1>,
+                                    base_floating_point_type<T2>>>>
+    : std::true_type {};
+
+template <typename InViewType, typename OutViewType>
+struct have_same_base_floating_point_type<
+    InViewType, OutViewType,
+    std::enable_if_t<
+        Kokkos::is_view_v<InViewType> && Kokkos::is_view_v<OutViewType> &&
+        std::is_same_v<
+            base_floating_point_type<typename InViewType::non_const_value_type>,
+            base_floating_point_type<
+                typename OutViewType::non_const_value_type>>>>
+    : std::true_type {};
+
+/// \brief Helper to check if two value have the same base floating point type.
+/// When applied to Kokkos::View, then check if values of views have the same
+/// base floating point type.
+template <typename T1, typename T2>
+inline constexpr bool have_same_base_floating_point_type_v =
+    have_same_base_floating_point_type<T1, T2>::value;
+
+template <typename InViewType, typename OutViewType, typename Enable = void>
+struct have_same_layout : std::false_type {};
+
+template <typename InViewType, typename OutViewType>
+struct have_same_layout<
+    InViewType, OutViewType,
+    std::enable_if_t<Kokkos::is_view_v<InViewType> &&
+                     Kokkos::is_view_v<OutViewType> &&
+                     std::is_same_v<typename InViewType::array_layout,
+                                    typename OutViewType::array_layout>>>
+    : std::true_type {};
+
+/// \brief Helper to check if two views have the same layout type.
+template <typename InViewType, typename OutViewType>
+inline constexpr bool have_same_layout_v =
+    have_same_layout<InViewType, OutViewType>::value;
+
+template <typename InViewType, typename OutViewType, typename Enable = void>
+struct have_same_rank : std::false_type {};
+
+template <typename InViewType, typename OutViewType>
+struct have_same_rank<
+    InViewType, OutViewType,
+    std::enable_if_t<Kokkos::is_view_v<InViewType> &&
+                     Kokkos::is_view_v<OutViewType> &&
+                     InViewType::rank() == OutViewType::rank()>>
+    : std::true_type {};
+
+/// \brief Helper to check if two views have the same rank.
+template <typename InViewType, typename OutViewType>
+inline constexpr bool have_same_rank_v =
+    have_same_rank<InViewType, OutViewType>::value;
+
+template <typename ExecutionSpace, typename InViewType, typename OutViewType,
+          typename Enable = void>
+struct are_operatable_views : std::false_type {};
+
+template <typename ExecutionSpace, typename InViewType, typename OutViewType>
+struct are_operatable_views<
+    ExecutionSpace, InViewType, OutViewType,
+    std::enable_if_t<
+        is_operatable_view_v<ExecutionSpace, InViewType> &&
+        is_operatable_view_v<ExecutionSpace, OutViewType> &&
+        have_same_base_floating_point_type_v<InViewType, OutViewType> &&
+        have_same_layout_v<InViewType, OutViewType> &&
+        have_same_rank_v<InViewType, OutViewType>>> : std::true_type {};
+
+/// \brief Helper to check if Views are acceptable View for Kokkos-FFT and
+/// memory space are accessible from the ExecutionSpace.
+/// In addition, precisions, layout and rank are checked to be identical.
+template <typename ExecutionSpace, typename InViewType, typename OutViewType>
+inline constexpr bool are_operatable_views_v =
+    are_operatable_views<ExecutionSpace, InViewType, OutViewType>::value;
+
+// Other traits
+
 /// \brief Helper to define a managable View type from the original view type
 template <typename T>
 struct managable_view_type {