Don't use get for std::complex slice specifier

and test that complex<double> works like pair as slice specifier *barf*
kokkos · Sep 4, 2024 · 4f551ba · 4f551ba
1 parent 192ffb2
commit 4f551ba
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/include/experimental/__p2630_bits/submdspan_extents.hpp b/include/experimental/__p2630_bits/submdspan_extents.hpp
@@ -105,6 +105,12 @@ constexpr auto first_of(const Slice &i) {
   return std::get<0>(i);
 }
 
+template<class T>
+MDSPAN_INLINE_FUNCTION
+constexpr auto first_of(const std::complex<T> &i) {
+  return i.real();
+}
+
 template <class OffsetType, class ExtentType, class StrideType>
 MDSPAN_INLINE_FUNCTION
 constexpr OffsetType
@@ -136,6 +142,12 @@ constexpr auto last_of(std::integral_constant<size_t, k>, const Extents &,
   return std::get<1>(i);
 }
 
+template<size_t k, class Extents, class T>
+MDSPAN_INLINE_FUNCTION
+constexpr auto last_of(std::integral_constant<size_t, k>, const Extents &, const std::complex<T> &i) {
+  return i.imag();
+}
+
 // Suppress spurious warning with NVCC about no return statement.
 // This is a known issue in NVCC and NVC++
 // Depending on the CUDA and GCC version we need both the builtin

diff --git a/tests/test_submdspan.cpp b/tests/test_submdspan.cpp
@@ -140,6 +140,7 @@ using submdspan_test_types =
     // layout_right to layout_right Check Extents Preservation
     , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10>,           args_t<10>,          Kokkos::extents<size_t,10>, Kokkos::full_extent_t>
     , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10>,           args_t<10>,          Kokkos::extents<size_t,dyn>, std::pair<int,int>>
+    , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10>,           args_t<10>,          Kokkos::extents<size_t,dyn>, std::complex<double>>
     , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10>,           args_t<10>,          Kokkos::extents<size_t>, int>
     , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10,20>,        args_t<10,20>,       Kokkos::extents<size_t,10,20>, Kokkos::full_extent_t, Kokkos::full_extent_t>
     , std::tuple<Kokkos::layout_right, Kokkos::layout_right, Kokkos::extents<size_t,10,20>,        args_t<10,20>,       Kokkos::extents<size_t,dyn,20>, std::pair<int,int>, Kokkos::full_extent_t>
@@ -274,6 +275,10 @@ struct TestSubMDSpan<
     return std::pair<int,int>(1,3);
   }
   MDSPAN_INLINE_FUNCTION
+  static auto create_slice_arg(std::complex<double>) {
+    return std::complex<double>{1.,3.};
+  }
+  MDSPAN_INLINE_FUNCTION
   static auto create_slice_arg(Kokkos::strided_slice<int,int,int>) {
     return Kokkos::strided_slice<int,int,int>{1,3,2};
   }
@@ -300,6 +305,12 @@ struct TestSubMDSpan<
   }
   template<class SrcMDSpan, class SubMDSpan, size_t ... SrcIdx, size_t ... SubIdx, class ... SliceArgs>
   MDSPAN_INLINE_FUNCTION
+  static bool check_submdspan_match(int src_idx, int sub_idx, SrcMDSpan src_mds, SubMDSpan sub_mds, std::index_sequence<SrcIdx...>, std::index_sequence<SubIdx...>, std::complex<double> p, SliceArgs ... slices) {
+    using idx_t = typename SubMDSpan::index_type;
+    return (sub_mds.extent(sub_idx)==static_cast<idx_t>(p.imag()-p.real())) && check_submdspan_match(++src_idx, ++sub_idx, src_mds, sub_mds, std::index_sequence<SrcIdx...,2>(), std::index_sequence<SubIdx...,1>(), slices...);
+  }
+  template<class SrcMDSpan, class SubMDSpan, size_t ... SrcIdx, size_t ... SubIdx, class ... SliceArgs>
+  MDSPAN_INLINE_FUNCTION
   static bool check_submdspan_match(int src_idx, int sub_idx, SrcMDSpan src_mds, SubMDSpan sub_mds, std::index_sequence<SrcIdx...>, std::index_sequence<SubIdx...>,
                                      Kokkos::strided_slice<int,int,int> p, SliceArgs ... slices) {
     using idx_t = typename SubMDSpan::index_type;