Skip to content

Commit

Permalink
kokkos#6805: don't allow arbitrary padding values
Browse files Browse the repository at this point in the history
  • Loading branch information
@nmm0
nmm0 committed May 15, 2024
1 parent 75bb319 commit ff819c8
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 40 deletions.
23 changes: 23 additions & 0 deletions core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,29 @@ KOKKOS_INLINE_FUNCTION auto array_layout_from_mapping(
rank > 7 ? mapping.stride(7) : 0,
};
} else {
// FIXME: Kokkos Layouts don't store stride (it's in the mapping)
// We could conceivably fix this by adding an extra ViewCtorProp for
// an abritrary padding. For now we will check for this.
if constexpr (rank > 1 &&
(std::is_same_v<
typename mapping_type::layout_type,
Experimental::layout_left_padded<dynamic_extent>> ||
std::is_same_v<
typename mapping_type::layout_type,
Experimental::layout_right_padded<dynamic_extent>>)) {
constexpr size_t strided_index =
std::is_same_v<typename mapping_type::layout_type,
Experimental::layout_left_padded<dynamic_extent>>
? 1
: rank - 2;
constexpr size_t extent_index =
std::is_same_v<typename mapping_type::layout_type,
Experimental::layout_left_padded<dynamic_extent>>
? 0
: rank - 1;
KOKKOS_ASSERT(mapping.stride(strided_index) == ext.extent(extent_index));
}

return ArrayLayout{rank > 0 ? ext.extent(0) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,
rank > 1 ? ext.extent(1) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,
rank > 2 ? ext.extent(2) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,
Expand Down
64 changes: 24 additions & 40 deletions core/unit_test/TestMDSpanConversion.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,100 +182,84 @@ struct TestViewMDSpanConversion {
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type *, Kokkos::LayoutLeft, ExecutionSpace>("ref",
7),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 1>>{
typename layout_left_padded<7>::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(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7>>{
typename layout_left_padded<7>::template mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::extents<std::size_t, 7>()});
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type[7], Kokkos::LayoutLeft, ExecutionSpace>("ref"),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 1>>{
typename layout_left_padded<7>::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 *, Kokkos::LayoutLeft, ExecutionSpace>("ref",
7),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7>>{
typename layout_left_padded<7>::template mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::extents<std::size_t, 7>()});

test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type **, Kokkos::LayoutLeft, ExecutionSpace>("ref",
7, 3),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 2>>{
typename layout_left_padded<7>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3)});
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type[7][3], Kokkos::LayoutLeft, ExecutionSpace>(
"ref"),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
typename layout_left_padded<7>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>()});
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type[7][3], Kokkos::LayoutLeft, ExecutionSpace>(
"ref"),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 2>>{
typename layout_left_padded<7>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3)});
test_conversion_from_mdspan<Kokkos::LayoutLeft>(
Kokkos::View<value_type **, Kokkos::LayoutLeft, ExecutionSpace>("ref",
7, 3),
typename layout_left_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
typename layout_left_padded<7>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>()});

// LayoutRight
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type *, Kokkos::LayoutRight, ExecutionSpace>("ref",
7),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 1>>{
typename layout_right_padded<7>::template mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::dextents<std::size_t, 1>(7)});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type[7], Kokkos::LayoutRight, ExecutionSpace>("ref"),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7>>{
typename layout_right_padded<7>::template mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::extents<std::size_t, 7>()});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type[7], Kokkos::LayoutRight, ExecutionSpace>("ref"),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 1>>{
typename layout_right_padded<7>::template mapping<Kokkos::dextents<std::size_t, 1>>{
Kokkos::dextents<std::size_t, 1>(7)});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type *, Kokkos::LayoutRight, ExecutionSpace>("ref",
7),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7>>{
typename layout_right_padded<7>::template mapping<Kokkos::extents<std::size_t, 7>>{
Kokkos::extents<std::size_t, 7>()});

test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type **, Kokkos::LayoutRight, ExecutionSpace>("ref",
7, 3),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3)});
3, 7),
typename layout_right_padded<7>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(3, 7)});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type[7][3], Kokkos::LayoutRight, ExecutionSpace>(
Kokkos::View<value_type[3][7], Kokkos::LayoutRight, ExecutionSpace>(
"ref"),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>()});
typename layout_right_padded<7>::template mapping<Kokkos::extents<std::size_t, 3, 7>>{
Kokkos::extents<std::size_t, 3, 7>()});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type[7][3], Kokkos::LayoutRight, ExecutionSpace>(
Kokkos::View<value_type[3][7], Kokkos::LayoutRight, ExecutionSpace>(
"ref"),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(7, 3)});
typename layout_right_padded<7>::template mapping<Kokkos::dextents<std::size_t, 2>>{
Kokkos::dextents<std::size_t, 2>(3, 7)});
test_conversion_from_mdspan<Kokkos::LayoutRight>(
Kokkos::View<value_type **, Kokkos::LayoutRight, ExecutionSpace>("ref",
7, 3),
typename layout_right_padded<sizeof(
value_type)>::template mapping<Kokkos::extents<std::size_t, 7, 3>>{
Kokkos::extents<std::size_t, 7, 3>()});
3, 7),
typename layout_right_padded<7>::template mapping<Kokkos::extents<std::size_t, 3, 7>>{
Kokkos::extents<std::size_t, 3, 7>()});

// LayoutStride
{
Expand Down

0 comments on commit ff819c8

Please sign in to comment.