Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

More traits for unary and binary operations #119

Merged
merged 4 commits into from
Jul 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 113 additions & 5 deletions common/src/KokkosFFT_traits.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@

namespace KokkosFFT {
namespace Impl {

// Traits for unary operation

template <typename T>
struct base_floating_point {
using value_type = T;
Expand Down Expand Up @@ -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;
Expand All @@ -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 {};
Expand All @@ -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 {};
Expand All @@ -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 {
Expand Down
Loading
Loading