team-level std algos: part 4 (kokkos#6208)

Team-level implementation and tests for MinMaxElement Co-authored-by: Cezary Skrzyński Co-authored-by: Jakub Strzebonski
masterleinad · Aug 30, 2023 · de4380e · de4380e
1 parent e76708a
commit de4380e
Show file tree

Hide file tree

Showing 8 changed files with 848 additions and 63 deletions.
diff --git a/algorithms/src/std_algorithms/Kokkos_MaxElement.hpp b/algorithms/src/std_algorithms/Kokkos_MaxElement.hpp
@@ -23,81 +23,148 @@
 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>
 auto max_element(const ExecutionSpace& ex, IteratorType first,
                  IteratorType last) {
-  return Impl::min_or_max_element_impl<MaxFirstLoc>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLoc>(
       "Kokkos::max_element_iterator_api_default", ex, first, last);
 }
 
-template <class ExecutionSpace, class IteratorType>
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto max_element(const std::string& label, const ExecutionSpace& ex,
                  IteratorType first, IteratorType last) {
-  return Impl::min_or_max_element_impl<MaxFirstLoc>(label, ex, first, last);
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLoc>(label, ex, first,
+                                                             last);
 }
 
-template <class ExecutionSpace, class IteratorType, class ComparatorType>
+template <
+    typename ExecutionSpace, typename IteratorType, typename ComparatorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto max_element(const ExecutionSpace& ex, IteratorType first,
                  IteratorType last, ComparatorType comp) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MaxFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLocCustomComparator>(
       "Kokkos::max_element_iterator_api_default", ex, first, last,
       std::move(comp));
 }
 
-template <class ExecutionSpace, class IteratorType, class ComparatorType>
+template <
+    typename ExecutionSpace, typename IteratorType, typename ComparatorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto max_element(const std::string& label, const ExecutionSpace& ex,
                  IteratorType first, IteratorType last, ComparatorType comp) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MaxFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLocCustomComparator>(
       label, ex, first, last, std::move(comp));
 }
 
-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 max_element(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
-  return Impl::min_or_max_element_impl<MaxFirstLoc>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLoc>(
       "Kokkos::max_element_view_api_default", ex, begin(v), end(v));
 }
 
-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 max_element(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
-  return Impl::min_or_max_element_impl<MaxFirstLoc>(label, ex, begin(v),
-                                                    end(v));
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLoc>(label, ex,
+                                                             begin(v), end(v));
 }
 
-template <class ExecutionSpace, class DataType, class ComparatorType,
-          class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename ComparatorType,
+    typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto max_element(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v,
                  ComparatorType comp) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MaxFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLocCustomComparator>(
       "Kokkos::max_element_view_api_default", ex, begin(v), end(v),
       std::move(comp));
 }
 
-template <class ExecutionSpace, class DataType, class ComparatorType,
-          class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename ComparatorType,
+    typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto max_element(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v,
                  ComparatorType comp) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MaxFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MaxFirstLocCustomComparator>(
       label, ex, begin(v), end(v), std::move(comp));
 }
 
+//
+// 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 auto max_element(const TeamHandleType& teamHandle,
+                                 IteratorType first, IteratorType last) {
+  return Impl::min_or_max_element_team_impl<MaxFirstLoc>(teamHandle, first,
+                                                         last);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto max_element(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& v) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+
+  return Impl::min_or_max_element_team_impl<MaxFirstLoc>(teamHandle, begin(v),
+                                                         end(v));
+}
+
+template <typename TeamHandleType, typename IteratorType,
+          typename ComparatorType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto max_element(const TeamHandleType& teamHandle,
+                                 IteratorType first, IteratorType last,
+                                 ComparatorType comp) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  return Impl::min_or_max_element_team_impl<MaxFirstLocCustomComparator>(
+      teamHandle, first, last, std::move(comp));
+}
+
+template <typename TeamHandleType, typename DataType, typename ComparatorType,
+          typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto max_element(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& v, ComparatorType comp) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  return Impl::min_or_max_element_team_impl<MaxFirstLocCustomComparator>(
+      teamHandle, begin(v), end(v), std::move(comp));
+}
+
 }  // namespace Experimental
 }  // namespace Kokkos
 

diff --git a/algorithms/src/std_algorithms/Kokkos_MinElement.hpp b/algorithms/src/std_algorithms/Kokkos_MinElement.hpp
@@ -23,81 +23,148 @@
 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>
 auto min_element(const ExecutionSpace& ex, IteratorType first,
                  IteratorType last) {
-  return Impl::min_or_max_element_impl<MinFirstLoc>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLoc>(
       "Kokkos::min_element_iterator_api_default", ex, first, last);
 }
 
-template <class ExecutionSpace, class IteratorType>
+template <
+    typename ExecutionSpace, typename IteratorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto min_element(const std::string& label, const ExecutionSpace& ex,
                  IteratorType first, IteratorType last) {
-  return Impl::min_or_max_element_impl<MinFirstLoc>(label, ex, first, last);
+  return Impl::min_or_max_element_exespace_impl<MinFirstLoc>(label, ex, first,
+                                                             last);
 }
 
-template <class ExecutionSpace, class IteratorType, class ComparatorType>
+template <
+    typename ExecutionSpace, typename IteratorType, typename ComparatorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto min_element(const ExecutionSpace& ex, IteratorType first,
                  IteratorType last, ComparatorType comp) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MinFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLocCustomComparator>(
       "Kokkos::min_element_iterator_api_default", ex, first, last,
       std::move(comp));
 }
 
-template <class ExecutionSpace, class IteratorType, class ComparatorType>
+template <
+    typename ExecutionSpace, typename IteratorType, typename ComparatorType,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto min_element(const std::string& label, const ExecutionSpace& ex,
                  IteratorType first, IteratorType last, ComparatorType comp) {
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MinFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLocCustomComparator>(
       label, ex, first, last, std::move(comp));
 }
 
-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 min_element(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
-  return Impl::min_or_max_element_impl<MinFirstLoc>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLoc>(
       "Kokkos::min_element_view_api_default", ex, begin(v), end(v));
 }
 
-template <class ExecutionSpace, class DataType, class ComparatorType,
-          class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename ComparatorType,
+    typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto min_element(const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v,
                  ComparatorType comp) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MinFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLocCustomComparator>(
       "Kokkos::min_element_view_api_default", ex, begin(v), end(v),
       std::move(comp));
 }
 
-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 min_element(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
 
-  return Impl::min_or_max_element_impl<MinFirstLoc>(label, ex, begin(v),
-                                                    end(v));
+  return Impl::min_or_max_element_exespace_impl<MinFirstLoc>(label, ex,
+                                                             begin(v), end(v));
 }
 
-template <class ExecutionSpace, class DataType, class ComparatorType,
-          class... Properties>
+template <
+    typename ExecutionSpace, typename DataType, typename ComparatorType,
+    typename... Properties,
+    std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
 auto min_element(const std::string& label, const ExecutionSpace& ex,
                  const ::Kokkos::View<DataType, Properties...>& v,
                  ComparatorType comp) {
   Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
   Impl::static_assert_is_not_openmptarget(ex);
 
-  return Impl::min_or_max_element_impl<MinFirstLocCustomComparator>(
+  return Impl::min_or_max_element_exespace_impl<MinFirstLocCustomComparator>(
       label, ex, begin(v), end(v), std::move(comp));
 }
 
+//
+// 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 auto min_element(const TeamHandleType& teamHandle,
+                                 IteratorType first, IteratorType last) {
+  return Impl::min_or_max_element_team_impl<MinFirstLoc>(teamHandle, first,
+                                                         last);
+}
+
+template <typename TeamHandleType, typename DataType, typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto min_element(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& v) {
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+
+  return Impl::min_or_max_element_team_impl<MinFirstLoc>(teamHandle, begin(v),
+                                                         end(v));
+}
+
+template <typename TeamHandleType, typename IteratorType,
+          typename ComparatorType,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto min_element(const TeamHandleType& teamHandle,
+                                 IteratorType first, IteratorType last,
+                                 ComparatorType comp) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  return Impl::min_or_max_element_team_impl<MinFirstLocCustomComparator>(
+      teamHandle, first, last, std::move(comp));
+}
+
+template <typename TeamHandleType, typename DataType, typename ComparatorType,
+          typename... Properties,
+          std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
+KOKKOS_FUNCTION auto min_element(
+    const TeamHandleType& teamHandle,
+    const ::Kokkos::View<DataType, Properties...>& v, ComparatorType comp) {
+  Impl::static_assert_is_not_openmptarget(teamHandle);
+  Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
+  return Impl::min_or_max_element_team_impl<MinFirstLocCustomComparator>(
+      teamHandle, begin(v), end(v), std::move(comp));
+}
+
 }  // namespace Experimental
 }  // namespace Kokkos