Skip to content

Commit

Permalink
kokkos#6805: make mdspan conversion compatible with device code
Browse files Browse the repository at this point in the history
  • Loading branch information
@nmm0
nmm0 committed May 15, 2024
1 parent be51f9d commit 2792eeb
Show file tree
Hide file tree
Showing 2 changed files with 140 additions and 94 deletions.
7 changes: 4 additions & 3 deletions core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,12 @@ KOKKOS_INLINE_FUNCTION auto mapping_from_view_mapping(const VM &view_mapping) {
// std::span is not available in C++17 (our current requirements),
// so we need to use the std::array constructor for layout mappings.
// When C++20 is available, we can use std::span here instead
std::array<std::size_t, VM::Rank> strides;
view_mapping.stride_fill(strides.data());
std::size_t strides[VM::Rank];
view_mapping.stride_fill(&strides[0]);
if constexpr (std::is_same_v<typename mapping_type::layout_type,
Kokkos::layout_stride>) {
return mapping_type(extents_from_view_mapping<extents_type>(view_mapping),
return mapping_type(Kokkos::mdspan_non_standard,
extents_from_view_mapping<extents_type>(view_mapping),
strides);
} else if constexpr (VM::Rank > 1 &&
std::is_same_v<typename mapping_type::layout_type,
Expand Down
227 changes: 136 additions & 91 deletions core/unit_test/TestMDSpanConversion.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,11 @@ struct TestViewMDSpanConversion {
static void test_conversion_from_mdspan(
Kokkos::View<DataType, RefViewProps...> ref,
const MDSpanLayoutMapping &mapping) {
using unmanaged_view_view_type =
using unmanaged_view_type =
Kokkos::View<DataType, KokkosLayout, ExecutionSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using natural_mdspan_type = typename Kokkos::Impl::MDSpanViewTraits<
typename unmanaged_view_view_type::traits>::mdspan_type;
typename unmanaged_view_type::traits>::mdspan_type;
using mapping_type = MDSpanLayoutMapping;
using mdspan_layout_type = typename MDSpanLayoutMapping::layout_type;
using extents_type = typename mapping_type::extents_type;
Expand All @@ -67,14 +67,15 @@ struct TestViewMDSpanConversion {

static_assert(std::is_constructible_v<natural_mdspan_type, mdspan_type>);
static_assert(std::is_convertible_v<mdspan_type, natural_mdspan_type> ==
std::is_convertible_v<mdspan_type, unmanaged_view_view_type>);

std::is_convertible_v<mdspan_type, unmanaged_view_type>);
// Manually create an mdspan from ref so we have a valid pointer to play
// with
const auto &exts = mapping.extents();
auto mds = mdspan_type{ref.data(), mapping};

auto test_view = unmanaged_view_view_type(mds);
auto test_view = unmanaged_view_type(mds);

#if 0

ASSERT_EQ(test_view.data(), ref.data());
ASSERT_EQ(test_view.data(), mds.data_handle());
Expand All @@ -83,6 +84,7 @@ struct TestViewMDSpanConversion {
ASSERT_EQ(test_view.extent(r), ref.extent(r));
ASSERT_EQ(test_view.extent(r), exts.extent(r));
}
#endif
}

template <class MDSpanLayoutMapping, class ViewType>
Expand Down Expand Up @@ -178,7 +180,6 @@ struct TestViewMDSpanConversion {
!std::is_convertible_v<Kokkos::View<value_type **, Kokkos::LayoutRight>,
natural_mdspan_type_for_view<Kokkos::View<
value_type **, Kokkos::LayoutLeft>>>);

// nvcc doesn't do CTAD properly here, making this way more verbose..
// LayoutLeft
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Expand All @@ -187,6 +188,7 @@ struct TestViewMDSpanConversion {
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::dextents<std::size_t, 1>(7)});

test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type[7], Kokkos::LayoutLeft, ExecutionSpace>("ref"),
typename layout_left_padded<sizeof(
Expand Down Expand Up @@ -279,53 +281,78 @@ struct TestViewMDSpanConversion {
Kokkos::extents<std::size_t, 7, 3>()});

// LayoutStride
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::dextents<std::size_t, 1>{7},
std::array<std::size_t, 1>{2}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7>>{
{}, std::array<std::size_t, 1>{2}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::dextents<std::size_t, 1>{7},
std::array<std::size_t, 1>{2}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::extents<std::size_t, 7>(), std::array<std::size_t, 1>{2}});

test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3),
std::array<std::size_t, 2>{2, 4}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7][3], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>(),
std::array<std::size_t, 2>{2, 4}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7][3], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3),
std::array<std::size_t, 2>{2, 4}});
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>(),
std::array<std::size_t, 2>{2, 4}});
{
const size_t strides[] = {2};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::mdspan_non_standard, Kokkos::dextents<std::size_t, 1>{7},
strides});
}
{
const size_t strides[] = {2};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::mdspan_non_standard, {}, strides});
}
{
const size_t strides[] = {2};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::mdspan_non_standard, Kokkos::dextents<std::size_t, 1>{7},
strides});
}
{
const size_t strides[] = {2};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::mdspan_non_standard, Kokkos::extents<std::size_t, 7>(),
strides});
}

{
const size_t strides[] = {2, 4};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::mdspan_non_standard,
Kokkos::dextents<std::size_t, 2>(7, 3), strides});
}
{
const size_t strides[] = {2, 4};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7][3], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::mdspan_non_standard, Kokkos::extents<std::size_t, 7, 3>(),
strides});
}
{
const size_t strides[] = {2, 4};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type[7][3], Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::mdspan_non_standard,
Kokkos::dextents<std::size_t, 2>(7, 3), strides});
}
{
const size_t strides[] = {2, 4};
test_conversion_from_mdspan<Kokkos::LayoutStride>(
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"ref", Kokkos::LayoutStride{7, 2, 3, 4}),
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::mdspan_non_standard, Kokkos::extents<std::size_t, 7, 3>(),
strides});
}

// Conversion to mdspan
test_conversion_to_mdspan(
Expand All @@ -349,29 +376,35 @@ struct TestViewMDSpanConversion {
Kokkos::View<value_type **, Kokkos::LayoutRight, ExecutionSpace>("v", 4,
7));

test_conversion_to_mdspan(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
{}, std::array<std::size_t, 1>{5}),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}));
test_conversion_to_mdspan(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
{}, std::array<std::size_t, 2>{5, 9}),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}));
{
const size_t strides[] = {5};
test_conversion_to_mdspan(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}));
}
{
const size_t strides[] = {5, 9};
test_conversion_to_mdspan(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}));
}

// Aligned types (for padded layouts)
test_conversion_to_mdspan(
layout_left_padded<Kokkos::dynamic_extent>::mapping<
Kokkos::extents<std::size_t, 31, 7>>({}, 32),
Kokkos::extents<std::size_t, 127, 7>>({}, 128),
Kokkos::View<value_type **, Kokkos::LayoutLeft, ExecutionSpace>(
Kokkos::view_alloc("v", Kokkos::AllowPadding), 31, 7));
Kokkos::view_alloc("v", Kokkos::AllowPadding), 127, 7));

test_conversion_to_mdspan(
layout_right_padded<Kokkos::dynamic_extent>::mapping<
Kokkos::extents<std::size_t, 7, 31>>({}, 32),
Kokkos::extents<std::size_t, 7, 127>>({}, 128),
Kokkos::View<value_type **, Kokkos::LayoutRight, ExecutionSpace>(
Kokkos::view_alloc("v", Kokkos::AllowPadding), 7, 31));
Kokkos::view_alloc("v", Kokkos::AllowPadding), 7, 127));

// Conversion with standard default_accessor

Expand Down Expand Up @@ -399,18 +432,24 @@ struct TestViewMDSpanConversion {
7),
Kokkos::default_accessor<value_type>{});

test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
{}, std::array<std::size_t, 1>{5}),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}),
Kokkos::default_accessor<value_type>{});
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
{}, std::array<std::size_t, 2>{5, 9}),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}),
Kokkos::default_accessor<value_type>{});
{
const size_t strides[] = {5};
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}),
Kokkos::default_accessor<value_type>{});
}
{
const size_t strides[] = {5, 9};
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}),
Kokkos::default_accessor<value_type>{});
}

// Conversion with a test accessor

Expand Down Expand Up @@ -438,18 +477,24 @@ struct TestViewMDSpanConversion {
7),
TestAccessor{});

test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
{}, std::array<std::size_t, 1>{5}),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}),
TestAccessor{});
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
{}, std::array<std::size_t, 2>{5, 9}),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}),
TestAccessor{});
{
const size_t strides[] = {5};
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type *, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5}),
TestAccessor{});
}
{
const size_t strides[] = {5, 9};
test_conversion_to_mdspan_with_accessor(
Kokkos::layout_stride::mapping<Kokkos::extents<std::size_t, 4, 7>>(
Kokkos::mdspan_non_standard, {}, strides),
Kokkos::View<value_type **, Kokkos::LayoutStride, ExecutionSpace>(
"v", Kokkos::LayoutStride{4, 5, 7, 9}),
TestAccessor{});
}
}
};

Expand Down

0 comments on commit 2792eeb

Please sign in to comment.