Add mandate check for strided_slice

kokkos · Jan 24, 2023 · 5702ddc · 5702ddc
1 parent d89cb19
commit 5702ddc
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/include/experimental/__p2630_bits/strided_slice.hpp b/include/experimental/__p2630_bits/strided_slice.hpp
@@ -15,9 +15,18 @@
 //
 //@HEADER
 
+#include <type_traits>
+
 namespace std {
 namespace experimental {
 
+namespace {
+  template<class T>
+  struct __mdspan_is_integral_constant: std::false_type {};
+
+  template<class T, T val>
+  struct __mdspan_is_integral_constant<integral_constant<T,val>>: std::true_type {};
+}
 // Slice Specifier allowing for strides and compile time extent
 template <class OffsetType, class ExtentType, class StrideType>
 struct strided_slice {
@@ -28,6 +37,10 @@ struct strided_slice {
   OffsetType offset;
   ExtentType extent;
   StrideType stride;
+
+  static_assert(is_integral_v<OffsetType> || __mdspan_is_integral_constant<OffsetType>::value);
+  static_assert(is_integral_v<ExtentType> || __mdspan_is_integral_constant<ExtentType>::value);
+  static_assert(is_integral_v<StrideType> || __mdspan_is_integral_constant<StrideType>::value);
 };
 
 } // experimental