Team-level std algos: part 7 (kokkos#6211)

Team level algorithms for move and shift Co-authored-by: Jakub Strzebonski Co-authored-by: antoinemeyer5 <antoine.meyer54@gmail.com> Co-authored-by: Cezary Skrzyński <cezary.skrzynski@ng-analytics.com>
nliber · Sep 7, 2023 · fc213ea · fc213ea
1 parent fd774d9
commit fc213ea
Show file tree

Hide file tree

Showing 14 changed files with 998 additions and 72 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_Move.hpp b/algorithms/src/std_algorithms/Kokkos_Move.hpp
@@ -23,41 +23,81 @@
 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 move(const ExecutionSpace& ex, InputIterator first,
                     InputIterator last, OutputIterator d_first) {
-  return Impl::move_impl("Kokkos::move_iterator_api_default", ex, first, last,
-                         d_first);
+  return Impl::move_exespace_impl("Kokkos::move_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 move(const std::string& label, const ExecutionSpace& ex,
                     InputIterator first, InputIterator last,
                     OutputIterator d_first) {
-  return Impl::move_impl(label, ex, first, last, d_first);
+  return Impl::move_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 move(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::move_impl("Kokkos::move_view_api_default", ex, begin(source),
-                         end(source), begin(dest));
+  return Impl::move_exespace_impl("Kokkos::move_view_api_default", ex,
+                                  begin(source), end(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 move(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::move_impl(label, ex, begin(source), end(source), begin(dest));
+  return Impl::move_exespace_impl(label, ex, begin(source), end(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 move(const TeamHandleType& teamHandle,
+                                    InputIterator first, InputIterator last,
+                                    OutputIterator d_first) {
+  return Impl::move_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 move(
+    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::move_team_impl(teamHandle, begin(source), end(source),
+                              begin(dest));
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_MoveBackward.hpp b/algorithms/src/std_algorithms/Kokkos_MoveBackward.hpp
@@ -23,42 +23,83 @@
 namespace Kokkos {
 namespace Experimental {
 
-template <class ExecutionSpace, class IteratorType1, class IteratorType2>
+//
+// overload set accepting execution space
+//
+template <
+    typename ExecutionSpace, typename IteratorType1, typename IteratorType2,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType2 move_backward(const ExecutionSpace& ex, IteratorType1 first,
                             IteratorType1 last, IteratorType2 d_last) {
-  return Impl::move_backward_impl("Kokkos::move_backward_iterator_api_default",
-                                  ex, first, last, d_last);
+  return Impl::move_backward_exespace_impl(
+      "Kokkos::move_backward_iterator_api_default", ex, first, last, d_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 move_backward(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::move_backward_impl("Kokkos::move_backward_view_api_default", ex,
-                                  begin(source), end(source), end(dest));
+  return Impl::move_backward_exespace_impl(
+      "Kokkos::move_backward_view_api_default", ex, begin(source), end(source),
+      end(dest));
 }
 
-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>
 IteratorType2 move_backward(const std::string& label, const ExecutionSpace& ex,
                             IteratorType1 first, IteratorType1 last,
                             IteratorType2 d_last) {
-  return Impl::move_backward_impl(label, ex, first, last, d_last);
+  return Impl::move_backward_exespace_impl(label, ex, first, last, d_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 move_backward(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::move_backward_impl(label, ex, begin(source), end(source),
-                                  end(dest));
+  return Impl::move_backward_exespace_impl(label, ex, begin(source),
+                                           end(source), end(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 IteratorType1,
+          typename IteratorType2,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION IteratorType2 move_backward(const TeamHandleType& teamHandle,
+                                            IteratorType1 first,
+                                            IteratorType1 last,
+                                            IteratorType2 d_last) {
+  return Impl::move_backward_team_impl(teamHandle, first, last, d_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 move_backward(
+    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::move_backward_team_impl(teamHandle, begin(source), end(source),
+                                       end(dest));
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_ShiftLeft.hpp b/algorithms/src/std_algorithms/Kokkos_ShiftLeft.hpp
@@ -23,36 +23,70 @@
 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 shift_left(const ExecutionSpace& ex, IteratorType first,
                         IteratorType last,
                         typename IteratorType::difference_type n) {
-  return Impl::shift_left_impl("Kokkos::shift_left_iterator_api_default", ex,
-                               first, last, n);
+  return Impl::shift_left_exespace_impl(
+      "Kokkos::shift_left_iterator_api_default", ex, first, last, n);
 }
 
-template <class ExecutionSpace, class IteratorType>
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType shift_left(const std::string& label, const ExecutionSpace& ex,
                         IteratorType first, IteratorType last,
                         typename IteratorType::difference_type n) {
-  return Impl::shift_left_impl(label, ex, first, last, n);
+  return Impl::shift_left_exespace_impl(label, ex, first, last, n);
 }
 
-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 shift_left(const ExecutionSpace& ex,
                 const ::Kokkos::View<DataType, Properties...>& view,
                 typename decltype(begin(view))::difference_type n) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::shift_left_impl("Kokkos::shift_left_view_api_default", ex,
-                               begin(view), end(view), n);
+  return Impl::shift_left_exespace_impl("Kokkos::shift_left_view_api_default",
+                                        ex, begin(view), end(view), n);
 }
 
-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 shift_left(const std::string& label, const ExecutionSpace& ex,
                 const ::Kokkos::View<DataType, Properties...>& view,
                 typename decltype(begin(view))::difference_type n) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::shift_left_impl(label, ex, begin(view), end(view), n);
+  return Impl::shift_left_exespace_impl(label, ex, begin(view), end(view), n);
+}
+
+//
+// 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
+shift_left(const TeamHandleType& teamHandle, IteratorType first,
+           IteratorType last, typename IteratorType::difference_type n) {
+  return Impl::shift_left_team_impl(teamHandle, first, last, n);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto shift_left(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& view,
+    typename decltype(begin(view))::difference_type n) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  return Impl::shift_left_team_impl(teamHandle, begin(view), end(view), n);
 }
 
 }  // namespace Experimental

diff --git a/algorithms/src/std_algorithms/Kokkos_ShiftRight.hpp b/algorithms/src/std_algorithms/Kokkos_ShiftRight.hpp
@@ -23,36 +23,70 @@
 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 shift_right(const ExecutionSpace& ex, IteratorType first,
                          IteratorType last,
                          typename IteratorType::difference_type n) {
-  return Impl::shift_right_impl("Kokkos::shift_right_iterator_api_default", ex,
-                                first, last, n);
+  return Impl::shift_right_exespace_impl(
+      "Kokkos::shift_right_iterator_api_default", ex, first, last, n);
 }
 
-template <class ExecutionSpace, class IteratorType>
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 IteratorType shift_right(const std::string& label, const ExecutionSpace& ex,
                          IteratorType first, IteratorType last,
                          typename IteratorType::difference_type n) {
-  return Impl::shift_right_impl(label, ex, first, last, n);
+  return Impl::shift_right_exespace_impl(label, ex, first, last, n);
 }
 
-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 shift_right(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& view,
                  typename decltype(begin(view))::difference_type n) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::shift_right_impl("Kokkos::shift_right_view_api_default", ex,
-                                begin(view), end(view), n);
+  return Impl::shift_right_exespace_impl("Kokkos::shift_right_view_api_default",
+                                         ex, begin(view), end(view), n);
 }
 
-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 shift_right(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& view,
                  typename decltype(begin(view))::difference_type n) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
-  return Impl::shift_right_impl(label, ex, begin(view), end(view), n);
+  return Impl::shift_right_exespace_impl(label, ex, begin(view), end(view), n);
+}
+
+//
+// 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
+shift_right(const TeamHandleType& teamHandle, IteratorType first,
+            IteratorType last, typename IteratorType::difference_type n) {
+  return Impl::shift_right_team_impl(teamHandle, first, last, n);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto shift_right(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& view,
+    typename decltype(begin(view))::difference_type n) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view);
+  return Impl::shift_right_team_impl(teamHandle, begin(view), end(view), n);
 }
 
 }  // namespace Experimental