New View: conversion to mdspan progress

nmm0 · Jul 13, 2024 · 9425083 · 9425083
1 parent 6919696
commit 9425083
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 0 deletions.
diff --git a/core/src/View/Kokkos_BasicView.hpp b/core/src/View/Kokkos_BasicView.hpp
@@ -158,6 +158,41 @@ class BasicView
       SliceSpecifiers... slices)
       : mdspan_type(submdspan(
             src_view, Impl::transform_kokkos_slice_to_mdspan_slice(slices)...)) {}
+
+ public: 
+  //----------------------------------------
+  // Conversion to MDSpan
+  template <class OtherElementType, class OtherExtents, class OtherLayoutPolicy,
+            class OtherAccessor,
+            typename = std::enable_if_t<std::is_assignable_v<
+                mdspan<OtherElementType, OtherExtents, OtherLayoutPolicy,
+                       OtherAccessor>, mdspan_type>>>
+  KOKKOS_INLINE_FUNCTION constexpr operator mdspan<
+      OtherElementType, OtherExtents, OtherLayoutPolicy, OtherAccessor>() {
+    return mdspan_type(*this);
+  }
+
+  template <class OtherAccessorType = 
+    Kokkos::default_accessor<typename mdspan_type::element_type>,
+      //Impl::SpaceAwareAccessor<
+      //          memory_space,
+      //          Kokkos::default_accessor<typename mdspan_type::element_type>>,
+            typename = std::enable_if_t<std::is_assignable_v<
+                typename mdspan_type::data_handle_type,
+                typename OtherAccessorType::data_handle_type>>>
+  KOKKOS_INLINE_FUNCTION constexpr auto to_mdspan(
+      const OtherAccessorType& other_accessor =
+          static_cast<OtherAccessorType>(mdspan_type::accessor())) {
+    using ret_mdspan_type =
+        mdspan<typename mdspan_type::element_type,
+               typename mdspan_type::extents_type,
+               typename mdspan_type::layout_type, OtherAccessorType>;
+    return ret_mdspan_type(static_cast<typename OtherAccessorType::data_handle_type>(mdspan_type::data_handle()), mdspan_type::mapping(), other_accessor);
+  }
+
+  void assign_data(element_type* ptr) {
+    mdspan_type::operator =(mdspan_type{typename mdspan_type::data_handle_type(ptr), mdspan_type::mapping(), mdspan_type::accessor()});
+  }
 
  private:
   template <typename E, bool AllowPadding, bool Initialize>

diff --git a/core/src/View/MDSpan/Kokkos_MDSpan_Accessor.hpp b/core/src/View/MDSpan/Kokkos_MDSpan_Accessor.hpp
@@ -70,6 +70,13 @@ struct SpaceAwareAccessor {
   KOKKOS_FUNCTION
   explicit operator NestedAccessor() const { return nested_acc; }
 
+  template<class OtherElementType,
+   class = std::enable_if_t<std::is_convertible_v<
+                element_type(*) [], OtherElementType (*)[]> &&
+           std::is_convertible_v<nested_accessor_type, Kokkos::default_accessor<OtherElementType>>>>
+  KOKKOS_FUNCTION
+  operator Kokkos::default_accessor<OtherElementType>() const { return nested_acc; }
+
   KOKKOS_FUNCTION
   constexpr reference access(data_handle_type p, size_t i) const noexcept {
     Kokkos::Impl::runtime_check_memory_access_violation<memory_space>(
@@ -137,6 +144,14 @@ struct SpaceAwareAccessor<AnonymousSpace, NestedAccessor> {
 
   KOKKOS_FUNCTION
   explicit operator NestedAccessor() const { return nested_acc; }
+
+  template<class OtherElementType,
+   class = std::enable_if_t<std::is_convertible_v<
+                element_type(*) [], OtherElementType (*)[]> &&
+           std::is_convertible_v<nested_accessor_type, Kokkos::default_accessor<OtherElementType>>>>
+  KOKKOS_FUNCTION
+  operator Kokkos::default_accessor<OtherElementType>() const { return nested_acc; }
+
 
   KOKKOS_FUNCTION
   constexpr reference access(data_handle_type p, size_t i) const noexcept {

diff --git a/core/unit_test/default/TestDefaultDeviceDevelop.cpp b/core/unit_test/default/TestDefaultDeviceDevelop.cpp
@@ -40,5 +40,8 @@ TEST(defaultdevicetype, development_test) {
   auto sub_a = Kokkos::submdspan(mds, std::pair{1,3}); 
   auto sub_b = Kokkos::submdspan(mds, std::array{1,3}); 
   auto sub_c = Kokkos::submdspan(mds, Kokkos::pair{1,3}); 
+  auto acc = c.accessor();
+  const decltype(acc) acc_const = acc;
+  const Kokkos::default_accessor<float> acc_def = acc_const;//static_cast<Kokkos::default_accessor<float>>(acc_const);
 }
 }  // namespace Test