team-level std algos: part 13 (kokkos#6351)

Add team level algorithms for inclusive scan
tcclevenger · Oct 5, 2023 · e40f026 · e40f026
1 parent 890148e
commit e40f026
Show file tree

Hide file tree

Showing 8 changed files with 1,299 additions and 181 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_InclusiveScan.hpp b/algorithms/src/std_algorithms/Kokkos_InclusiveScan.hpp
diff --git a/algorithms/src/std_algorithms/Kokkos_TransformInclusiveScan.hpp b/algorithms/src/std_algorithms/Kokkos_TransformInclusiveScan.hpp
@@ -23,40 +23,53 @@
 namespace Kokkos {
 namespace Experimental {
 
+//
+// overload set accepting execution space
+//
+
 // overload set 1 (no init value)
-template <class ExecutionSpace, class InputIteratorType,
-          class OutputIteratorType, class BinaryOpType, class UnaryOpType>
-std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
-                     InputIteratorType, OutputIteratorType>::value,
-                 OutputIteratorType>
-transform_inclusive_scan(const ExecutionSpace& ex, InputIteratorType first,
-                         InputIteratorType last, OutputIteratorType first_dest,
-                         BinaryOpType binary_op, UnaryOpType unary_op) {
+template <typename ExecutionSpace, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType,
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
+                  Kokkos::is_execution_space_v<ExecutionSpace>,
+              int> = 0>
+OutputIteratorType transform_inclusive_scan(const ExecutionSpace& ex,
+                                            InputIteratorType first,
+                                            InputIteratorType last,
+                                            OutputIteratorType first_dest,
+                                            BinaryOpType binary_op,
+                                            UnaryOpType unary_op) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::transform_inclusive_scan_impl(
+  return Impl::transform_inclusive_scan_exespace_impl(
       "Kokkos::transform_inclusive_scan_custom_functors_iterator_api", ex,
       first, last, first_dest, binary_op, unary_op);
 }
 
-template <class ExecutionSpace, class InputIteratorType,
-          class OutputIteratorType, class BinaryOpType, class UnaryOpType>
-std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
-                     InputIteratorType, OutputIteratorType>::value,
-                 OutputIteratorType>
-transform_inclusive_scan(const std::string& label, const ExecutionSpace& ex,
-                         InputIteratorType first, InputIteratorType last,
-                         OutputIteratorType first_dest, BinaryOpType binary_op,
-                         UnaryOpType unary_op) {
+template <typename ExecutionSpace, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType,
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
+                  Kokkos::is_execution_space_v<ExecutionSpace>,
+              int> = 0>
+OutputIteratorType transform_inclusive_scan(
+    const std::string& label, const ExecutionSpace& ex, InputIteratorType first,
+    InputIteratorType last, OutputIteratorType first_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::transform_inclusive_scan_impl(label, ex, first, last, first_dest,
-                                             binary_op, unary_op);
+  return Impl::transform_inclusive_scan_exespace_impl(
+      label, ex, first, last, first_dest, binary_op, unary_op);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryOpType,
-          class UnaryOpType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryOpType,
+    typename UnaryOpType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto transform_inclusive_scan(
     const ExecutionSpace& ex,
     const ::Kokkos::View<DataType1, Properties1...>& view_from,
@@ -66,15 +79,17 @@ auto transform_inclusive_scan(
   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::transform_inclusive_scan_impl(
+  return Impl::transform_inclusive_scan_exespace_impl(
       "Kokkos::transform_inclusive_scan_custom_functors_view_api", ex,
       KE::cbegin(view_from), KE::cend(view_from), KE::begin(view_dest),
       binary_op, unary_op);
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryOpType,
-          class UnaryOpType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryOpType,
+    typename UnaryOpType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto transform_inclusive_scan(
     const std::string& label, const ExecutionSpace& ex,
     const ::Kokkos::View<DataType1, Properties1...>& view_from,
@@ -84,46 +99,59 @@ auto transform_inclusive_scan(
   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::transform_inclusive_scan_impl(
+  return Impl::transform_inclusive_scan_exespace_impl(
       label, ex, KE::cbegin(view_from), KE::cend(view_from),
       KE::begin(view_dest), binary_op, unary_op);
 }
 
 // overload set 2 (init value)
-template <class ExecutionSpace, class InputIteratorType,
-          class OutputIteratorType, class BinaryOpType, class UnaryOpType,
-          class ValueType>
-std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
-                     InputIteratorType, OutputIteratorType>::value,
-                 OutputIteratorType>
-transform_inclusive_scan(const ExecutionSpace& ex, InputIteratorType first,
-                         InputIteratorType last, OutputIteratorType first_dest,
-                         BinaryOpType binary_op, UnaryOpType unary_op,
-                         ValueType init_value) {
+template <typename ExecutionSpace, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType, typename ValueType,
+
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
+                  Kokkos::is_execution_space_v<ExecutionSpace>,
+              int> = 0>
+OutputIteratorType transform_inclusive_scan(
+    const ExecutionSpace& ex, InputIteratorType first, InputIteratorType last,
+    OutputIteratorType first_dest, BinaryOpType binary_op, UnaryOpType unary_op,
+    ValueType init_value) {
   Impl::static_assert_is_not_openmptarget(ex);
-  return Impl::transform_inclusive_scan_impl(
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
+  return Impl::transform_inclusive_scan_exespace_impl(
       "Kokkos::transform_inclusive_scan_custom_functors_iterator_api", ex,
-      first, last, first_dest, binary_op, unary_op, init_value);
+      first, last, first_dest, binary_op, unary_op, std::move(init_value));
 }
 
-template <class ExecutionSpace, class InputIteratorType,
-          class OutputIteratorType, class BinaryOpType, class UnaryOpType,
-          class ValueType>
-std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
-                     InputIteratorType, OutputIteratorType>::value,
-                 OutputIteratorType>
-transform_inclusive_scan(const std::string& label, const ExecutionSpace& ex,
-                         InputIteratorType first, InputIteratorType last,
-                         OutputIteratorType first_dest, BinaryOpType binary_op,
-                         UnaryOpType unary_op, ValueType init_value) {
+template <typename ExecutionSpace, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType, typename ValueType,
+
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
+                  Kokkos::is_execution_space_v<ExecutionSpace>,
+              int> = 0>
+OutputIteratorType transform_inclusive_scan(
+    const std::string& label, const ExecutionSpace& ex, InputIteratorType first,
+    InputIteratorType last, OutputIteratorType first_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op, ValueType init_value) {
   Impl::static_assert_is_not_openmptarget(ex);
-  return Impl::transform_inclusive_scan_impl(label, ex, first, last, first_dest,
-                                             binary_op, unary_op, init_value);
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
+  return Impl::transform_inclusive_scan_exespace_impl(
+      label, ex, first, last, first_dest, binary_op, unary_op,
+      std::move(init_value));
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryOpType,
-          class UnaryOpType, class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryOpType,
+    typename UnaryOpType, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto transform_inclusive_scan(
     const ExecutionSpace& ex,
     const ::Kokkos::View<DataType1, Properties1...>& view_from,
@@ -132,16 +160,21 @@ auto transform_inclusive_scan(
   Impl::static_assert_is_not_openmptarget(ex);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
   namespace KE = ::Kokkos::Experimental;
-  return Impl::transform_inclusive_scan_impl(
+  return Impl::transform_inclusive_scan_exespace_impl(
       "Kokkos::transform_inclusive_scan_custom_functors_view_api", ex,
       KE::cbegin(view_from), KE::cend(view_from), KE::begin(view_dest),
-      binary_op, unary_op, init_value);
+      binary_op, unary_op, std::move(init_value));
 }
 
-template <class ExecutionSpace, class DataType1, class... Properties1,
-          class DataType2, class... Properties2, class BinaryOpType,
-          class UnaryOpType, class ValueType>
+template <
+    typename ExecutionSpace, typename DataType1, typename... Properties1,
+    typename DataType2, typename... Properties2, typename BinaryOpType,
+    typename UnaryOpType, typename ValueType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto transform_inclusive_scan(
     const std::string& label, const ExecutionSpace& ex,
     const ::Kokkos::View<DataType1, Properties1...>& view_from,
@@ -150,10 +183,97 @@ auto transform_inclusive_scan(
   Impl::static_assert_is_not_openmptarget(ex);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
   namespace KE = ::Kokkos::Experimental;
-  return Impl::transform_inclusive_scan_impl(
+  return Impl::transform_inclusive_scan_exespace_impl(
       label, ex, KE::cbegin(view_from), KE::cend(view_from),
-      KE::begin(view_dest), binary_op, unary_op, init_value);
+      KE::begin(view_dest), binary_op, unary_op, std::move(init_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.
+//
+
+// overload set 1 (no init value)
+template <typename TeamHandleType, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType,
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType> &&
+                  Kokkos::is_team_handle_v<TeamHandleType>,
+              int> = 0>
+KOKKOS_FUNCTION OutputIteratorType transform_inclusive_scan(
+    const TeamHandleType& teamHandle, InputIteratorType first,
+    InputIteratorType last, OutputIteratorType first_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+
+  return Impl::transform_inclusive_scan_team_impl(
+      teamHandle, first, last, first_dest, binary_op, unary_op);
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename DataType2, typename... Properties2, typename BinaryOpType,
+          typename UnaryOpType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto transform_inclusive_scan(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view_from,
+    const ::Kokkos::View<DataType2, Properties2...>& view_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  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::transform_inclusive_scan_team_impl(
+      teamHandle, KE::cbegin(view_from), KE::cend(view_from),
+      KE::begin(view_dest), binary_op, unary_op);
+}
+
+// overload set 2 (init value)
+template <typename TeamHandleType, typename InputIteratorType,
+          typename OutputIteratorType, typename BinaryOpType,
+          typename UnaryOpType, typename ValueType,
+          std::enable_if_t<
+              Impl::are_iterators_v<InputIteratorType, OutputIteratorType> &&
+                  Kokkos::is_team_handle_v<TeamHandleType>,
+              int> = 0>
+KOKKOS_FUNCTION OutputIteratorType transform_inclusive_scan(
+    const TeamHandleType& teamHandle, InputIteratorType first,
+    InputIteratorType last, OutputIteratorType first_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op, ValueType init_value) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
+  return Impl::transform_inclusive_scan_team_impl(
+      teamHandle, first, last, first_dest, binary_op, unary_op,
+      std::move(init_value));
+}
+
+template <typename TeamHandleType, typename DataType1, typename... Properties1,
+          typename DataType2, typename... Properties2, typename BinaryOpType,
+          typename UnaryOpType, typename ValueType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto transform_inclusive_scan(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType1, Properties1...>& view_from,
+    const ::Kokkos::View<DataType2, Properties2...>& view_dest,
+    BinaryOpType binary_op, UnaryOpType unary_op, ValueType init_value) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
+  static_assert(std::is_move_constructible_v<ValueType>,
+                "ValueType must be move constructible.");
+
+  namespace KE = ::Kokkos::Experimental;
+  return Impl::transform_inclusive_scan_team_impl(
+      teamHandle, KE::cbegin(view_from), KE::cend(view_from),
+      KE::begin(view_dest), binary_op, unary_op, std::move(init_value));
 }
 
 }  // namespace Experimental