kokkos#6805: fix constraints on implicit mdspan conversion and test f…

…or undesirable implicit conversions
nmm0 · May 15, 2024 · 5f4b5fb · 5f4b5fb
1 parent 47d121d
commit 5f4b5fb
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/core/src/Kokkos_View.hpp b/core/src/Kokkos_View.hpp
@@ -1790,7 +1790,11 @@ class View : public ViewTraits<DataType, Properties...> {
   //----------------------------------------
   // Conversion to MDSpan
   template <class OtherElementType, class OtherExtents, class OtherLayoutPolicy,
-            class OtherAccessor>
+            class OtherAccessor,
+            typename = std::enable_if_t<std::is_assignable_v<
+                mdspan<OtherElementType, OtherExtents, OtherLayoutPolicy,
+                       OtherAccessor>,
+                typename Impl::MDSpanViewTraits<traits>::mdspan_type>>>
   KOKKOS_INLINE_FUNCTION constexpr operator mdspan<
       OtherElementType, OtherExtents, OtherLayoutPolicy, OtherAccessor>() {
     using mdspan_type = typename Impl::MDSpanViewTraits<traits>::mdspan_type;

diff --git a/core/unit_test/TestMDSpanConversion.hpp b/core/unit_test/TestMDSpanConversion.hpp
@@ -133,7 +133,40 @@ struct TestViewMDSpanConversion {
     ASSERT_EQ(cvt.mapping(), ref_layout_mapping);
   }
 
+  template <typename ViewType>
+  using natural_mdspan_type_for_view = typename Kokkos::Impl::MDSpanViewTraits<
+      typename ViewType::traits>::mdspan_type;
+
   static void run_test() {
+    // Verify we can only convert to compatible mdspans
+    static_assert(std::is_convertible_v<
+                  Kokkos::View<value_type *>,
+                  natural_mdspan_type_for_view<Kokkos::View<value_type *>>>);
+    static_assert(
+        std::is_convertible_v<
+            Kokkos::View<value_type *>,
+            natural_mdspan_type_for_view<Kokkos::View<const value_type *>>>);
+
+    // Do not cast const away
+    static_assert(!std::is_convertible_v<
+                  Kokkos::View<const value_type *>,
+                  natural_mdspan_type_for_view<Kokkos::View<value_type *>>>);
+
+    // Mismatched dim
+    static_assert(!std::is_convertible_v<
+                  Kokkos::View<value_type *>,
+                  natural_mdspan_type_for_view<Kokkos::View<value_type **>>>);
+
+    // Mismatched layouts
+    static_assert(
+        !std::is_convertible_v<Kokkos::View<value_type **, Kokkos::LayoutLeft>,
+                               natural_mdspan_type_for_view<Kokkos::View<
+                                   value_type **, Kokkos::LayoutRight>>>);
+    static_assert(
+        !std::is_convertible_v<Kokkos::View<value_type **, Kokkos::LayoutRight>,
+                               natural_mdspan_type_for_view<Kokkos::View<
+                                   value_type **, Kokkos::LayoutLeft>>>);
+
     // nvcc doesn't do CTAD properly here, making this way more verbose..
     // LayoutLeft
     test_conversion_from_mdspan<Kokkos::LayoutLeft>(