Skip to content

Commit

Permalink
Some cleaning regarding the Kokkos layer (#70)
Browse files Browse the repository at this point in the history 
* Some cleaning regarding the Kokkos layer

* Update internal doc
  • Loading branch information
@tpadioleau
tpadioleau authored Jul 6, 2022
1 parent 7d5e765 commit fcac1f7
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 77 deletions.
4 changes: 2 additions & 2 deletions include/ddc/chunk.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
s.mapping(),
std::make_index_sequence<sizeof...(DDims)> {});
return Kokkos::View<
detail::mdspan_to_kokkos_element_type_t<ElementType, sizeof...(DDims)>,
detail::mdspan_to_kokkos_element_t<ElementType, sizeof...(DDims)>,
decltype(kokkos_layout),
typename Allocator::memory_space>(s.data(), kokkos_layout);
}
Expand All @@ -275,7 +275,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
s.mapping(),
std::make_index_sequence<sizeof...(DDims)> {});
return Kokkos::View<
detail::mdspan_to_kokkos_element_type_t<ElementType, sizeof...(DDims)>,
detail::mdspan_to_kokkos_element_t<ElementType, sizeof...(DDims)>,
decltype(kokkos_layout),
typename Allocator::memory_space>(s.data(), kokkos_layout);
}
Expand Down
6 changes: 3 additions & 3 deletions include/ddc/chunk_span.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ class ChunkSpan<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy, Memo
s.mapping(),
std::make_index_sequence<sizeof...(DDims)> {});
return Kokkos::View<
detail::mdspan_to_kokkos_element_type_t<ElementType, sizeof...(DDims)>,
detail::mdspan_to_kokkos_element_t<ElementType, sizeof...(DDims)>,
decltype(kokkos_layout),
MemorySpace>(s.data(), kokkos_layout);
}
Expand All @@ -325,9 +325,9 @@ template <
class... DDims,
class = std::enable_if_t<Kokkos::is_view<KokkosView>::value>>
ChunkSpan(KokkosView const& view, DiscreteDomain<DDims...> domain) -> ChunkSpan<
detail::kokkos_to_mdspan_element_type_t<typename KokkosView::data_type>,
detail::kokkos_to_mdspan_element_t<typename KokkosView::data_type>,
DiscreteDomain<DDims...>,
detail::mdspan_layout_t<typename KokkosView::array_layout>,
detail::kokkos_to_mdspan_layout_t<typename KokkosView::array_layout>,
typename KokkosView::memory_space>;

template <
Expand Down
119 changes: 63 additions & 56 deletions include/ddc/detail/kokkos.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,55 +8,101 @@

namespace detail {

template <class KokkosLP>
struct mdspan_layout;
template <class T>
struct type_holder
{
using type = T;
};

template <class KokkosLP>
using mdspan_layout_t = typename mdspan_layout<KokkosLP>::type;
struct kokkos_to_mdspan_layout
{
static_assert(
std::is_same_v<KokkosLP, KokkosLP>,
"Usage of non-specialized kokkos_to_mdspan_layout struct is not allowed");
};

template <>
struct mdspan_layout<Kokkos::LayoutLeft>
struct kokkos_to_mdspan_layout<Kokkos::LayoutLeft>
{
using type = std::experimental::layout_left;
};

template <>
struct mdspan_layout<Kokkos::LayoutRight>
struct kokkos_to_mdspan_layout<Kokkos::LayoutRight>
{
using type = std::experimental::layout_right;
};

template <>
struct mdspan_layout<Kokkos::LayoutStride>
struct kokkos_to_mdspan_layout<Kokkos::LayoutStride>
{
using type = std::experimental::layout_stride;
};

/// Alias template to transform a mdspan layout type to a Kokkos layout type
template <class KokkosLP>
using kokkos_to_mdspan_layout_t = typename kokkos_to_mdspan_layout<KokkosLP>::type;

template <class mdspanLP>
struct kokkos_layout;

template <class mdspanLP>
using kokkos_layout_t = typename kokkos_layout<mdspanLP>::type;
struct mdspan_to_kokkos_layout
{
static_assert(
std::is_same_v<mdspanLP, mdspanLP>,
"Usage of non-specialized mdspan_to_kokkos_layout struct is not allowed");
};

template <>
struct kokkos_layout<std::experimental::layout_left>
struct mdspan_to_kokkos_layout<std::experimental::layout_left>
{
using type = Kokkos::LayoutLeft;
};

template <>
struct kokkos_layout<std::experimental::layout_right>
struct mdspan_to_kokkos_layout<std::experimental::layout_right>
{
using type = Kokkos::LayoutRight;
};

template <>
struct kokkos_layout<std::experimental::layout_stride>
struct mdspan_to_kokkos_layout<std::experimental::layout_stride>
{
using type = Kokkos::LayoutStride;
};

/// Alias template to transform a Kokkos layout type to a mdspan layout type
template <class mdspanLP>
using mdspan_to_kokkos_layout_t = typename mdspan_to_kokkos_layout<mdspanLP>::type;

template <class ET, std::size_t N>
struct mdspan_to_kokkos_element
: std::conditional_t<
N == 0,
type_holder<ET>,
mdspan_to_kokkos_element<std::add_pointer_t<ET>, N - 1>>
{
};

/// Alias template to transform a mdspan element type to a Kokkos element type
/// Only dynamic dimensions is supported for now i.e. `double[4]*` is not yet covered.
template <class ET, std::size_t N>
using mdspan_to_kokkos_element_t = typename mdspan_to_kokkos_element<ET, N>::type;

template <class ET>
struct kokkos_to_mdspan_element
: std::conditional_t<
std::is_pointer_v<std::decay_t<ET>>,
kokkos_to_mdspan_element<std::remove_pointer_t<std::decay_t<ET>>>,
type_holder<ET>>
{
};

/// Alias template to transform a Kokkos element type to a mdspan element type
/// Only dynamic dimensions is supported for now i.e. `double[4]*` is not yet covered.
template <class ET>
using kokkos_to_mdspan_element_t = typename kokkos_to_mdspan_element<ET>::type;


template <std::size_t... Is>
Kokkos::LayoutStride make_layout_stride(
Expand All @@ -67,13 +113,13 @@ Kokkos::LayoutStride make_layout_stride(
}

template <class EP, class MP, std::size_t... Is>
kokkos_layout_t<typename MP::layout_type> build_kokkos_layout(
mdspan_to_kokkos_layout_t<typename MP::layout_type> build_kokkos_layout(
EP const& ep,
MP const& mapping,
std::index_sequence<Is...>)
{
DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
using kokkos_layout_type = kokkos_layout_t<typename MP::layout_type>;
using kokkos_layout_type = mdspan_to_kokkos_layout_t<typename MP::layout_type>;
if constexpr (std::is_same_v<kokkos_layout_type, Kokkos::LayoutStride>) {
std::array<std::size_t, sizeof...(Is) * 2> storage;
std::experimental::mdspan<
Expand All @@ -91,53 +137,14 @@ kokkos_layout_t<typename MP::layout_type> build_kokkos_layout(
DDC_IF_NVCC_THEN_POP
}

/// Recursively add a pointer
template <class ET, std::size_t N>
struct mdspan_to_kokkos_element_type : mdspan_to_kokkos_element_type<std::add_pointer_t<ET>, N - 1>
{
};

template <class ET>
struct mdspan_to_kokkos_element_type<ET, 0>
{
using type = ET;
};

template <class ET, std::size_t N>
using mdspan_to_kokkos_element_type_t = typename mdspan_to_kokkos_element_type<ET, N>::type;

template <class T, std::size_t N>
struct final_type
{
using type = T;
static constexpr std::size_t rank = N;
};

/// Recursively remove a pointer
template <class ET, std::size_t N>
struct kokkos_to_mdspan_element_type
: std::conditional_t<
std::is_pointer_v<std::decay_t<ET>>,
kokkos_to_mdspan_element_type<std::remove_pointer_t<std::decay_t<ET>>, N + 1>,
final_type<ET, N>>
{
};

template <class ET>
using kokkos_to_mdspan_element_type_t = typename kokkos_to_mdspan_element_type<ET, 0>::type;

template <class ET>
constexpr inline std::size_t kokkos_to_mdspan_element_type_rank
= kokkos_to_mdspan_element_type<ET, 0>::rank;

template <class DataType, class... Properties, std::size_t... Is>
auto build_mdspan(Kokkos::View<DataType, Properties...> const view, std::index_sequence<Is...>)
{
DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
using element_type = kokkos_to_mdspan_element_type_t<DataType>;
using element_type = kokkos_to_mdspan_element_t<DataType>;
using extents_type = std::experimental::dextents<Kokkos::View<DataType, Properties...>::rank>;
using layout_type
= mdspan_layout_t<typename Kokkos::View<DataType, Properties...>::array_layout>;
using layout_type = kokkos_to_mdspan_layout_t<
typename Kokkos::View<DataType, Properties...>::array_layout>;
using mapping_type = typename layout_type::template mapping<extents_type>;
extents_type exts(view.extent(Is)...);
if constexpr (std::is_same_v<layout_type, std::experimental::layout_stride>) {
Expand Down
33 changes: 17 additions & 16 deletions include/ddc/transform_reduce.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,71 +15,72 @@
namespace detail {

template <class Reducer>
struct KokkosReducer;

template <class Reducer>
using KokkosReducer_t = typename KokkosReducer<Reducer>::type;
struct ddc_to_kokkos_reducer;

template <class T>
struct KokkosReducer<reducer::sum<T>>
struct ddc_to_kokkos_reducer<reducer::sum<T>>
{
using type = Kokkos::Sum<T>;
};

template <class T>
struct KokkosReducer<reducer::prod<T>>
struct ddc_to_kokkos_reducer<reducer::prod<T>>
{
using type = Kokkos::Prod<T>;
};

template <>
struct KokkosReducer<reducer::land>
struct ddc_to_kokkos_reducer<reducer::land>
{
using type = Kokkos::LAnd<void>;
};

template <>
struct KokkosReducer<reducer::lor>
struct ddc_to_kokkos_reducer<reducer::lor>
{
using type = Kokkos::LOr<void>;
};

template <class T>
struct KokkosReducer<reducer::band<T>>
struct ddc_to_kokkos_reducer<reducer::band<T>>
{
using type = Kokkos::BAnd<T>;
};

template <class T>
struct KokkosReducer<reducer::bor<T>>
struct ddc_to_kokkos_reducer<reducer::bor<T>>
{
using type = Kokkos::BOr<T>;
};

template <class T>
struct KokkosReducer<reducer::bxor<T>>
struct ddc_to_kokkos_reducer<reducer::bxor<T>>
{
static_assert(std::is_same_v<T, T>, "This reducer is not yet implemented");
};

template <class T>
struct KokkosReducer<reducer::min<T>>
struct ddc_to_kokkos_reducer<reducer::min<T>>
{
using type = Kokkos::Min<T>;
};

template <class T>
struct KokkosReducer<reducer::max<T>>
struct ddc_to_kokkos_reducer<reducer::max<T>>
{
using type = Kokkos::Max<T>;
};

template <class T>
struct KokkosReducer<reducer::minmax<T>>
struct ddc_to_kokkos_reducer<reducer::minmax<T>>
{
using type = Kokkos::MinMax<T>;
};

/// Alias template to transform a DDC reducer type to a Kokkos reducer type
template <class Reducer>
using ddc_to_kokkos_reducer_t = typename ddc_to_kokkos_reducer<Reducer>::type;

/** A serial reduction over a nD domain
* @param[in] domain the range over which to apply the algorithm
* @param[in] neutral the neutral element of the reduction operation
Expand Down Expand Up @@ -165,7 +166,7 @@ inline T transform_reduce_kokkos(
BinaryReductionOp,
UnaryTransformOp,
DDim0>(reduce, transform),
KokkosReducer_t<BinaryReductionOp>(result));
ddc_to_kokkos_reducer_t<BinaryReductionOp>(result));
return result;
}

Expand Down Expand Up @@ -208,7 +209,7 @@ inline T transform_reduce_kokkos(
DDim0,
DDim1,
DDims...>(reduce, transform),
KokkosReducer_t<BinaryReductionOp>(result));
ddc_to_kokkos_reducer_t<BinaryReductionOp>(result));
return result;
}

Expand Down

0 comments on commit fcac1f7

Please sign in to comment.