mdspan: add conversion from extents to datatype

kokkos · Nov 17, 2022 · e36a0aa · e36a0aa
1 parent c3948fa
commit e36a0aa
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 5 deletions.
diff --git a/core/src/View/MDSpan/Kokkos_MDSpan_Extents.hpp b/core/src/View/MDSpan/Kokkos_MDSpan_Extents.hpp
@@ -65,8 +65,12 @@ namespace Kokkos {
 namespace Impl {
 template <class DataType>
 struct ViewArrayAnalysis;
+
 template <std::size_t... Vals>
 struct ViewDimension;
+
+template <class T, class Dim>
+struct ViewDataType;
 }
 
 namespace Experimental::Impl {
@@ -91,28 +95,51 @@ struct ExtentFromDimension<std::size_t{0}> {
   static constexpr inline std::size_t value = std::experimental::dynamic_extent;
 };
 
+template<std::size_t N>
+struct DimensionFromExtent {
+  static constexpr inline std::size_t value = N;
+};
+
+template<>
+struct DimensionFromExtent<std::experimental::dynamic_extent> {
+  static constexpr inline std::size_t value = std::size_t{0};
+};
+
 template<class SizeType, class Dimension, class Indices>
 struct ExtentsFromDimension;
 
 template<class SizeType, class Dimension, std::size_t... Indices>
 struct ExtentsFromDimension<SizeType, Dimension, std::index_sequence<Indices...>> {
   using dimension_type = Dimension;
   using type = std::experimental::extents<SizeType, ExtentFromDimension<Dimension::static_extent(Indices)>::value...>;
+};
+
+template<class Extents, class Indices>
+struct DimensionsFromExtent;
 
-  static constexpr type construct( const dimension_type &_dim ) noexcept {
-    return type{ _dim.extent( _dim.extent( Indices )... ) };
-  };
+template<class Extents, std::size_t... Indices>
+struct DimensionsFromExtent<Extents, std::index_sequence<Indices...>> {
+  using extents_type = Extents;
+  using type = ::Kokkos::Impl::ViewDimension<DimensionFromExtent<extents_type::static_extent(Indices)>::value...>;
 };
 
 template<class DataType>
-struct ExtentsFromDataType
-{
+struct ExtentsFromDataType {
   using array_analysis = ::Kokkos::Impl::ViewArrayAnalysis<DataType>;
   using size_type = std::size_t; // Mirrors Kokkos::View's size type
   using dimension_type = typename array_analysis::dimension;
 
   using type = typename ExtentsFromDimension<size_type, dimension_type, std::make_index_sequence<dimension_type::rank>>::type;
 };
+
+template<class T, class Extents>
+struct DataTypeFromExtents {
+  using extents_type = Extents;
+  using dimension_type = typename DimensionsFromExtent<Extents, std::make_index_sequence<extents_type::rank()>>::type;
+
+  // Will cause a compile error if it is malformed (i.e. dynamic after static)
+  using type = typename ::Kokkos::Impl::ViewDataType<T, dimension_type>::type;
+};
 }
 }
 

diff --git a/core/unit_test/TestViewMDSpan.hpp b/core/unit_test/TestViewMDSpan.hpp
@@ -8,6 +8,7 @@
 #ifdef KOKKOS_ENABLE_IMPL_MDSPAN
 
 namespace Test {
+// Conversion from DataType to extents
 // 0-rank view
 static_assert(std::is_same_v<typename Kokkos::Experimental::Impl::ExtentsFromDataType<double>::type,
                              std::experimental::extents<std::size_t>>);
@@ -31,6 +32,45 @@ static_assert(std::is_same_v<typename Kokkos::Experimental::Impl::ExtentsFromDat
                                                         std::size_t{3},
                                                         std::size_t{2},
                                                         std::size_t{8}>>);
+
+// Conversion from extents to DataType
+// 0-rank extents
+static_assert(std::is_same_v<double,
+                             typename Kokkos::Experimental::Impl::DataTypeFromExtents<double,
+                               std::experimental::extents<std::size_t>>::type
+                               >);
+
+// only dynamic
+static_assert(std::is_same_v<double****,
+                             typename Kokkos::Experimental::Impl::DataTypeFromExtents<double,
+                               std::experimental::extents<std::size_t,
+                                std::experimental::dynamic_extent,
+                                std::experimental::dynamic_extent,
+                                std::experimental::dynamic_extent,
+                                std::experimental::dynamic_extent
+                               >
+                             >::type
+                             >);
+
+// only static
+static_assert(std::is_same_v<double[7][5][3],
+                             typename Kokkos::Experimental::Impl::DataTypeFromExtents<double,
+                               std::experimental::extents<std::size_t, 7, 5, 3>
+                             >::type
+                             >);
+
+// both dynamic and static
+static_assert(std::is_same_v<double***[20][45],
+                             typename Kokkos::Experimental::Impl::DataTypeFromExtents<double,
+                               std::experimental::extents<std::size_t,
+                                std::experimental::dynamic_extent,
+                                std::experimental::dynamic_extent,
+                                std::experimental::dynamic_extent,
+                                20,
+                                45
+                               >
+                             >::type
+                             >);
 }  // namespace Test
 
 #endif  // KOKKOS_ENABLE_IMPL_MDSPAN