team-level std algos: part 6 (kokkos#6210)

Team-level implementation and tests for reverse and rotate

Co-authored-by: Cezary Skrzyński
Co-authored-by: Jakub Strzebonski
Co-authored-by: antoinemeyer5 <antoine.meyer54@gmail.com>

algorithms/src/std_algorithms/Kokkos_Reverse.hpp

@@ -23,34 +23,67 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void reverse(const ExecutionSpace& ex, InputIterator first,
              InputIterator last) {
-  return Impl::reverse_impl("Kokkos::reverse_iterator_api_default", ex, first,
-                            last);
+  return Impl::reverse_exespace_impl("Kokkos::reverse_iterator_api_default", ex,
+                                     first, last);
 }
 
-template <class ExecutionSpace, class InputIterator>
+template <
+    typename ExecutionSpace, typename InputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void reverse(const std::string& label, const ExecutionSpace& ex,
              InputIterator first, InputIterator last) {
-  return Impl::reverse_impl(label, ex, first, last);
+  return Impl::reverse_exespace_impl(label, ex, first, last);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void reverse(const ExecutionSpace& ex,
              const ::Kokkos::View<DataType, Properties...>& view) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::reverse_impl("Kokkos::reverse_view_api_default", ex,
-                            KE::begin(view), KE::end(view));
+  return Impl::reverse_exespace_impl("Kokkos::reverse_view_api_default", ex,
+                                     KE::begin(view), KE::end(view));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 void reverse(const std::string& label, const ExecutionSpace& ex,
              const ::Kokkos::View<DataType, Properties...>& view) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
   namespace KE = ::Kokkos::Experimental;
-  return Impl::reverse_impl(label, ex, KE::begin(view), KE::end(view));
+  return Impl::reverse_exespace_impl(label, ex, KE::begin(view), KE::end(view));
+}
+
+//
+// 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,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION void reverse(const TeamHandleType& teamHandle,
+                             InputIterator first, InputIterator last) {
+  return Impl::reverse_team_impl(teamHandle, first, last);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION void reverse(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& view) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::reverse_team_impl(teamHandle, KE::begin(view), KE::end(view));
 }
 
 }  // namespace Experimental

algorithms/src/std_algorithms/Kokkos_ReverseCopy.hpp

@@ -23,42 +23,83 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator reverse_copy(const ExecutionSpace& ex, InputIterator first,
                             InputIterator last, OutputIterator d_first) {
-  return Impl::reverse_copy_impl("Kokkos::reverse_copy_iterator_api_default",
-                                 ex, first, last, d_first);
+  return Impl::reverse_copy_exespace_impl(
+      "Kokkos::reverse_copy_iterator_api_default", ex, first, last, d_first);
 }
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator>
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator reverse_copy(const std::string& label, const ExecutionSpace& ex,
                             InputIterator first, InputIterator last,
                             OutputIterator d_first) {
-  return Impl::reverse_copy_impl(label, ex, first, last, d_first);
+  return Impl::reverse_copy_exespace_impl(label, ex, first, last, d_first);
 }
 
-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 reverse_copy(const ExecutionSpace& ex,
                   const ::Kokkos::View<DataType1, Properties1...>& source,
                   ::Kokkos::View<DataType2, Properties2...>& dest) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
 
-  return Impl::reverse_copy_impl("Kokkos::reverse_copy_view_api_default", ex,
-                                 cbegin(source), cend(source), begin(dest));
+  return Impl::reverse_copy_exespace_impl(
+      "Kokkos::reverse_copy_view_api_default", ex, cbegin(source), cend(source),
+      begin(dest));
 }
 
-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 reverse_copy(const std::string& label, const ExecutionSpace& ex,
                   const ::Kokkos::View<DataType1, Properties1...>& source,
                   ::Kokkos::View<DataType2, Properties2...>& dest) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
 
-  return Impl::reverse_copy_impl(label, ex, cbegin(source), cend(source),
-                                 begin(dest));
+  return Impl::reverse_copy_exespace_impl(label, ex, cbegin(source),
+                                          cend(source), begin(dest));
+}
+
+//
+// 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,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION OutputIterator reverse_copy(const TeamHandleType& teamHandle,
+                                            InputIterator first,
+                                            InputIterator last,
+                                            OutputIterator d_first) {
+  return Impl::reverse_copy_team_impl(teamHandle, first, last, d_first);
+}
+
+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 reverse_copy(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& source,
+    ::Kokkos::View<DataType2, Properties2...>& dest) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
+
+  return Impl::reverse_copy_team_impl(teamHandle, cbegin(source), cend(source),
+                                      begin(dest));
 }
 
 }  // namespace Experimental

algorithms/src/std_algorithms/Kokkos_Rotate.hpp

@@ -23,36 +23,71 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class IteratorType>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType rotate(const ExecutionSpace& ex, IteratorType first,
                     IteratorType n_first, IteratorType last) {
-  return Impl::rotate_impl("Kokkos::rotate_iterator_api_default", ex, first,
-                           n_first, last);
+  return Impl::rotate_exespace_impl("Kokkos::rotate_iterator_api_default", ex,
+                                    first, n_first, last);
 }
 
-template <class ExecutionSpace, class IteratorType>
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType rotate(const std::string& label, const ExecutionSpace& ex,
                     IteratorType first, IteratorType n_first,
                     IteratorType last) {
-  return Impl::rotate_impl(label, ex, first, n_first, last);
+  return Impl::rotate_exespace_impl(label, ex, first, n_first, last);
 }
 
-template <class ExecutionSpace, class DataType, class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto rotate(const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& view,
             std::size_t n_location) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::rotate_impl("Kokkos::rotate_view_api_default", ex, begin(view),
-                           begin(view) + n_location, end(view));
+  return Impl::rotate_exespace_impl("Kokkos::rotate_view_api_default", ex,
+                                    begin(view), begin(view) + n_location,
+                                    end(view));
 }
 
-template <class ExecutionSpace, class DataType, class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto rotate(const std::string& label, const ExecutionSpace& ex,
             const ::Kokkos::View<DataType, Properties...>& view,
             std::size_t n_location) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::rotate_impl(label, ex, begin(view), begin(view) + n_location,
-                           end(view));
+  return Impl::rotate_exespace_impl(label, ex, begin(view),
+                                    begin(view) + n_location, end(view));
+}
+
+//
+// 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,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION IteratorType rotate(const TeamHandleType& teamHandle,
+                                    IteratorType first, IteratorType n_first,
+                                    IteratorType last) {
+  return Impl::rotate_team_impl(teamHandle, first, n_first, last);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto rotate(const TeamHandleType& teamHandle,
+                            const ::Kokkos::View<DataType, Properties...>& view,
+                            std::size_t n_location) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  return Impl::rotate_team_impl(teamHandle, begin(view),
+                                begin(view) + n_location, end(view));
 }
 
 }  // namespace Experimental

algorithms/src/std_algorithms/Kokkos_RotateCopy.hpp

@@ -23,47 +23,92 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator rotate_copy(const ExecutionSpace& ex, InputIterator first,
                            InputIterator n_first, InputIterator last,
                            OutputIterator d_first) {
-  return Impl::rotate_copy_impl("Kokkos::rotate_copy_iterator_api_default", ex,
-                                first, n_first, last, d_first);
+  return Impl::rotate_copy_exespace_impl(
+      "Kokkos::rotate_copy_iterator_api_default", ex, first, n_first, last,
+      d_first);
 }
 
-template <class ExecutionSpace, class InputIterator, class OutputIterator>
+template <
+    typename ExecutionSpace, typename InputIterator, typename OutputIterator,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 OutputIterator rotate_copy(const std::string& label, const ExecutionSpace& ex,
                            InputIterator first, InputIterator n_first,
                            InputIterator last, OutputIterator d_first) {
-  return Impl::rotate_copy_impl(label, ex, first, n_first, last, d_first);
+  return Impl::rotate_copy_exespace_impl(label, ex, first, n_first, last,
+                                         d_first);
 }
 
-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 rotate_copy(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType1, Properties1...>& source,
                  std::size_t n_location,
                  const ::Kokkos::View<DataType2, Properties2...>& dest) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
 
-  return Impl::rotate_copy_impl("Kokkos::rotate_copy_view_api_default", ex,
-                                cbegin(source), cbegin(source) + n_location,
-                                cend(source), begin(dest));
+  return Impl::rotate_copy_exespace_impl(
+      "Kokkos::rotate_copy_view_api_default", ex, cbegin(source),
+      cbegin(source) + n_location, cend(source), begin(dest));
 }
 
-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 rotate_copy(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType1, Properties1...>& source,
                  std::size_t n_location,
                  const ::Kokkos::View<DataType2, Properties2...>& dest) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
 
-  return Impl::rotate_copy_impl(label, ex, cbegin(source),
-                                cbegin(source) + n_location, cend(source),
-                                begin(dest));
+  return Impl::rotate_copy_exespace_impl(label, ex, cbegin(source),
+                                         cbegin(source) + n_location,
+                                         cend(source), begin(dest));
+}
+
+//
+// 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,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION OutputIterator rotate_copy(const TeamHandleType& teamHandle,
+                                           InputIterator first,
+                                           InputIterator n_first,
+                                           InputIterator last,
+                                           OutputIterator d_first) {
+  return Impl::rotate_copy_team_impl(teamHandle, first, n_first, last, d_first);
+}
+
+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 rotate_copy(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& source,
+    std::size_t n_location,
+    const ::Kokkos::View<DataType2, Properties2...>& dest) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(source);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(dest);
+
+  return Impl::rotate_copy_team_impl(teamHandle, cbegin(source),
+                                     cbegin(source) + n_location, cend(source),
+                                     begin(dest));
 }
 
 }  // namespace Experimental