kokkos#6805: add conversion to mdspan

nmm0 · Mar 7, 2024 · 9e38f34 · 9e38f34
1 parent c173db3
commit 9e38f34
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 0 deletions.
diff --git a/core/src/Kokkos_View.hpp b/core/src/Kokkos_View.hpp
@@ -1770,6 +1770,26 @@ class View : public ViewTraits<DataType, Properties...> {
                  typename traits::array_layout,
                  typename Experimental::Impl::MDSpanViewTraits<
                      traits>::mdspan_type>(mds.mapping())) {}
+
+  //----------------------------------------
+  // Conversion to MDSpan
+private:
+ template <typename U = typename Experimental::Impl::MDSpanViewTraits<
+               traits>::mdspan_type>
+ constexpr auto make_natural_mdspan(
+     std::enable_if_t<!std::is_same_v<
+         Experimental::Impl::UnsupportedKokkosArrayLayout, U>>* = nullptr) {
+    using mdspan_type =
+        typename Experimental::Impl::MDSpanViewTraits<traits>::mdspan_type;
+    return mdspan_type{data(), Experimental::Impl::mapping_from_view_mapping<mdspan_type>(m_map)};
+ }
+
+public:
+
+  template <class OtherElementType, class OtherExtents, class OtherLayoutPolicy, class OtherAccessor>
+  operator mdspan<OtherElementType, OtherExtents, OtherLayoutPolicy, OtherAccessor> () {
+    return make_natural_mdspan();
+  }
 #endif  // KOKKOS_ENABLE_IMPL_MDSPAN
 };
 

diff --git a/core/src/View/MDSpan/Kokkos_MDSpan_Extents.hpp b/core/src/View/MDSpan/Kokkos_MDSpan_Extents.hpp
@@ -115,6 +115,20 @@ KOKKOS_INLINE_FUNCTION auto dimension_from_extent(const Extents &e,
                                                   std::size_t r) noexcept {
   return Extents::static_extent(r) == dynamic_extent ? e.extent(r) : 0;
 }
+
+template <class Extents, class VM, std::size_t... Indices>
+constexpr KOKKOS_INLINE_FUNCTION auto extents_from_view_mapping_impl(
+    const VM &view_mapping, std::index_sequence<Indices...>) {
+  return Extents{view_mapping.extent(Indices)...};
+}
+
+template <class Extents, class VM>
+constexpr KOKKOS_INLINE_FUNCTION auto extents_from_view_mapping(
+    const VM &view_mapping) {
+  static_assert(Extents::rank() == VM::Rank);
+  return extents_from_view_mapping_impl<Extents>(
+      view_mapping, std::make_index_sequence<Extents::rank()>{});
+}
 }  // namespace Kokkos::Experimental::Impl
 
 #endif  // KOKKOS_EXPERIMENTAL_MDSPAN_EXTENTS_HPP
diff --git a/core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp b/core/src/View/MDSpan/Kokkos_MDSpan_Layout.hpp
@@ -96,6 +96,14 @@ KOKKOS_INLINE_FUNCTION auto array_layout_from_mapping(const typename MDSpanType:
       rank > 7 ? dimension_from_extent(ext, 7) : KOKKOS_IMPL_CTOR_DEFAULT_ARG};
 }
 
+template<class MDSpanType, class VM>
+KOKKOS_INLINE_FUNCTION auto mapping_from_view_mapping(const VM &view_mapping) {
+  using mapping_type = typename MDSpanType::mapping_type;
+  using extents_type = typename mapping_type::extents_type;
+
+  return mapping_type(extents_from_view_mapping<extents_type>(view_mapping));
+}
+
 template <class ElementType, class Extents, class LayoutPolicy,
           class AccessorPolicy>
 KOKKOS_INLINE_FUNCTION auto view_offset_from_mdspan(

diff --git a/core/unit_test/view/TestMDSpanConversion.hpp b/core/unit_test/view/TestMDSpanConversion.hpp
@@ -50,6 +50,8 @@ struct TestViewMDSpanConversion {
       ASSERT_EQ(test_view.extent(r), ref.extent(r));
       ASSERT_EQ(test_view.extent(r), exts.extent(r));
     }
+
+    natural_mdspan_type cvt = test_view;
   }
 
   static void run_test() {