More traits for unary and binary operations

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;
@@ -107,6 +110,113 @@ 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<
+        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_precision : std::false_type {};
+
+template <typename T1, typename T2>
+struct have_same_precision<
+    T1, T2,
+    std::enable_if_t<!Kokkos::is_view<T1>::value &&
+                     !Kokkos::is_view<T2>::value &&
+                     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_precision<
+    InViewType, OutViewType,
+    std::enable_if_t<
+        Kokkos::is_view<InViewType>::value &&
+        Kokkos::is_view<OutViewType>::value &&
+        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 precision type.
+///      When applied to Kokkos::View, then check if values of views have the
+///      same base precision type.
+template <typename T1, typename T2>
+inline constexpr bool have_same_precision_v =
+    have_same_precision<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<InViewType>::value &&
+                     Kokkos::is_view<OutViewType>::value &&
+                     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<InViewType>::value &&
+                     Kokkos::is_view<OutViewType>::value &&
+                     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_precision_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 {