kokkos#6805: add mdspan fixes from kokkos/mdspan#327

nmm0 · Apr 24, 2024 · f0a1687 · f0a1687
1 parent 3f99a8c
commit f0a1687
Show file tree

Hide file tree

Showing 3 changed files with 305 additions and 228 deletions.
diff --git a/tpls/mdspan/include/experimental/__p0009_bits/layout_stride.hpp b/tpls/mdspan/include/experimental/__p0009_bits/layout_stride.hpp
@@ -199,6 +199,12 @@ struct layout_stride {
         return __strides_storage_t{static_cast<index_type>(s[Idxs])...};
       }
 
+      template<class IntegralType>
+      MDSPAN_INLINE_FUNCTION
+      static constexpr const __strides_storage_t fill_strides(mdspan_non_standard_tag, const IntegralType (&s)[extents_type::rank()]) {
+        return __strides_storage_t{static_cast<index_type>(s[Idxs])...};
+      }
+
 #ifdef __cpp_lib_span
       template<class IntegralType>
       MDSPAN_INLINE_FUNCTION
@@ -309,6 +315,44 @@ struct layout_stride {
        */
     }
 
+    MDSPAN_TEMPLATE_REQUIRES(
+      class IntegralTypes,
+      /* requires */ (
+        // MSVC 19.32 does not like using index_type here, requires the typename Extents::index_type
+        // error C2641: cannot deduce template arguments for 'MDSPAN_IMPL_STANDARD_NAMESPACE::layout_stride::mapping'
+        _MDSPAN_TRAIT(std::is_convertible, const std::remove_const_t<IntegralTypes>&, typename Extents::index_type) &&
+        _MDSPAN_TRAIT(std::is_nothrow_constructible, typename Extents::index_type, const std::remove_const_t<IntegralTypes>&)
+      )
+    )
+    MDSPAN_INLINE_FUNCTION
+    constexpr
+    mapping(
+      mdspan_non_standard_tag,
+      extents_type const& e,
+      IntegralTypes (&s)[extents_type::rank()]
+    ) noexcept
+#if defined(_MDSPAN_USE_ATTRIBUTE_NO_UNIQUE_ADDRESS)
+      : __members{
+#else
+      : __base_t(__base_t{__member_pair_t(
+#endif
+          e, __strides_storage_t(__impl::fill_strides(mdspan_non_standard, s))
+#if defined(_MDSPAN_USE_ATTRIBUTE_NO_UNIQUE_ADDRESS)
+        }
+#else
+        )})
+#endif
+    {
+      /*
+       * TODO: check preconditions
+       * - s[i] > 0 is true for all i in the range [0, rank_ ).
+       * - REQUIRED-SPAN-SIZE(e, s) is a representable value of type index_type ([basic.fundamental]).
+       * - If rank_ is greater than 0, then there exists a permutation P of the integers in the
+       *   range [0, rank_), such that s[ pi ] >= s[ pi − 1 ] * e.extent( pi − 1 ) is true for
+       *   all i in the range [1, rank_ ), where pi is the ith element of P.
+       */
+    }
+
 #ifdef __cpp_lib_span
     MDSPAN_TEMPLATE_REQUIRES(
       class IntegralTypes,

diff --git a/tpls/mdspan/include/experimental/__p0009_bits/macros.hpp b/tpls/mdspan/include/experimental/__p0009_bits/macros.hpp
@@ -629,3 +629,8 @@ struct __bools;
 
 // </editor-fold> end Pre-C++14 constexpr }}}1
 //==============================================================================
+
+namespace MDSPAN_IMPL_STANDARD_NAMESPACE {
+constexpr struct mdspan_non_standard_tag {
+} mdspan_non_standard;
+} // namespace MDSPAN_IMPL_STANDARD_NAMESPACE