Fix incompatible dynamic_extent definition in Kokkos

crtrott · Jun 18, 2024 · a1f1255 · a1f1255
1 parent 0d5cc92
commit a1f1255
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 39 deletions.
diff --git a/core/src/Kokkos_Extents.hpp b/core/src/Kokkos_Extents.hpp
@@ -25,33 +25,40 @@ static_assert(false,
 #include <cstddef>
 #include <type_traits>
 #include <Kokkos_Macros.hpp>
+#ifdef KOKKOS_ENABLE_IMPL_MDSPAN
+#include <mdspan/mdspan.hpp>
+#else
+#include <limits>
+#endif
 
 namespace Kokkos {
-namespace Experimental {
 
-constexpr ptrdiff_t dynamic_extent = -1;
+#ifndef KOKKOS_ENABLE_IMPL_MDSPAN
+constexpr size_t dynamic_extent = std::numeric_limits<size_t>::max();
+#endif
 
-template <ptrdiff_t... ExtentSpecs>
+namespace Experimental {
+
+template <size_t... ExtentSpecs>
 struct Extents {
   /* TODO @enhancement flesh this out more */
 };
 
-template <class Exts, ptrdiff_t NewExtent>
+template <class Exts, size_t NewExtent>
 struct PrependExtent;
 
-template <ptrdiff_t... Exts, ptrdiff_t NewExtent>
+template <size_t... Exts, size_t NewExtent>
 struct PrependExtent<Extents<Exts...>, NewExtent> {
   using type = Extents<NewExtent, Exts...>;
 };
 
-template <class Exts, ptrdiff_t NewExtent>
+template <class Exts, size_t NewExtent>
 struct AppendExtent;
 
-template <ptrdiff_t... Exts, ptrdiff_t NewExtent>
+template <size_t... Exts, size_t NewExtent>
 struct AppendExtent<Extents<Exts...>, NewExtent> {
   using type = Extents<Exts..., NewExtent>;
 };
-
 }  // end namespace Experimental
 
 namespace Impl {
@@ -75,33 +82,32 @@ struct _parse_impl {
 
 // We have to treat the case of int**[x] specially, since it *doesn't* go
 // backwards
-template <class T, ptrdiff_t... ExtentSpec>
+template <class T, size_t... ExtentSpec>
 struct _parse_impl<T*, Kokkos::Experimental::Extents<ExtentSpec...>,
                    std::enable_if_t<_all_remaining_extents_dynamic<T>::value>>
-    : _parse_impl<T, Kokkos::Experimental::Extents<
-                         Kokkos::Experimental::dynamic_extent, ExtentSpec...>> {
-};
+    : _parse_impl<T, Kokkos::Experimental::Extents<Kokkos::dynamic_extent,
+                                                   ExtentSpec...>> {};
 
 // int*(*[x])[y] should still work also (meaning int[][x][][y])
-template <class T, ptrdiff_t... ExtentSpec>
+template <class T, size_t... ExtentSpec>
 struct _parse_impl<
     T*, Kokkos::Experimental::Extents<ExtentSpec...>,
     std::enable_if_t<!_all_remaining_extents_dynamic<T>::value>> {
   using _next = Kokkos::Experimental::AppendExtent<
       typename _parse_impl<T, Kokkos::Experimental::Extents<ExtentSpec...>,
                            void>::type,
-      Kokkos::Experimental::dynamic_extent>;
+      Kokkos::dynamic_extent>;
   using type = typename _next::type;
 };
 
-template <class T, ptrdiff_t... ExtentSpec, unsigned N>
+template <class T, size_t... ExtentSpec, unsigned N>
 struct _parse_impl<T[N], Kokkos::Experimental::Extents<ExtentSpec...>, void>
-    : _parse_impl<
-          T, Kokkos::Experimental::Extents<ExtentSpec...,
-                                           ptrdiff_t(N)>  // TODO @pedantic this
-                                                          // could be a
-                                                          // narrowing cast
-          > {};
+    : _parse_impl<T,
+                  Kokkos::Experimental::Extents<ExtentSpec...,
+                                                size_t(N)>  // TODO @pedantic
+                                                            // this could be a
+                                                            // narrowing cast
+                  > {};
 
 }  // end namespace _parse_view_extents_impl
 
@@ -111,38 +117,34 @@ struct ParseViewExtents {
       DataType, Kokkos::Experimental::Extents<>>::type;
 };
 
-template <class ValueType, ptrdiff_t Ext>
+template <class ValueType, size_t Ext>
 struct ApplyExtent {
   using type = ValueType[Ext];
 };
 
 template <class ValueType>
-struct ApplyExtent<ValueType, Kokkos::Experimental::dynamic_extent> {
+struct ApplyExtent<ValueType, Kokkos::dynamic_extent> {
   using type = ValueType*;
 };
 
-template <class ValueType, unsigned N, ptrdiff_t Ext>
+template <class ValueType, unsigned N, size_t Ext>
 struct ApplyExtent<ValueType[N], Ext> {
   using type = typename ApplyExtent<ValueType, Ext>::type[N];
 };
 
-template <class ValueType, ptrdiff_t Ext>
+template <class ValueType, size_t Ext>
 struct ApplyExtent<ValueType*, Ext> {
   using type = ValueType * [Ext];
 };
 
 template <class ValueType>
-struct ApplyExtent<ValueType*, Kokkos::Experimental::dynamic_extent> {
-  using type =
-      typename ApplyExtent<ValueType,
-                           Kokkos::Experimental::dynamic_extent>::type*;
+struct ApplyExtent<ValueType*, dynamic_extent> {
+  using type = typename ApplyExtent<ValueType, dynamic_extent>::type*;
 };
 
 template <class ValueType, unsigned N>
-struct ApplyExtent<ValueType[N], Kokkos::Experimental::dynamic_extent> {
-  using type =
-      typename ApplyExtent<ValueType,
-                           Kokkos::Experimental::dynamic_extent>::type[N];
+struct ApplyExtent<ValueType[N], dynamic_extent> {
+  using type = typename ApplyExtent<ValueType, dynamic_extent>::type[N];
 };
 
 }  // end namespace Impl

diff --git a/core/src/impl/Kokkos_ViewMapping.hpp b/core/src/impl/Kokkos_ViewMapping.hpp
@@ -3463,7 +3463,7 @@ struct SubViewDataTypeImpl<void, ValueType, Kokkos::Experimental::Extents<>> {
 };
 
 /* for integral args, subview doesn't have that dimension */
-template <class ValueType, ptrdiff_t Ext, ptrdiff_t... Exts, class Integral,
+template <class ValueType, size_t Ext, size_t... Exts, class Integral,
           class... Args>
 struct SubViewDataTypeImpl<
     std::enable_if_t<std::is_integral<std::decay_t<Integral>>::value>,
@@ -3472,7 +3472,7 @@ struct SubViewDataTypeImpl<
                           Kokkos::Experimental::Extents<Exts...>, Args...> {};
 
 /* for ALL slice, subview has the same dimension */
-template <class ValueType, ptrdiff_t Ext, ptrdiff_t... Exts, class... Args>
+template <class ValueType, size_t Ext, size_t... Exts, class... Args>
 struct SubViewDataTypeImpl<void, ValueType,
                            Kokkos::Experimental::Extents<Ext, Exts...>,
                            Kokkos::ALL_t, Args...>
@@ -3483,7 +3483,7 @@ struct SubViewDataTypeImpl<void, ValueType,
  * static sizes */
 /* Since we don't allow interleaving of dynamic and static extents, make all of
  * the dimensions to the left dynamic  */
-template <class ValueType, ptrdiff_t Ext, ptrdiff_t... Exts, class PairLike,
+template <class ValueType, size_t Ext, size_t... Exts, class PairLike,
           class... Args>
 struct SubViewDataTypeImpl<
     std::enable_if_t<is_pair_like<PairLike>::value>, ValueType,

diff --git a/core/unit_test/TestViewSubview.hpp b/core/unit_test/TestViewSubview.hpp
@@ -2294,9 +2294,8 @@ template <class Space>
 struct TestExtentsStaticTests {
   using test1 = typename static_expect_same<
       /* expected */
-      Kokkos::Experimental::Extents<Kokkos::Experimental::dynamic_extent,
-                                    Kokkos::Experimental::dynamic_extent, 1, 2,
-                                    3>,
+      Kokkos::Experimental::Extents<Kokkos::dynamic_extent,
+                                    Kokkos::dynamic_extent, 1, 2, 3>,
       /* actual */
       typename Kokkos::Impl::ParseViewExtents<double* * [1][2][3]>::type>::type;