team-level std algos: part 3 (kokkos#6207)

Adding team level std algorithms for find, find_if, find_if_not, all_of, any_of, none_of, search_n
masterleinad · Aug 29, 2023 · 578dafa · 578dafa
1 parent 43d3e53
commit 578dafa
Show file tree

Hide file tree

Showing 18 changed files with 1,996 additions and 130 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_AllOf.hpp b/algorithms/src/std_algorithms/Kokkos_AllOf.hpp
@@ -23,41 +23,79 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class Predicate>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool all_of(const ExecutionSpace& ex, InputIterator first, InputIterator last,
             Predicate predicate) {
-  return Impl::all_of_impl("Kokkos::all_of_iterator_api_default", ex, first,
-                           last, predicate);
+  return Impl::all_of_exespace_impl("Kokkos::all_of_iterator_api_default", ex,
+                                    first, last, predicate);
 }
 
-template <class ExecutionSpace, class InputIterator, class Predicate>
+template <
+    typename ExecutionSpace, typename InputIterator, typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool all_of(const std::string& label, const ExecutionSpace& ex,
             InputIterator first, InputIterator last, Predicate predicate) {
-  return Impl::all_of_impl(label, ex, first, last, predicate);
+  return Impl::all_of_exespace_impl(label, ex, first, last, predicate);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool all_of(const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& v,
             Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::all_of_impl("Kokkos::all_of_view_api_default", ex, KE::cbegin(v),
-                           KE::cend(v), std::move(predicate));
+  return Impl::all_of_exespace_impl("Kokkos::all_of_view_api_default", ex,
+                                    KE::cbegin(v), KE::cend(v),
+                                    std::move(predicate));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool all_of(const std::string& label, const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& v,
             Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::all_of_impl(label, ex, KE::cbegin(v), KE::cend(v),
-                           std::move(predicate));
+  return Impl::all_of_exespace_impl(label, ex, KE::cbegin(v), KE::cend(v),
+                                    std::move(predicate));
+}
+
+//
+// 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.
+//
+template <typename TeamHandleType, typename InputIterator, typename Predicate,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION bool all_of(const TeamHandleType& teamHandle,
+                            InputIterator first, InputIterator last,
+                            Predicate predicate) {
+  return Impl::all_of_team_impl(teamHandle, first, last, predicate);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          typename Predicate,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION bool all_of(const TeamHandleType& teamHandle,
+                            const ::Kokkos::View<DataType, Properties...>& v,
+                            Predicate predicate) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::all_of_team_impl(teamHandle, KE::cbegin(v), KE::cend(v),
+                                std::move(predicate));
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_AnyOf.hpp b/algorithms/src/std_algorithms/Kokkos_AnyOf.hpp
@@ -23,41 +23,79 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class Predicate>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool any_of(const ExecutionSpace& ex, InputIterator first, InputIterator last,
             Predicate predicate) {
-  return Impl::any_of_impl("Kokkos::any_of_view_api_default", ex, first, last,
-                           predicate);
+  return Impl::any_of_exespace_impl("Kokkos::any_of_view_api_default", ex,
+                                    first, last, predicate);
 }
 
-template <class ExecutionSpace, class InputIterator, class Predicate>
+template <
+    typename ExecutionSpace, typename InputIterator, typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool any_of(const std::string& label, const ExecutionSpace& ex,
             InputIterator first, InputIterator last, Predicate predicate) {
-  return Impl::any_of_impl(label, ex, first, last, predicate);
+  return Impl::any_of_exespace_impl(label, ex, first, last, predicate);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool any_of(const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& v,
             Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::any_of_impl("Kokkos::any_of_view_api_default", ex, KE::cbegin(v),
-                           KE::cend(v), std::move(predicate));
+  return Impl::any_of_exespace_impl("Kokkos::any_of_view_api_default", ex,
+                                    KE::cbegin(v), KE::cend(v),
+                                    std::move(predicate));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename Predicate,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 bool any_of(const std::string& label, const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& v,
             Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::any_of_impl(label, ex, KE::cbegin(v), KE::cend(v),
-                           std::move(predicate));
+  return Impl::any_of_exespace_impl(label, ex, KE::cbegin(v), KE::cend(v),
+                                    std::move(predicate));
+}
+
+//
+// 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.
+//
+template <typename TeamHandleType, typename InputIterator, typename Predicate,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION bool any_of(const TeamHandleType& teamHandle,
+                            InputIterator first, InputIterator last,
+                            Predicate predicate) {
+  return Impl::any_of_team_impl(teamHandle, first, last, predicate);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          typename Predicate,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION bool any_of(const TeamHandleType& teamHandle,
+                            const ::Kokkos::View<DataType, Properties...>& v,
+                            Predicate predicate) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::any_of_team_impl(teamHandle, KE::cbegin(v), KE::cend(v),
+                                std::move(predicate));
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_Find.hpp b/algorithms/src/std_algorithms/Kokkos_Find.hpp
@@ -23,36 +23,76 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class T>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 InputIterator find(const ExecutionSpace& ex, InputIterator first,
                    InputIterator last, const T& value) {
-  return Impl::find_impl("Kokkos::find_iterator_api_default", ex, first, last,
-                         value);
+  return Impl::find_exespace_impl("Kokkos::find_iterator_api_default", ex,
+                                  first, last, value);
 }
 
-template <class ExecutionSpace, class InputIterator, class T>
+template <
+    typename ExecutionSpace, typename InputIterator, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 InputIterator find(const std::string& label, const ExecutionSpace& ex,
                    InputIterator first, InputIterator last, const T& value) {
-  return Impl::find_impl(label, ex, first, last, value);
+  return Impl::find_exespace_impl(label, ex, first, last, value);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties, class T>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find(const ExecutionSpace& ex,
           const ::Kokkos::View<DataType, Properties...>& view, const T& value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_impl("Kokkos::find_view_api_default", ex, KE::begin(view),
-                         KE::end(view), value);
+  return Impl::find_exespace_impl("Kokkos::find_view_api_default", ex,
+                                  KE::begin(view), KE::end(view), value);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties, class T>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto find(const std::string& label, const ExecutionSpace& ex,
           const ::Kokkos::View<DataType, Properties...>& view, const T& value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
 
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_impl(label, ex, KE::begin(view), KE::end(view), value);
+  return Impl::find_exespace_impl(label, ex, KE::begin(view), KE::end(view),
+                                  value);
+}
+
+//
+// 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.
+//
+template <typename TeamHandleType, typename InputIterator, typename T,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION InputIterator find(const TeamHandleType& teamHandle,
+                                   InputIterator first, InputIterator last,
+                                   const T& value) {
+  return Impl::find_team_impl(teamHandle, first, last, value);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          typename T,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto find(const TeamHandleType& teamHandle,
+                          const ::Kokkos::View<DataType, Properties...>& view,
+                          const T& value) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::find_team_impl(teamHandle, KE::begin(view), KE::end(view),
+                              value);
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_FindIf.hpp b/algorithms/src/std_algorithms/Kokkos_FindIf.hpp
@@ -23,42 +23,82 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class IteratorType, class PredicateType>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename IteratorType, typename PredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType find_if(const ExecutionSpace& ex, IteratorType first,
                      IteratorType last, PredicateType predicate) {
-  return Impl::find_if_or_not_impl<true>("Kokkos::find_if_iterator_api_default",
-                                         ex, first, last, std::move(predicate));
+  return Impl::find_if_or_not_exespace_impl<true>(
+      "Kokkos::find_if_iterator_api_default", ex, first, last,
+      std::move(predicate));
 }
 
-template <class ExecutionSpace, class IteratorType, class PredicateType>
+template <
+    typename ExecutionSpace, typename IteratorType, typename PredicateType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType find_if(const std::string& label, const ExecutionSpace& ex,
                      IteratorType first, IteratorType last,
                      PredicateType predicate) {
-  return Impl::find_if_or_not_impl<true>(label, ex, first, last,
-                                         std::move(predicate));
+  return Impl::find_if_or_not_exespace_impl<true>(label, ex, first, last,
+                                                  std::move(predicate));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <typename ExecutionSpace, typename DataType, typename... Properties,
+          typename Predicate,
+          std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
+                           int> = 0>
 auto find_if(const ExecutionSpace& ex,
              const ::Kokkos::View<DataType, Properties...>& v,
              Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_if_or_not_impl<true>("Kokkos::find_if_view_api_default", ex,
-                                         KE::begin(v), KE::end(v),
-                                         std::move(predicate));
+  return Impl::find_if_or_not_exespace_impl<true>(
+      "Kokkos::find_if_view_api_default", ex, KE::begin(v), KE::end(v),
+      std::move(predicate));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class Predicate>
+template <typename ExecutionSpace, typename DataType, typename... Properties,
+          typename Predicate,
+          std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
+                           int> = 0>
 auto find_if(const std::string& label, const ExecutionSpace& ex,
              const ::Kokkos::View<DataType, Properties...>& v,
              Predicate predicate) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::find_if_or_not_impl<true>(label, ex, KE::begin(v), KE::end(v),
-                                         std::move(predicate));
+  return Impl::find_if_or_not_exespace_impl<true>(
+      label, ex, KE::begin(v), KE::end(v), std::move(predicate));
+}
+
+//
+// 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.
+//
+template <typename TeamHandleType, typename IteratorType,
+          typename PredicateType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION IteratorType find_if(const TeamHandleType& teamHandle,
+                                     IteratorType first, IteratorType last,
+                                     PredicateType predicate) {
+  return Impl::find_if_or_not_team_impl<true>(teamHandle, first, last,
+                                              std::move(predicate));
+}
+
+template <
+    typename TeamHandleType, typename DataType, typename... Properties,
+    typename Predicate,
+    std::enable_if_t<::Kokkos::is_team_handle<TeamHandleType>::value, int> = 0>
+KOKKOS_FUNCTION auto find_if(const TeamHandleType& teamHandle,
+                             const ::Kokkos::View<DataType, Properties...>& v,
+                             Predicate predicate) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::find_if_or_not_team_impl<true>(teamHandle, KE::begin(v),
+                                              KE::end(v), std::move(predicate));
 }
 
 }  // namespace Experimental