kokkos#6805: fix constraints on mdspan converting constructor and all…

…ow runtime-unmanaged (but compile-time managed) views to be constructed from mdspans
nmm0 · Feb 29, 2024 · 21b422b · 21b422b
1 parent 3dec933
commit 21b422b
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 52 deletions.
diff --git a/core/src/Kokkos_View.hpp b/core/src/Kokkos_View.hpp
@@ -373,6 +373,33 @@ struct ViewTraits {
   //------------------------------------
 };
 
+#ifdef KOKKOS_ENABLE_IMPL_MDSPAN
+namespace Experimental::Impl {
+struct UnsupportedKokkosArrayLayout;
+
+template <class Traits, class Enabled = void>
+struct MDSpanViewTraits {
+  using mdspan_type = UnsupportedKokkosArrayLayout;
+};
+
+/**
+ * "Natural" mdspan for a view if the View's ArrayLayout is supported.
+ */
+template <class Traits>
+struct MDSpanViewTraits<
+    Traits, std::void_t<typename Experimental::Impl::LayoutFromArrayLayout<
+                typename Traits::array_layout>::type>> {
+  using index_type   = typename Traits::execution_space::size_type;
+  using extents_type = typename Experimental::Impl::ExtentsFromDataType<
+      index_type, typename Traits::data_type>::type;
+  using mdspan_layout_type = typename Experimental::Impl::LayoutFromArrayLayout<
+      typename Traits::array_layout>::type;
+  using mdspan_type =
+      mdspan<typename Traits::value_type, extents_type, mdspan_layout_type>;
+};
+}  // namespace Experimental::Impl
+#endif  // KOKKOS_ENABLE_IMPL_MDSPAN
+
 /** \class View
  *  \brief View to an array of data.
  *
@@ -1730,33 +1757,17 @@ class View : public ViewTraits<DataType, Properties...> {
   //----------------------------------------
   // MDSpan converting constructors
 #ifdef KOKKOS_ENABLE_IMPL_MDSPAN
-#ifndef KOKKOS_IMPL_TEST_ACCESS
- private:
-#endif  // KOKKOS_IMPL_TEST_ACCESS
-  template <class OtherExtents, class OtherLayoutPolicy, class Accessor>
-  static constexpr inline bool mdspan_conversion_constraints =
-      !traits::is_managed &&
-      std::is_constructible_v<typename traits::value_type,
-                              typename Accessor::reference> &&
-      std::is_assignable_v<typename Accessor::reference,
-                           typename traits::value_type> &&
-      std::is_default_constructible_v<typename traits::value_type> &&
-      std::is_constructible_v<
-          typename traits::array_layout,
-          const typename Experimental::Impl::ArrayLayoutFromLayout<
-              OtherLayoutPolicy>::type&>;
-
  public:
-  template <class OtherElementType, class OtherExtents, class OtherLayoutPolicy,
-            class = std::enable_if_t<mdspan_conversion_constraints<
-                OtherExtents, OtherLayoutPolicy,
-                Kokkos::default_accessor<OtherElementType>>>>
-  MDSPAN_CONDITIONAL_EXPLICIT(
-      !(std::is_convertible_v<Kokkos::default_accessor,
-                              typename traits::value_type>))
-  KOKKOS_INLINE_FUNCTION
-      View(mdspan<OtherElementType, OtherExtents, OtherLayoutPolicy> const& mds)
-      : m_track(), m_map(mds) {}
+  template <typename U = typename Impl::MDSpanViewTraits<traits>::mdspan_type>
+  KOKKOS_INLINE_FUNCTION MDSPAN_CONDITIONAL_EXPLICIT(traits::is_managed) View(
+      const typename Impl::MDSpanViewTraits<traits>::mdspan_type& mds,
+      std::enable_if_t<
+          !std::is_same_v<Impl::UnsupportedKokkosArrayLayout, U>>* = nullptr)
+      : View(mds.data_handle(),
+             Experimental::Impl::array_layout_from_mapping<
+                 typename traits::array_layout,
+                 typename Impl::MDSpanViewTraits<traits>::mdspan_type>(
+                 mds.mapping())) {}
 #endif  // KOKKOS_ENABLE_IMPL_MDSPAN
 };
 

diff --git a/core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp b/core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp
@@ -46,6 +46,24 @@ struct ArrayLayoutFromLayout<Experimental::layout_right_padded<padding_value>> {
       padding = {};
 };
 
+template <class ArrayLayout>
+struct LayoutFromArrayLayout;
+
+template <>
+struct LayoutFromArrayLayout<Kokkos::LayoutLeft> {
+  using type = Experimental::layout_left_padded<dynamic_extent>;
+};
+
+template <>
+struct LayoutFromArrayLayout<Kokkos::LayoutRight> {
+  using type = Experimental::layout_right_padded<dynamic_extent>;
+};
+
+template <>
+struct LayoutFromArrayLayout<Kokkos::LayoutStride> {
+  using type = layout_stride;
+};
+
 template <class T, class Extents, class Layout>
 struct ViewOffsetFromExtents {
   using value_type   = T;
@@ -57,18 +75,17 @@ struct ViewOffsetFromExtents {
       Kokkos::Impl::ViewOffset<typename data_analysis::dimension, array_layout>;
 };
 
-template <class ElementType, class Extents, class LayoutPolicy,
-          class AccessorPolicy>
-KOKKOS_INLINE_FUNCTION auto array_layout_from_mdspan(
-    const mdspan<ElementType, Extents, LayoutPolicy, AccessorPolicy> &mds) {
-  using layout_type = typename ArrayLayoutFromLayout<LayoutPolicy>::type;
-  const auto &ext   = mds.extents();
+template<class ArrayLayout, class MDSpanType>
+KOKKOS_INLINE_FUNCTION auto array_layout_from_mapping(const typename MDSpanType::mapping_type &mapping) {
+  using mapping_type = typename MDSpanType::mapping_type;
+  using extents_type = typename mapping_type::extents_type;
 
-  static constexpr auto rank = Extents::rank();
+  static constexpr auto rank = extents_type::rank();
+  const auto &ext = mapping.extents();
 
   static_assert(rank <= ARRAY_LAYOUT_MAX_RANK,
                 "Unsupported rank for mdspan (must be <= 8)");
-  return layout_type{
+  return ArrayLayout{
       rank > 0 ? dimension_from_extent(ext, 0) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,
       rank > 1 ? dimension_from_extent(ext, 1) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,
       rank > 2 ? dimension_from_extent(ext, 2) : KOKKOS_IMPL_CTOR_DEFAULT_ARG,

diff --git a/core/src/impl/Kokkos_ViewMapping.hpp b/core/src/impl/Kokkos_ViewMapping.hpp
@@ -2971,15 +2971,6 @@ class ViewMapping<
       : m_impl_handle(Impl::get_property<Impl::PointerTag>(arg_prop)),
         m_impl_offset(std::integral_constant<unsigned, 0>(), arg_layout) {}
 
-#ifdef KOKKOS_ENABLE_IMPL_MDSPAN
-  template <class OtherElementType, class OtherExtents, class OtherLayoutPolicy>
-  explicit KOKKOS_INLINE_FUNCTION ViewMapping(
-      Kokkos::mdspan<OtherElementType, OtherExtents, OtherLayoutPolicy> const&
-          mds)
-      : m_impl_handle(mds.data_handle()),
-        m_impl_offset(Experimental::Impl::view_offset_from_mdspan(mds)) {}
-#endif
-
   /**\brief  Assign data */
   KOKKOS_INLINE_FUNCTION
   void assign_data(pointer_type arg_ptr) {

diff --git a/core/unit_test/view/TestMDSpanConversion.hpp b/core/unit_test/view/TestMDSpanConversion.hpp
@@ -15,6 +15,7 @@
 //@HEADER
 
 #include <gtest/gtest.h>
+#include <type_traits>
 
 #define KOKKOS_IMPL_TEST_ACCESS
 
@@ -32,17 +33,10 @@ struct TestViewMDSpanConversion {
                                           const MDSpanExtents &exts) {
     using view_type   = Kokkos::View<DataType, KokkosLayout,
                                    Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
+    using natural_mdspan_type = typename Kokkos::Impl::MDSpanViewTraits<typename view_type::traits>::mdspan_type;
     using mdspan_type = Kokkos::mdspan<value_type, MDSpanExtents, MDSpanLayout>;
 
-    using converted_layout_type =
-        typename Kokkos::Experimental::Impl::ArrayLayoutFromLayout<
-            MDSpanLayout>::type;
-    static_assert(
-        std::is_constructible_v<typename view_type::traits::array_layout,
-                                const converted_layout_type &>);
-    static_assert(
-        view_type::template mdspan_conversion_constraints<
-            MDSpanExtents, MDSpanLayout, Kokkos::default_accessor<value_type>>);
+    static_assert(std::is_convertible_v<mdspan_type, natural_mdspan_type>);
 
     // Manually create an mdspan from ref so we have a valid pointer to play
     // with