Skip to content

Commit

Permalink
Check matching static extents in View constructor (kokkos#5190)
Browse files Browse the repository at this point in the history 
* Also check static extents

* Use separate bools for the two checks and improve error messages

* Improve error messages

* Refactor test

* Add another test

* Extend existing tests

* Use make_index_sequence

* Remove LIVE and DIE

* Try removing guards for OpenMPTarget

* Use different style in view_construction_with_wrong_static_extents

* Remove template template parameter

* Also check unmanaged views

* Fix DynRankView

* Add FIXME

* Pass label by const char *

* fix terminating define

* Deduce rank and dynamic_rank from View

* Only call runtime_check_rank when KOKKOS_ENABLE_DEBUG_BOUNDS_CHECK is defined

* Guard with if constexpr

---------

Co-authored-by: Francesco Rizzi <fnrizzi@sandia.gov>
  • Loading branch information
@masterleinad @fnrizzi
masterleinad and fnrizzi committed Jan 9, 2024
1 parent 27286c3 commit 0e4a158
Show file tree
Hide file tree
Showing 5 changed files with 265 additions and 383 deletions.
13 changes: 11 additions & 2 deletions containers/src/Kokkos_DynRankView.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1653,8 +1653,17 @@ KOKKOS_FUNCTION auto as_view_of_rank_n(
Kokkos::abort("Converting DynRankView to a View of mis-matched rank!");)
}

return View<typename RankDataType<T, N>::type, Args...>(
v.data(), v.impl_map().layout());
auto layout = v.impl_map().layout();

if constexpr (std::is_same_v<decltype(layout), Kokkos::LayoutLeft> ||
std::is_same_v<decltype(layout), Kokkos::LayoutRight> ||
std::is_same_v<decltype(layout), Kokkos::LayoutStride> ||
is_layouttiled<decltype(layout)>::value) {
for (int i = N; i < 7; ++i)
layout.dimension[i] = KOKKOS_IMPL_CTOR_DEFAULT_ARG;
}

return View<typename RankDataType<T, N>::type, Args...>(v.data(), layout);
}

template <typename Function, typename... Args>
Expand Down
116 changes: 91 additions & 25 deletions core/src/Kokkos_View.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,25 +75,59 @@ constexpr KOKKOS_INLINE_FUNCTION std::size_t count_valid_integers(
(i6 != KOKKOS_INVALID_INDEX) + (i7 != KOKKOS_INVALID_INDEX);
}

KOKKOS_INLINE_FUNCTION
void runtime_check_rank(const size_t rank, const size_t dyn_rank,
const bool is_void_spec, const size_t i0,
const size_t i1, const size_t i2, const size_t i3,
const size_t i4, const size_t i5, const size_t i6,
const size_t i7, const std::string& label) {
// FIXME Ideally, we would not instantiate this function for every possible View
// type. We should be able to only pass "extent" when we use mdspan.
template <typename View>
KOKKOS_INLINE_FUNCTION void runtime_check_rank(
const View&, const bool is_void_spec, const size_t i0, const size_t i1,
const size_t i2, const size_t i3, const size_t i4, const size_t i5,
const size_t i6, const size_t i7, const char* label) {
(void)(label);

if (is_void_spec) {
const size_t num_passed_args =
count_valid_integers(i0, i1, i2, i3, i4, i5, i6, i7);
// We either allow to pass as many extents as the dynamic rank is, or
// as many extents as the total rank is. In the latter case, the given
// extents for the static dimensions must match the
// compile-time extents.
constexpr int rank = View::rank();
constexpr int dyn_rank = View::rank_dynamic();
const bool n_args_is_dyn_rank = num_passed_args == dyn_rank;
const bool n_args_is_rank = num_passed_args == rank;

if constexpr (rank != dyn_rank) {
if (n_args_is_rank) {
size_t new_extents[8] = {i0, i1, i2, i3, i4, i5, i6, i7};
for (int i = dyn_rank; i < rank; ++i)
if (new_extents[i] != View::static_extent(i)) {
KOKKOS_IF_ON_HOST(
const std::string message =
"The specified run-time extent for Kokkos::View '" +
std::string(label) +
"' does not match the compile-time extent in dimension " +
std::to_string(i) + ". The given extent is " +
std::to_string(new_extents[i]) + " but should be " +
std::to_string(View::static_extent(i)) + ".\n";
Kokkos::abort(message.c_str());)
KOKKOS_IF_ON_DEVICE(
Kokkos::abort(
"The specified run-time extents for a Kokkos::View "
"do not match the compile-time extents.");)
}
}
}

if (num_passed_args != dyn_rank && num_passed_args != rank) {
if (!n_args_is_dyn_rank && !n_args_is_rank) {
KOKKOS_IF_ON_HOST(
const std::string message =
"Constructor for Kokkos View '" + label +
"' has mismatched number of arguments. Number of arguments = " +
"Constructor for Kokkos::View '" + std::string(label) +
"' has mismatched number of arguments. The number "
"of arguments = " +
std::to_string(num_passed_args) +
" but dynamic rank = " + std::to_string(dyn_rank) + " \n";
" neither matches the dynamic rank = " +
std::to_string(dyn_rank) +
" nor the total rank = " + std::to_string(rank) + "\n";
Kokkos::abort(message.c_str());)
KOKKOS_IF_ON_DEVICE(Kokkos::abort("Constructor for Kokkos View has "
"mismatched number of arguments.");)
Expand Down Expand Up @@ -1402,21 +1436,30 @@ class View : public ViewTraits<DataType, Properties...> {
"execution space");
}

size_t i0 = arg_layout.dimension[0];
size_t i1 = arg_layout.dimension[1];
size_t i2 = arg_layout.dimension[2];
size_t i3 = arg_layout.dimension[3];
size_t i4 = arg_layout.dimension[4];
size_t i5 = arg_layout.dimension[5];
size_t i6 = arg_layout.dimension[6];
size_t i7 = arg_layout.dimension[7];

const std::string& alloc_name =
Impl::get_property<Impl::LabelTag>(prop_copy);
Impl::runtime_check_rank(
rank, rank_dynamic,
std::is_same<typename traits::specialize, void>::value, i0, i1, i2, i3,
i4, i5, i6, i7, alloc_name);
#ifdef KOKKOS_ENABLE_DEBUG_BOUNDS_CHECK
if constexpr (std::is_same_v<typename traits::array_layout,
Kokkos::LayoutLeft> ||
std::is_same_v<typename traits::array_layout,
Kokkos::LayoutRight> ||
std::is_same_v<typename traits::array_layout,
Kokkos::LayoutStride> ||
is_layouttiled<typename traits::array_layout>::value) {
size_t i0 = arg_layout.dimension[0];
size_t i1 = arg_layout.dimension[1];
size_t i2 = arg_layout.dimension[2];
size_t i3 = arg_layout.dimension[3];
size_t i4 = arg_layout.dimension[4];
size_t i5 = arg_layout.dimension[5];
size_t i6 = arg_layout.dimension[6];
size_t i7 = arg_layout.dimension[7];

const std::string& alloc_name =
Impl::get_property<Impl::LabelTag>(prop_copy);
Impl::runtime_check_rank(
*this, std::is_same<typename traits::specialize, void>::value, i0, i1,
i2, i3, i4, i5, i6, i7, alloc_name.c_str());
}
#endif

Kokkos::Impl::SharedAllocationRecord<>* record = m_map.allocate_shared(
prop_copy, arg_layout, Impl::ViewCtorProp<P...>::has_execution_space);
Expand Down Expand Up @@ -1445,6 +1488,29 @@ class View : public ViewTraits<DataType, Properties...> {
typename Impl::ViewCtorProp<P...>::pointer_type>::value,
"Constructing View to wrap user memory must supply matching pointer "
"type");

#ifdef KOKKOS_ENABLE_DEBUG_BOUNDS_CHECK
if constexpr (std::is_same_v<typename traits::array_layout,
Kokkos::LayoutLeft> ||
std::is_same_v<typename traits::array_layout,
Kokkos::LayoutRight> ||
std::is_same_v<typename traits::array_layout,
Kokkos::LayoutStride> ||
is_layouttiled<typename traits::array_layout>::value) {
size_t i0 = arg_layout.dimension[0];
size_t i1 = arg_layout.dimension[1];
size_t i2 = arg_layout.dimension[2];
size_t i3 = arg_layout.dimension[3];
size_t i4 = arg_layout.dimension[4];
size_t i5 = arg_layout.dimension[5];
size_t i6 = arg_layout.dimension[6];
size_t i7 = arg_layout.dimension[7];

Impl::runtime_check_rank(
*this, std::is_same<typename traits::specialize, void>::value, i0, i1,
i2, i3, i4, i5, i6, i7, "UNMANAGED");
}
#endif
}

// Simple dimension-only layout
Expand Down
6 changes: 6 additions & 0 deletions core/src/impl/Kokkos_ViewArray.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,12 @@ class ViewMapping<Traits, Kokkos::Array<>> {
return m_impl_offset.m_dim.extent(r);
}

static KOKKOS_INLINE_FUNCTION constexpr size_t static_extent(
const unsigned r) noexcept {
using dim_type = typename offset_type::dimension_type;
return dim_type::static_extent(r);
}

KOKKOS_INLINE_FUNCTION constexpr typename Traits::array_layout layout()
const {
return m_impl_offset.layout();
Expand Down
12 changes: 6 additions & 6 deletions core/unit_test/TestViewAPI.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1004,25 +1004,25 @@ class TestViewAPI {
hView3 hv_3("dView3::HostMirror", N0);
hView4 hv_4("dView4::HostMirror", N0);

dView0 dv_0_1(nullptr, 0);
dView0 dv_0_1(nullptr);
dView0 dv_0_2(hv_0.label(), hv_0.layout());

dView1 dv_1_1(nullptr, 0);
dView1 dv_1_1(nullptr, N0);
dView1 dv_1_2(hv_1.label(), hv_1.layout());

dView2 dv_2_1(nullptr, 0);
dView2 dv_2_1(nullptr, N0);
dView2 dv_2_2(hv_2.label(), hv_2.layout());

dView3 dv_3_1(nullptr, 0);
dView3 dv_3_1(nullptr, N0);
dView3 dv_3_2(hv_3.label(), hv_3.layout());

dView4 dv_4_1(nullptr, 0);
dView4 dv_4_1(nullptr, N0);
dView4 dv_4_2(hv_4.label(), hv_4.layout());
}

static void run_test_contruction_from_layout_2() {
using dView3_0 = Kokkos::View<T ***, device>;
using dView3_1 = Kokkos::View<T * * [N1], device>;
using dView3_1 = Kokkos::View<T * * [N2], device>;
using dView3_2 = Kokkos::View<T * [N1][N2], device>;
using dView3_3 = Kokkos::View<T[N0][N1][N2], device>;

Expand Down

0 comments on commit 0e4a158

Please sign in to comment.