Skip to content

Commit bab1620

Browse files
yasahi-hpcYuuichi Asahi
yasahi-hpc
and
Yuuichi Asahi
authored
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]>
1 parent aff3416 commit bab1620

File tree

3 files changed

+694
-22
lines changed

3 files changed

+694
-22
lines changed

‎common/src/KokkosFFT_traits.hpp

+113-5
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99

1010
namespace KokkosFFT {
1111
namespace Impl {
12+
13+
// Traits for unary operation
14+
1215
template <typename T>
1316
struct base_floating_point {
1417
using value_type = T;
@@ -60,15 +63,15 @@ struct is_admissible_value_type<
6063

6164
template <typename T>
6265
struct is_admissible_value_type<
63-
T, std::enable_if_t<Kokkos::is_view<T>::value &&
66+
T, std::enable_if_t<Kokkos::is_view_v<T> &&
6467
(is_real_v<typename T::non_const_value_type> ||
6568
is_complex_v<typename T::non_const_value_type>)>>
6669
: std::true_type {};
6770

6871
/// \brief Helper to check if a type is an acceptable value type
6972
/// (float/double/Kokkos::complex<float>/Kokkos::complex<double>) for Kokkos-FFT
70-
/// When applied to Kokkos::View, then check if a value type is an
71-
/// acceptable real/complex type.
73+
/// When applied to Kokkos::View, then check if a value type is an acceptable
74+
/// real/complex type.
7275
template <typename T>
7376
inline constexpr bool is_admissible_value_type_v =
7477
is_admissible_value_type<T>::value;
@@ -80,7 +83,7 @@ template <typename ViewType>
8083
struct is_layout_left_or_right<
8184
ViewType,
8285
std::enable_if_t<
83-
Kokkos::is_view<ViewType>::value &&
86+
Kokkos::is_view_v<ViewType> &&
8487
(std::is_same_v<typename ViewType::array_layout, Kokkos::LayoutLeft> ||
8588
std::is_same_v<typename ViewType::array_layout, Kokkos::LayoutRight>)>>
8689
: std::true_type {};
@@ -96,7 +99,7 @@ struct is_admissible_view : std::false_type {};
9699

97100
template <typename ViewType>
98101
struct is_admissible_view<
99-
ViewType, std::enable_if_t<Kokkos::is_view<ViewType>::value &&
102+
ViewType, std::enable_if_t<Kokkos::is_view_v<ViewType> &&
100103
is_layout_left_or_right_v<ViewType> &&
101104
is_admissible_value_type_v<ViewType>>>
102105
: std::true_type {};
@@ -107,6 +110,111 @@ template <typename ViewType>
107110
inline constexpr bool is_admissible_view_v =
108111
is_admissible_view<ViewType>::value;
109112

113+
template <typename ExecutionSpace, typename ViewType, typename Enable = void>
114+
struct is_operatable_view : std::false_type {};
115+
116+
template <typename ExecutionSpace, typename ViewType>
117+
struct is_operatable_view<
118+
ExecutionSpace, ViewType,
119+
std::enable_if_t<
120+
Kokkos::is_execution_space_v<ExecutionSpace> &&
121+
is_admissible_view_v<ViewType> &&
122+
Kokkos::SpaceAccessibility<
123+
ExecutionSpace, typename ViewType::memory_space>::accessible>>
124+
: std::true_type {};
125+
126+
/// \brief Helper to check if a View is an acceptable View for Kokkos-FFT and
127+
/// memory space is accessible from the ExecutionSpace
128+
template <typename ExecutionSpace, typename ViewType>
129+
inline constexpr bool is_operatable_view_v =
130+
is_operatable_view<ExecutionSpace, ViewType>::value;
131+
132+
// Traits for binary operations
133+
template <typename T1, typename T2, typename Enable = void>
134+
struct have_same_base_floating_point_type : std::false_type {};
135+
136+
template <typename T1, typename T2>
137+
struct have_same_base_floating_point_type<
138+
T1, T2,
139+
std::enable_if_t<!Kokkos::is_view_v<T1> && !Kokkos::is_view_v<T2> &&
140+
std::is_same_v<base_floating_point_type<T1>,
141+
base_floating_point_type<T2>>>>
142+
: std::true_type {};
143+
144+
template <typename InViewType, typename OutViewType>
145+
struct have_same_base_floating_point_type<
146+
InViewType, OutViewType,
147+
std::enable_if_t<
148+
Kokkos::is_view_v<InViewType> && Kokkos::is_view_v<OutViewType> &&
149+
std::is_same_v<
150+
base_floating_point_type<typename InViewType::non_const_value_type>,
151+
base_floating_point_type<
152+
typename OutViewType::non_const_value_type>>>>
153+
: std::true_type {};
154+
155+
/// \brief Helper to check if two value have the same base floating point type.
156+
/// When applied to Kokkos::View, then check if values of views have the same
157+
/// base floating point type.
158+
template <typename T1, typename T2>
159+
inline constexpr bool have_same_base_floating_point_type_v =
160+
have_same_base_floating_point_type<T1, T2>::value;
161+
162+
template <typename InViewType, typename OutViewType, typename Enable = void>
163+
struct have_same_layout : std::false_type {};
164+
165+
template <typename InViewType, typename OutViewType>
166+
struct have_same_layout<
167+
InViewType, OutViewType,
168+
std::enable_if_t<Kokkos::is_view_v<InViewType> &&
169+
Kokkos::is_view_v<OutViewType> &&
170+
std::is_same_v<typename InViewType::array_layout,
171+
typename OutViewType::array_layout>>>
172+
: std::true_type {};
173+
174+
/// \brief Helper to check if two views have the same layout type.
175+
template <typename InViewType, typename OutViewType>
176+
inline constexpr bool have_same_layout_v =
177+
have_same_layout<InViewType, OutViewType>::value;
178+
179+
template <typename InViewType, typename OutViewType, typename Enable = void>
180+
struct have_same_rank : std::false_type {};
181+
182+
template <typename InViewType, typename OutViewType>
183+
struct have_same_rank<
184+
InViewType, OutViewType,
185+
std::enable_if_t<Kokkos::is_view_v<InViewType> &&
186+
Kokkos::is_view_v<OutViewType> &&
187+
InViewType::rank() == OutViewType::rank()>>
188+
: std::true_type {};
189+
190+
/// \brief Helper to check if two views have the same rank.
191+
template <typename InViewType, typename OutViewType>
192+
inline constexpr bool have_same_rank_v =
193+
have_same_rank<InViewType, OutViewType>::value;
194+
195+
template <typename ExecutionSpace, typename InViewType, typename OutViewType,
196+
typename Enable = void>
197+
struct are_operatable_views : std::false_type {};
198+
199+
template <typename ExecutionSpace, typename InViewType, typename OutViewType>
200+
struct are_operatable_views<
201+
ExecutionSpace, InViewType, OutViewType,
202+
std::enable_if_t<
203+
is_operatable_view_v<ExecutionSpace, InViewType> &&
204+
is_operatable_view_v<ExecutionSpace, OutViewType> &&
205+
have_same_base_floating_point_type_v<InViewType, OutViewType> &&
206+
have_same_layout_v<InViewType, OutViewType> &&
207+
have_same_rank_v<InViewType, OutViewType>>> : std::true_type {};
208+
209+
/// \brief Helper to check if Views are acceptable View for Kokkos-FFT and
210+
/// memory space are accessible from the ExecutionSpace.
211+
/// In addition, precisions, layout and rank are checked to be identical.
212+
template <typename ExecutionSpace, typename InViewType, typename OutViewType>
213+
inline constexpr bool are_operatable_views_v =
214+
are_operatable_views<ExecutionSpace, InViewType, OutViewType>::value;
215+
216+
// Other traits
217+
110218
/// \brief Helper to define a managable View type from the original view type
111219
template <typename T>
112220
struct managable_view_type {

0 commit comments

Comments
 (0)