team-level std algos: part 5 (kokkos#6209)

Add team level std algorithms for fill and replace Co-authored-by: Cezary Skrzyński Co-authored-by: Jakub Strzebonski Co-authored-by: antoinemeyer5 <antoine.meyer54@gmail.com>
nliber · Aug 31, 2023 · 393abd3 · 393abd3
1 parent 8aae5ea
commit 393abd3
Show file tree

Hide file tree

Showing 18 changed files with 1,403 additions and 146 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_Fill.hpp b/algorithms/src/std_algorithms/Kokkos_Fill.hpp
@@ -23,33 +23,67 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class IteratorType, class T>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename IteratorType, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void fill(const ExecutionSpace& ex, IteratorType first, IteratorType last,
           const T& value) {
-  Impl::fill_impl("Kokkos::fill_iterator_api_default", ex, first, last, value);
+  Impl::fill_exespace_impl("Kokkos::fill_iterator_api_default", ex, first, last,
+                           value);
 }
 
-template <class ExecutionSpace, class IteratorType, class T>
+template <
+    typename ExecutionSpace, typename IteratorType, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void fill(const std::string& label, const ExecutionSpace& ex,
           IteratorType first, IteratorType last, const T& value) {
-  Impl::fill_impl(label, ex, first, last, value);
+  Impl::fill_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>
 void fill(const ExecutionSpace& ex,
           const ::Kokkos::View<DataType, Properties...>& view, const T& value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-
-  Impl::fill_impl("Kokkos::fill_view_api_default", ex, begin(view), end(view),
-                  value);
+  Impl::fill_exespace_impl("Kokkos::fill_view_api_default", ex, begin(view),
+                           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>
 void fill(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);
+  Impl::fill_exespace_impl(label, ex, begin(view), end(view), value);
+}
 
-  Impl::fill_impl(label, ex, begin(view), 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 IteratorType, typename T,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION void fill(const TeamHandleType& th, IteratorType first,
+                          IteratorType last, const T& value) {
+  Impl::fill_team_impl(th, 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 void fill(const TeamHandleType& th,
+                          const ::Kokkos::View<DataType, Properties...>& view,
+                          const T& value) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  Impl::fill_team_impl(th, begin(view), end(view), value);
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_FillN.hpp b/algorithms/src/std_algorithms/Kokkos_FillN.hpp
@@ -23,38 +23,72 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class IteratorType, class SizeType, class T>
+template <
+    typename ExecutionSpace, typename IteratorType, typename SizeType,
+    typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType fill_n(const ExecutionSpace& ex, IteratorType first, SizeType n,
                     const T& value) {
-  return Impl::fill_n_impl("Kokkos::fill_n_iterator_api_default", ex, first, n,
-                           value);
+  return Impl::fill_n_exespace_impl("Kokkos::fill_n_iterator_api_default", ex,
+                                    first, n, value);
 }
 
-template <class ExecutionSpace, class IteratorType, class SizeType, class T>
+template <
+    typename ExecutionSpace, typename IteratorType, typename SizeType,
+    typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType fill_n(const std::string& label, const ExecutionSpace& ex,
                     IteratorType first, SizeType n, const T& value) {
-  return Impl::fill_n_impl(label, ex, first, n, value);
+  return Impl::fill_n_exespace_impl(label, ex, first, n, value);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class SizeType, class T>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename SizeType, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto fill_n(const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& view, SizeType n,
             const T& value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
 
-  return Impl::fill_n_impl("Kokkos::fill_n_view_api_default", ex, begin(view),
-                           n, value);
+  return Impl::fill_n_exespace_impl("Kokkos::fill_n_view_api_default", ex,
+                                    begin(view), n, value);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties,
-          class SizeType, class T>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    typename SizeType, typename T,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto fill_n(const std::string& label, const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& view, SizeType n,
             const T& value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
 
-  return Impl::fill_n_impl(label, ex, begin(view), n, value);
+  return Impl::fill_n_exespace_impl(label, ex, begin(view), n, 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 IteratorType, typename SizeType,
+          typename T,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION IteratorType fill_n(const TeamHandleType& th,
+                                    IteratorType first, SizeType n,
+                                    const T& value) {
+  return Impl::fill_n_team_impl(th, first, n, value);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          typename SizeType, typename T,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto fill_n(const TeamHandleType& th,
+                            const ::Kokkos::View<DataType, Properties...>& view,
+                            SizeType n, const T& value) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  return Impl::fill_n_team_impl(th, begin(view), n, value);
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_Replace.hpp b/algorithms/src/std_algorithms/Kokkos_Replace.hpp
@@ -23,40 +23,77 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class Iterator, class ValueType>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename Iterator, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void replace(const ExecutionSpace& ex, Iterator first, Iterator last,
              const ValueType& old_value, const ValueType& new_value) {
-  return Impl::replace_impl("Kokkos::replace_iterator_api", ex, first, last,
-                            old_value, new_value);
+  Impl::replace_exespace_impl("Kokkos::replace_iterator_api", ex, first, last,
+                              old_value, new_value);
 }
 
-template <class ExecutionSpace, class Iterator, class ValueType>
+template <
+    typename ExecutionSpace, typename Iterator, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void replace(const std::string& label, const ExecutionSpace& ex, Iterator first,
              Iterator last, const ValueType& old_value,
              const ValueType& new_value) {
-  return Impl::replace_impl(label, ex, first, last, old_value, new_value);
+  Impl::replace_exespace_impl(label, ex, first, last, old_value, new_value);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void replace(const ExecutionSpace& ex,
              const ::Kokkos::View<DataType1, Properties1...>& view,
              const ValueType& old_value, const ValueType& new_value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::replace_impl("Kokkos::replace_view_api", ex, KE::begin(view),
-                            KE::end(view), old_value, new_value);
+  Impl::replace_exespace_impl("Kokkos::replace_view_api", ex, KE::begin(view),
+                              KE::end(view), old_value, new_value);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void replace(const std::string& label, const ExecutionSpace& ex,
              const ::Kokkos::View<DataType1, Properties1...>& view,
              const ValueType& old_value, const ValueType& new_value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::replace_impl(label, ex, KE::begin(view), KE::end(view),
-                            old_value, new_value);
+  Impl::replace_exespace_impl(label, ex, KE::begin(view), KE::end(view),
+                              old_value, new_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 Iterator, typename ValueType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION void replace(const TeamHandleType& teamHandle, Iterator first,
+                             Iterator last, const ValueType& old_value,
+                             const ValueType& new_value) {
+  Impl::replace_team_impl(teamHandle, first, last, old_value, new_value);
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename ValueType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION void replace(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view,
+    const ValueType& old_value, const ValueType& new_value) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  namespace KE = ::Kokkos::Experimental;
+  Impl::replace_team_impl(teamHandle, KE::begin(view), KE::end(view), old_value,
+                          new_value);
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_ReplaceCopy.hpp b/algorithms/src/std_algorithms/Kokkos_ReplaceCopy.hpp
@@ -23,54 +23,99 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator,
-          class ValueType>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator replace_copy(const ExecutionSpace& ex, InputIterator first_from,
                             InputIterator last_from, OutputIterator first_dest,
                             const ValueType& old_value,
                             const ValueType& new_value) {
-  return Impl::replace_copy_impl("Kokkos::replace_copy_iterator_api", ex,
-                                 first_from, last_from, first_dest, old_value,
-                                 new_value);
+  return Impl::replace_copy_exespace_impl("Kokkos::replace_copy_iterator_api",
+                                          ex, first_from, last_from, first_dest,
+                                          old_value, new_value);
 }
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator,
-          class ValueType>
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator replace_copy(const std::string& label, const ExecutionSpace& ex,
                             InputIterator first_from, InputIterator last_from,
                             OutputIterator first_dest,
                             const ValueType& old_value,
                             const ValueType& new_value) {
-  return Impl::replace_copy_impl(label, ex, first_from, last_from, first_dest,
-                                 old_value, new_value);
+  return Impl::replace_copy_exespace_impl(label, ex, first_from, last_from,
+                                          first_dest, old_value, new_value);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto replace_copy(const ExecutionSpace& ex,
                   const ::Kokkos::View<DataType1, Properties1...>& view_from,
                   const ::Kokkos::View<DataType2, Properties2...>& view_dest,
                   const ValueType& old_value, const ValueType& new_value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::replace_copy_impl("Kokkos::replace_copy_view_api", ex,
-                                 KE::cbegin(view_from), KE::cend(view_from),
-                                 KE::begin(view_dest), old_value, new_value);
+  return Impl::replace_copy_exespace_impl(
+      "Kokkos::replace_copy_view_api", ex, KE::cbegin(view_from),
+      KE::cend(view_from), KE::begin(view_dest), old_value, new_value);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto replace_copy(const std::string& label, const ExecutionSpace& ex,
                   const ::Kokkos::View<DataType1, Properties1...>& view_from,
                   const ::Kokkos::View<DataType2, Properties2...>& view_dest,
                   const ValueType& old_value, const ValueType& new_value) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::replace_copy_impl(label, ex, KE::cbegin(view_from),
-                                 KE::cend(view_from), KE::begin(view_dest),
-                                 old_value, new_value);
+  return Impl::replace_copy_exespace_impl(
+      label, ex, KE::cbegin(view_from), KE::cend(view_from),
+      KE::begin(view_dest), old_value, new_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 OutputIterator, typename ValueType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION OutputIterator replace_copy(const TeamHandleType& teamHandle,
+                                            InputIterator first_from,
+                                            InputIterator last_from,
+                                            OutputIterator first_dest,
+                                            const ValueType& old_value,
+                                            const ValueType& new_value) {
+  return Impl::replace_copy_team_impl(teamHandle, first_from, last_from,
+                                      first_dest, old_value, new_value);
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename DataType2, typename... Properties2, typename ValueType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto replace_copy(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view_from,
+    const ::Kokkos::View<DataType2, Properties2...>& view_dest,
+    const ValueType& old_value, const ValueType& new_value) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::replace_copy_team_impl(teamHandle, KE::cbegin(view_from),
+                                      KE::cend(view_from), KE::begin(view_dest),
+                                      old_value, new_value);
 }
 
 }  // namespace Experimental