Fix LWG Issue 4060 (fix submdspan to prevent invalid pointer creation) (

#323) * Add failing test for Issue 4060 LWG Issue 4060 ("submdspan preconditions do not forbid creating invalid pointer") currently fails, because submdspan does not correctly handle the case where first_ of one or more slices is out of bounds. This commit adds two tests taken straight from the Issue. As expected, one test passes and the other test fails. However, both tests trigger undefined behavior for the layout mapping. * layout_right::operator(): Add debug-only bounds checking * Implement fix for LWG Issue 4060 * Make index bounds checking C++17 - friendly * Fix previous fix for only some slices at upper bound * Remove unnecessary tuple construction * Remove explicit use of fold expressions ... in favor of macros. In theory, this would help us backport submdspan to C++14. * Add bounds checking to layout_left & layout_stride * Add bounds checking to layout_left_padded & layout_right_padded * Remove no-longer-needed compiler feature tests * Fix Issue 4060 test; test layout_left & layout_stride too * Improve Issue 4060 test * Guard use of assert with _MDSPAN_DEBUG ...thus following the pattern of other uses of assert in the same header file. * Respond to review feedback * Add MDSPAN_INLINE_FUNCTION to bounds-checking functions * Pass indices by value to bounds-checking functions * Don't use std::forward for slices (use const ref instead) * Use required_span_size() in test * Replace bool_constant (since it's C++17) Co-authored-by: Damien L-G <[email protected]> --------- Co-authored-by: Damien L-G <[email protected]>
kokkos · May 3, 2024 · a9c54cc · ldionne · Nov 4, 2024 · a9c54cc
1 parent f0ba298
commit a9c54cc
Show file tree

Hide file tree

Showing 7 changed files with 235 additions and 9 deletions.
diff --git a/include/experimental/__p0009_bits/extents.hpp b/include/experimental/__p0009_bits/extents.hpp
@@ -22,6 +22,7 @@
 #endif
 #include <array>
 
+#include <cassert>
 #include <cinttypes>
 
 namespace MDSPAN_IMPL_STANDARD_NAMESPACE {
@@ -614,5 +615,75 @@ static
 #endif
 constexpr bool __is_extents_v = __is_extents<T>::value;
 
+template<class InputIndexType, class ExtentsIndexType>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_lower_bound(InputIndexType user_index,
+                  ExtentsIndexType /* current_extent */,
+                  std::true_type /* is_signed */)
+{
+  (void) user_index; // prevent unused variable warning
+#ifdef _MDSPAN_DEBUG
+  assert(static_cast<ExtentsIndexType>(user_index) >= 0);
+#endif
+}
+
+template<class InputIndexType, class ExtentsIndexType>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_lower_bound(InputIndexType /* user_index */,
+                  ExtentsIndexType /* current_extent */,
+                  std::false_type /* is_signed */)
+{}
+
+template<class InputIndexType, class ExtentsIndexType>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_upper_bound(InputIndexType user_index,
+                  ExtentsIndexType current_extent)
+{
+  (void) user_index; // prevent unused variable warnings
+  (void) current_extent;
+#ifdef _MDSPAN_DEBUG
+  assert(static_cast<ExtentsIndexType>(user_index) < current_extent);
+#endif
+}
+
+template<class InputIndex, class ExtentsIndexType>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_one_index(InputIndex user_index,
+                ExtentsIndexType current_extent)
+{
+  check_lower_bound(user_index, current_extent,
+    std::integral_constant<bool, std::is_signed_v<ExtentsIndexType>>{});
+  check_upper_bound(user_index, current_extent);
+}
+
+template<size_t ... RankIndices,
+         class ExtentsIndexType, size_t ... Exts,
+         class ... Indices>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_all_indices_helper(std::index_sequence<RankIndices...>,
+                         const extents<ExtentsIndexType, Exts...>& exts,
+                         Indices... indices)
+{
+  _MDSPAN_FOLD_COMMA(
+    (check_one_index(indices, exts.extent(RankIndices)))
+  );
+}
+
+template<class ExtentsIndexType, size_t ... Exts,
+         class ... Indices>
+MDSPAN_INLINE_FUNCTION
+constexpr void
+check_all_indices(const extents<ExtentsIndexType, Exts...>& exts,
+                  Indices... indices)
+{
+  check_all_indices_helper(std::make_index_sequence<sizeof...(Indices)>(),
+                           exts, indices...);
+}
+
 } // namespace detail
 } // namespace MDSPAN_IMPL_STANDARD_NAMESPACE
diff --git a/include/experimental/__p0009_bits/layout_left.hpp b/include/experimental/__p0009_bits/layout_left.hpp
@@ -194,6 +194,9 @@ class layout_left::mapping {
     )
     _MDSPAN_HOST_DEVICE
     constexpr index_type operator()(Indices... idxs) const noexcept {
+#if ! defined(NDEBUG)
+      detail::check_all_indices(this->extents(), idxs...);
+#endif // ! NDEBUG
       return __compute_offset(__rank_count<0, extents_type::rank()>(), static_cast<index_type>(idxs)...);
     }
 

diff --git a/include/experimental/__p0009_bits/layout_right.hpp b/include/experimental/__p0009_bits/layout_right.hpp
@@ -195,6 +195,9 @@ class layout_right::mapping {
     )
     _MDSPAN_HOST_DEVICE
     constexpr index_type operator()(Indices... idxs) const noexcept {
+#if ! defined(NDEBUG)
+      detail::check_all_indices(this->extents(), idxs...);
+#endif // ! NDEBUG
       return __compute_offset(__rank_count<0, extents_type::rank()>(), static_cast<index_type>(idxs)...);
     }
 

diff --git a/include/experimental/__p0009_bits/layout_stride.hpp b/include/experimental/__p0009_bits/layout_stride.hpp
@@ -434,6 +434,9 @@ struct layout_stride {
     )
     MDSPAN_FORCE_INLINE_FUNCTION
     constexpr index_type operator()(Indices... idxs) const noexcept {
+#if ! defined(NDEBUG)
+      detail::check_all_indices(this->extents(), idxs...);
+#endif // ! NDEBUG
       return static_cast<index_type>(__impl::_call_op_impl(*this, static_cast<index_type>(idxs)...));
     }
 

diff --git a/include/experimental/__p2630_bits/submdspan_mapping.hpp b/include/experimental/__p2630_bits/submdspan_mapping.hpp
@@ -31,6 +31,47 @@ template <class LayoutMapping> struct submdspan_mapping_result {
 };
 
 namespace detail {
+
+// We use const Slice& and not Slice&& because the various
+// submdspan_mapping_impl overloads use their slices arguments
+// multiple times.  This makes perfect forwarding not useful, but we
+// still don't want to pass those (possibly of size 64 x 3 bits)
+// objects by value.
+template<class IndexType,
+         class Slice>
+MDSPAN_INLINE_FUNCTION
+constexpr bool
+one_slice_out_of_bounds(const IndexType& extent, const Slice& slice)
+{
+  return detail::first_of(slice) == extent;
+}
+
+template<size_t ... RankIndices,
+         class IndexType, size_t ... Exts,
+         class ... Slices>
+MDSPAN_INLINE_FUNCTION
+constexpr bool
+any_slice_out_of_bounds_helper(std::index_sequence<RankIndices...>,
+                               const extents<IndexType, Exts...>& exts,
+                               const Slices& ... slices)
+{
+  return _MDSPAN_FOLD_OR(
+    (one_slice_out_of_bounds(exts.extent(RankIndices), slices))
+  );
+}
+
+template<class IndexType, size_t ... Exts,
+         class ... Slices>
+MDSPAN_INLINE_FUNCTION
+constexpr bool
+any_slice_out_of_bounds(const extents<IndexType, Exts...>& exts,
+                        const Slices& ... slices)
+{
+  return any_slice_out_of_bounds_helper(
+    std::make_index_sequence<sizeof...(Slices)>(),
+    exts, slices...);
+}
+
 // constructs sub strides
 template <class SrcMapping, class... slice_strides, size_t... InvMapIdxs>
 MDSPAN_INLINE_FUNCTION
@@ -111,11 +152,19 @@ layout_left::mapping<Extents>::submdspan_mapping_impl(SliceSpecifiers... slices)
       std::conditional_t<preserve_layout, layout_left, layout_stride>;
   using dst_mapping_t = typename dst_layout_t::template mapping<dst_ext_t>;
 
+  // Figure out if any slice's lower bound equals the corresponding extent.
+  // If so, bypass evaluating the layout mapping.  This fixes LWG Issue 4060.
+  const bool out_of_bounds =
+    detail::any_slice_out_of_bounds(this->extents(), slices...);
+  auto offset = static_cast<size_t>(
+    out_of_bounds ?
+    this->required_span_size() :
+    this->operator()(detail::first_of(slices)...)
+  );
+
   if constexpr (std::is_same_v<dst_layout_t, layout_left>) {
     // layout_left case
-    return submdspan_mapping_result<dst_mapping_t>{
-        dst_mapping_t(dst_ext),
-        static_cast<size_t>(this->operator()(detail::first_of(slices)...))};
+    return submdspan_mapping_result<dst_mapping_t>{dst_mapping_t(dst_ext), offset};
   } else {
     // layout_stride case
     auto inv_map = detail::inv_map_rank(
@@ -132,7 +181,7 @@ layout_left::mapping<Extents>::submdspan_mapping_impl(SliceSpecifiers... slices)
     #else
                                    std::tuple{detail::stride_of(slices)...})),
     #endif
-        static_cast<size_t>(this->operator()(detail::first_of(slices)...))};
+        offset};
   }
 #if defined(__NVCC__) && !defined(__CUDA_ARCH__) && defined(__GNUC__)
   __builtin_unreachable();
@@ -218,11 +267,19 @@ layout_right::mapping<Extents>::submdspan_mapping_impl(
       std::conditional_t<preserve_layout, layout_right, layout_stride>;
   using dst_mapping_t = typename dst_layout_t::template mapping<dst_ext_t>;
 
+  // Figure out if any slice's lower bound equals the corresponding extent.
+  // If so, bypass evaluating the layout mapping.  This fixes LWG Issue 4060.
+  const bool out_of_bounds =
+    detail::any_slice_out_of_bounds(this->extents(), slices...);
+  auto offset = static_cast<size_t>(
+    out_of_bounds ?
+    this->required_span_size() :
+    this->operator()(detail::first_of(slices)...)
+  );
+
   if constexpr (std::is_same_v<dst_layout_t, layout_right>) {
     // layout_right case
-    return submdspan_mapping_result<dst_mapping_t>{
-        dst_mapping_t(dst_ext),
-        static_cast<size_t>(this->operator()(detail::first_of(slices)...))};
+    return submdspan_mapping_result<dst_mapping_t>{dst_mapping_t(dst_ext), offset};
   } else {
     // layout_stride case
     auto inv_map = detail::inv_map_rank(
@@ -239,7 +296,7 @@ layout_right::mapping<Extents>::submdspan_mapping_impl(
     #else
                                    std::tuple{detail::stride_of(slices)...})),
     #endif
-        static_cast<size_t>(this->operator()(detail::first_of(slices)...))};
+        offset};
   }
 #if defined(__NVCC__) && !defined(__CUDA_ARCH__) && defined(__GNUC__)
   __builtin_unreachable();
@@ -273,6 +330,17 @@ layout_stride::mapping<Extents>::submdspan_mapping_impl(
       std::index_sequence<>(),
       slices...);
   using dst_mapping_t = typename layout_stride::template mapping<dst_ext_t>;
+
+  // Figure out if any slice's lower bound equals the corresponding extent.
+  // If so, bypass evaluating the layout mapping.  This fixes LWG Issue 4060.
+  const bool out_of_bounds =
+    detail::any_slice_out_of_bounds(this->extents(), slices...);
+  auto offset = static_cast<size_t>(
+    out_of_bounds ?
+    this->required_span_size() :
+    this->operator()(detail::first_of(slices)...)
+  );
+
   return submdspan_mapping_result<dst_mapping_t>{
       dst_mapping_t(dst_ext, detail::construct_sub_strides(
                                  *this, inv_map,
@@ -283,7 +351,7 @@ layout_stride::mapping<Extents>::submdspan_mapping_impl(
 #else
                                  std::tuple(detail::stride_of(slices)...))),
 #endif
-      static_cast<size_t>(this->operator()(detail::first_of(slices)...))};
+      offset};
 }
 
 } // namespace MDSPAN_IMPL_STANDARD_NAMESPACE
diff --git a/include/experimental/__p2642_bits/layout_padded.hpp b/include/experimental/__p2642_bits/layout_padded.hpp
@@ -382,6 +382,9 @@ class layout_left_padded<PaddingValue>::mapping {
   )
   constexpr size_t operator()(_Indices... idxs) const noexcept
   {
+#if ! defined(NDEBUG)
+    ::MDSPAN_IMPL_STANDARD_NAMESPACE::detail::check_all_indices(this->extents(), idxs...);
+#endif // ! NDEBUG
     return compute_offset(std::index_sequence_for<_Indices...>{}, idxs...);
   }
 

diff --git a/tests/test_submdspan.cpp b/tests/test_submdspan.cpp
@@ -275,3 +275,78 @@ TYPED_TEST(TestSubMDSpan, submdspan_return_type) {
                 "SubMDSpan: wrong return type");
   __MDSPAN_TESTS_RUN_TEST(TestFixture::run());
 }
+
+TEST(TestSubmdspanIssue4060, Rank1) {
+  auto x = std::array<int, 3>{};
+  auto A = Kokkos::mdspan{x.data(), Kokkos::extents{3}};
+  ASSERT_EQ(A.mapping().required_span_size(), 3);
+  auto B = Kokkos::submdspan(A, std::tuple{3, 3});
+
+  ASSERT_EQ(B.rank(), 1u);
+  EXPECT_EQ(B.extent(0), 0);
+  EXPECT_EQ(B.data_handle(), x.data() + A.mapping().required_span_size());
+}
+
+template<class MappingType>
+void test_submdspan_issue4060_rank2_all(const MappingType& mapping)
+{
+  auto y = std::array<int, 9>{};
+  ASSERT_EQ(mapping.extents().rank(), 2u);
+  ASSERT_EQ(mapping.required_span_size(), y.size());
+  auto C = Kokkos::mdspan{y.data(), mapping};
+  auto D = Kokkos::submdspan(C, std::tuple{3u, 3u}, std::tuple{3u, 3u});
+
+  ASSERT_EQ(D.rank(), 2u);
+  EXPECT_EQ(D.extent(0), 0);
+  EXPECT_EQ(D.extent(1), 0);
+  EXPECT_EQ(D.data_handle(), y.data() + mapping.required_span_size());
+}
+
+TEST(TestSubmdspanIssue4060, Rank2_all) {
+  Kokkos::dextents<size_t, 2> exts{3u, 3u};
+  {
+    using mapping_type = Kokkos::layout_left::mapping<Kokkos::dextents<size_t, 2>>;
+    test_submdspan_issue4060_rank2_all(mapping_type{exts});
+  }
+  {
+    using mapping_type = Kokkos::layout_right::mapping<Kokkos::dextents<size_t, 2>>;
+    test_submdspan_issue4060_rank2_all(mapping_type{exts});
+  }
+  {
+    using mapping_type = Kokkos::layout_stride::mapping<Kokkos::dextents<size_t, 2>>;
+    std::array<size_t, 2> strides{1u, 3u};
+    test_submdspan_issue4060_rank2_all(mapping_type{exts, strides});
+  }
+}
+
+template<class MappingType>
+void test_submdspan_issue4060_rank2_one(const MappingType& mapping)
+{
+  auto y = std::array<int, 9>{};
+  ASSERT_EQ(mapping.extents().rank(), 2u);
+  ASSERT_EQ(mapping.required_span_size(), y.size());
+  auto C = Kokkos::mdspan{y.data(), mapping};
+  auto D = Kokkos::submdspan(C, std::tuple{0u, 3u}, std::tuple{3u, 3u});
+
+  ASSERT_EQ(D.rank(), 2u);
+  EXPECT_EQ(D.extent(0), 3u);
+  EXPECT_EQ(D.extent(1), 0);
+  EXPECT_EQ(D.data_handle(), y.data() + mapping.required_span_size());
+}
+
+TEST(TestSubmdspanIssue4060, Rank2_one) {
+  Kokkos::dextents<size_t, 2> exts{3u, 3u};
+  {
+    using mapping_type = Kokkos::layout_left::mapping<Kokkos::dextents<size_t, 2>>;
+    test_submdspan_issue4060_rank2_one(mapping_type{exts});
+  }
+  {
+    using mapping_type = Kokkos::layout_right::mapping<Kokkos::dextents<size_t, 2>>;
+    test_submdspan_issue4060_rank2_one(mapping_type{exts});
+  }
+  {
+    using mapping_type = Kokkos::layout_stride::mapping<Kokkos::dextents<size_t, 2>>;
+    std::array<size_t, 2> strides{1u, 3u};
+    test_submdspan_issue4060_rank2_one(mapping_type{exts, strides});
+  }
+}