team-level std algos: part 2 (kokkos#6205)

* add part2 Co-authored-by: Cezary Skrzyński Co-authored-by: Jakub Strzebonski * add tests * update API based on comments * address review * format again * fix lambda capture * remove maybe unused * remove maybe_unused * remove maybe_unused * formatting
nliber · Aug 24, 2023 · 15d5217 · 15d5217
1 parent 7e45f17
commit 15d5217
Show file tree

Hide file tree

Showing 13 changed files with 1,992 additions and 255 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_Equal.hpp b/algorithms/src/std_algorithms/Kokkos_Equal.hpp
diff --git a/algorithms/src/std_algorithms/Kokkos_FindEnd.hpp b/algorithms/src/std_algorithms/Kokkos_FindEnd.hpp
@@ -24,70 +24,89 @@
 namespace Kokkos {
 namespace Experimental {
 
+//
+// overload set accepting execution space
+//
+
 // overload set 1: no binary predicate passed
-template <class ExecutionSpace, class IteratorType1, class IteratorType2>
+template <
+    typename ExecutionSpace, typename IteratorType1, typename IteratorType2,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType1 find_end(const ExecutionSpace& ex, IteratorType1 first,
                        IteratorType1 last, IteratorType2 s_first,
                        IteratorType2 s_last) {
-  return Impl::find_end_impl("Kokkos::find_end_iterator_api_default", ex, first,
-                             last, s_first, s_last);
+  return Impl::find_end_exespace_impl("Kokkos::find_end_iterator_api_default",
+                                      ex, first, last, s_first, s_last);
 }
 
-template <class ExecutionSpace, class IteratorType1, class IteratorType2>
+template <
+    typename ExecutionSpace, typename IteratorType1, typename IteratorType2,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType1 find_end(const std::string& label, const ExecutionSpace& ex,
                        IteratorType1 first, IteratorType1 last,
                        IteratorType2 s_first, IteratorType2 s_last) {
-  return Impl::find_end_impl(label, ex, first, last, s_first, s_last);
+  return Impl::find_end_exespace_impl(label, ex, first, last, s_first, s_last);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find_end(const ExecutionSpace& ex,
               const ::Kokkos::View<DataType1, Properties1...>& view,
               const ::Kokkos::View<DataType2, Properties2...>& s_view) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_end_impl("Kokkos::find_end_view_api_default", ex,
-                             KE::begin(view), KE::end(view), KE::begin(s_view),
-                             KE::end(s_view));
+  return Impl::find_end_exespace_impl("Kokkos::find_end_view_api_default", ex,
+                                      KE::begin(view), KE::end(view),
+                                      KE::begin(s_view), KE::end(s_view));
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find_end(const std::string& label, const ExecutionSpace& ex,
               const ::Kokkos::View<DataType1, Properties1...>& view,
               const ::Kokkos::View<DataType2, Properties2...>& s_view) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_end_impl(label, ex, KE::begin(view), KE::end(view),
-                             KE::begin(s_view), KE::end(s_view));
+  return Impl::find_end_exespace_impl(label, ex, KE::begin(view), KE::end(view),
+                                      KE::begin(s_view), KE::end(s_view));
 }
 
 // overload set 2: binary predicate passed
-template <class ExecutionSpace, class IteratorType1, class IteratorType2,
-          class BinaryPredicateType>
+template <
+    typename ExecutionSpace, typename IteratorType1, typename IteratorType2,
+    typename BinaryPredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType1 find_end(const ExecutionSpace& ex, IteratorType1 first,
                        IteratorType1 last, IteratorType2 s_first,
                        IteratorType2 s_last, const BinaryPredicateType& pred) {
-  return Impl::find_end_impl("Kokkos::find_end_iterator_api_default", ex, first,
-                             last, s_first, s_last, pred);
+  return Impl::find_end_exespace_impl("Kokkos::find_end_iterator_api_default",
+                                      ex, first, last, s_first, s_last, pred);
 }
 
-template <class ExecutionSpace, class IteratorType1, class IteratorType2,
-          class BinaryPredicateType>
+template <
+    typename ExecutionSpace, typename IteratorType1, typename IteratorType2,
+    typename BinaryPredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType1 find_end(const std::string& label, const ExecutionSpace& ex,
                        IteratorType1 first, IteratorType1 last,
                        IteratorType2 s_first, IteratorType2 s_last,
                        const BinaryPredicateType& pred) {
-  return Impl::find_end_impl(label, ex, first, last, s_first, s_last, pred);
+  return Impl::find_end_exespace_impl(label, ex, first, last, s_first, s_last,
+                                      pred);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryPredicateType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryPredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find_end(const ExecutionSpace& ex,
               const ::Kokkos::View<DataType1, Properties1...>& view,
               const ::Kokkos::View<DataType2, Properties2...>& s_view,
@@ -96,13 +115,15 @@ auto find_end(const ExecutionSpace& ex,
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_end_impl("Kokkos::find_end_view_api_default", ex,
-                             KE::begin(view), KE::end(view), KE::begin(s_view),
-                             KE::end(s_view), pred);
+  return Impl::find_end_exespace_impl("Kokkos::find_end_view_api_default", ex,
+                                      KE::begin(view), KE::end(view),
+                                      KE::begin(s_view), KE::end(s_view), pred);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryPredicateType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryPredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find_end(const std::string& label, const ExecutionSpace& ex,
               const ::Kokkos::View<DataType1, Properties1...>& view,
               const ::Kokkos::View<DataType2, Properties2...>& s_view,
@@ -111,8 +132,71 @@ auto find_end(const std::string& label, const ExecutionSpace& ex,
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_end_impl(label, ex, KE::begin(view), KE::end(view),
-                             KE::begin(s_view), KE::end(s_view), pred);
+  return Impl::find_end_exespace_impl(label, ex, KE::begin(view), KE::end(view),
+                                      KE::begin(s_view), KE::end(s_view), pred);
+}
+
+//
+// overload set accepting a team handle
+// Note: for now omit the overloads accepting a label
+// since they cause issues on device because of the string allocation.
+//
+
+// overload set 1: no binary predicate passed
+template <typename TeamHandleType, typename IteratorType1,
+          typename IteratorType2,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION IteratorType1 find_end(const TeamHandleType& teamHandle,
+                                       IteratorType1 first, IteratorType1 last,
+                                       IteratorType2 s_first,
+                                       IteratorType2 s_last) {
+  return Impl::find_end_team_impl(teamHandle, first, last, s_first, s_last);
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename DataType2, typename... Properties2,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto find_end(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view,
+    const ::Kokkos::View<DataType2, Properties2...>& s_view) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::find_end_team_impl(teamHandle, KE::begin(view), KE::end(view),
+                                  KE::begin(s_view), KE::end(s_view));
+}
+
+// overload set 2: binary predicate passed
+template <typename TeamHandleType, typename IteratorType1,
+          typename IteratorType2, typename BinaryPredicateType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+
+KOKKOS_FUNCTION IteratorType1 find_end(const TeamHandleType& teamHandle,
+                                       IteratorType1 first, IteratorType1 last,
+                                       IteratorType2 s_first,
+                                       IteratorType2 s_last,
+                                       const BinaryPredicateType& pred) {
+  return Impl::find_end_team_impl(teamHandle, first, last, s_first, s_last,
+                                  pred);
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename DataType2, typename... Properties2,
+          typename BinaryPredicateType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto find_end(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view,
+    const ::Kokkos::View<DataType2, Properties2...>& s_view,
+    const BinaryPredicateType& pred) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(s_view);
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::find_end_team_impl(teamHandle, KE::begin(view), KE::end(view),
+                                  KE::begin(s_view), KE::end(s_view), pred);
 }
 
 }  // namespace Experimental