Support and add test for compile time zero length/stride index range

kokkos · Jan 9, 2023 · 3d40324 · 3d40324
1 parent 4fef2da
commit 3d40324
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/include/experimental/__p2630_bits/submdspan_extents.hpp b/include/experimental/__p2630_bits/submdspan_extents.hpp
@@ -173,7 +173,9 @@ template <class IndexT, class T0, T0 v0, class T1, T1 v1>
 MDSPAN_INLINE_FUNCTION
 constexpr auto divide(const integral_constant<T0, v0> &,
                       const integral_constant<T1, v1> &) {
-  return integral_constant<IndexT, v0 / v1>();
+  // cutting short division by zero
+  // this is used for strided_index_range with zero extent/stride
+  return integral_constant<IndexT, v0 == 0 ? 0 : v0 / v1>();
 }
 
 // multiply which can deal with integral constant preservation

diff --git a/tests/test_submdspan.cpp b/tests/test_submdspan.cpp
@@ -157,6 +157,7 @@ using submdspan_test_types =
     , std::tuple<stdex::layout_left, stdex::layout_stride,  stdex::extents<size_t,6,4,5,6,7,8>, args_t<6,4,5,6,7,8>, stdex::extents<size_t,4,dyn,7>, int, stdex::full_extent_t, std::pair<int,int>, int, stdex::full_extent_t, int>
     // layout_right to layout_stride
     , std::tuple<stdex::layout_right, stdex::layout_stride, stdex::dextents<size_t,1>,          args_t<10>,          stdex::dextents<size_t,1>, stdex::strided_index_range<int,int,int>>
+    , std::tuple<stdex::layout_right, stdex::layout_stride, stdex::dextents<size_t,1>,          args_t<10>,          stdex::extents<size_t,0>,  stdex::strided_index_range<int,std::integral_constant<int,0>,std::integral_constant<int,0>>>
     , std::tuple<stdex::layout_right, stdex::layout_stride, stdex::dextents<size_t,2>,          args_t<10,20>,       stdex::dextents<size_t,1>, stdex::strided_index_range<int,int,int>, int>
     , std::tuple<stdex::layout_right, stdex::layout_stride, stdex::dextents<size_t,2>,          args_t<10,20>,       stdex::dextents<size_t,2>, stdex::full_extent_t, std::pair<int,int>>
     , std::tuple<stdex::layout_right, stdex::layout_stride, stdex::dextents<size_t,2>,          args_t<10,20>,       stdex::dextents<size_t,2>, std::pair<int,int>, stdex::strided_index_range<int,int,int>>
@@ -199,15 +200,20 @@ struct TestSubMDSpan<
     return 2;
   }
   MDSPAN_INLINE_FUNCTION
-  static std::pair<int,int> create_slice_arg(std::pair<int,int>) {
+  static auto create_slice_arg(std::pair<int,int>) {
     return std::pair<int,int>(1,3);
   }
   MDSPAN_INLINE_FUNCTION
-  static stdex::strided_index_range<int,int,int> create_slice_arg(stdex::strided_index_range<int,int,int>) {
+  static auto create_slice_arg(stdex::strided_index_range<int,int,int>) {
     return stdex::strided_index_range<int,int,int>{1,3,2};
   }
+  template<int Ext, int Stride>
   MDSPAN_INLINE_FUNCTION
-  static stdex::full_extent_t create_slice_arg(stdex::full_extent_t) {
+  static auto create_slice_arg(stdex::strided_index_range<int,std::integral_constant<int, Ext>, std::integral_constant<int, Stride>>) {
+    return stdex::strided_index_range<int,std::integral_constant<int, Ext>, std::integral_constant<int, Stride>>{1};
+  }
+  MDSPAN_INLINE_FUNCTION
+  static auto create_slice_arg(stdex::full_extent_t) {
     return stdex::full_extent;
   }
 
@@ -223,11 +229,18 @@ struct TestSubMDSpan<
   }
   template<class SrcExtents, class SubExtents, class ... SliceArgs>
   MDSPAN_INLINE_FUNCTION
-  static bool match_expected_extents(int src_idx, int sub_idx, SrcExtents src_ext, SubExtents sub_ext, stdex::strided_index_range<int,int,int> p, SliceArgs ... slices) {
+  static bool match_expected_extents(int src_idx, int sub_idx, SrcExtents src_ext, SubExtents sub_ext,
+                                     stdex::strided_index_range<int,int,int> p, SliceArgs ... slices) {
     return (sub_ext.extent(sub_idx)==(p.extent+p.stride-1)/p.stride) && match_expected_extents(++src_idx, ++sub_idx, src_ext, sub_ext, slices...);
   }
   template<class SrcExtents, class SubExtents, class ... SliceArgs>
   MDSPAN_INLINE_FUNCTION
+  static bool match_expected_extents(int src_idx, int sub_idx, SrcExtents src_ext, SubExtents sub_ext,
+                                     stdex::strided_index_range<int,std::integral_constant<int, 0>,std::integral_constant<int,0>> p, SliceArgs ... slices) {
+    return (sub_ext.extent(sub_idx)==0) && match_expected_extents(++src_idx, ++sub_idx, src_ext, sub_ext, slices...);
+  }
+  template<class SrcExtents, class SubExtents, class ... SliceArgs>
+  MDSPAN_INLINE_FUNCTION
   static bool match_expected_extents(int src_idx, int sub_idx, SrcExtents src_ext, SubExtents sub_ext, stdex::full_extent_t, SliceArgs ... slices) {
     return (sub_ext.extent(sub_idx)==src_ext.extent(src_idx)) && match_expected_extents(++src_idx, ++sub_idx, src_ext, sub_ext, slices...);
   }