team-level std algos: common code needed (kokkos#6199)

* add common code needed for all team level * remove unused include
nliber · Jun 13, 2023 · 7cc2e45 · 7cc2e45
1 parent 63d695c
commit 7cc2e45
Show file tree

Hide file tree

Showing 4 changed files with 424 additions and 21 deletions.
diff --git a/algorithms/src/std_algorithms/impl/Kokkos_Constraints.hpp b/algorithms/src/std_algorithms/impl/Kokkos_Constraints.hpp
@@ -113,16 +113,18 @@ struct iterators_are_accessible_from<ExeSpace, Head, Tail...> {
       iterators_are_accessible_from<ExeSpace, Tail...>::value;
 };
 
-template <class ExecutionSpace, class... IteratorTypes>
+template <class ExecutionSpaceOrTeamHandleType, class... IteratorTypes>
 KOKKOS_INLINE_FUNCTION constexpr void
-static_assert_random_access_and_accessible(const ExecutionSpace& /* ex */,
-                                           IteratorTypes... /* iterators */) {
+static_assert_random_access_and_accessible(
+    const ExecutionSpaceOrTeamHandleType& /* ex_or_th*/,
+    IteratorTypes... /* iterators */) {
   static_assert(
       are_random_access_iterators<IteratorTypes...>::value,
       "Currently, Kokkos standard algorithms require random access iterators.");
-  static_assert(
-      iterators_are_accessible_from<ExecutionSpace, IteratorTypes...>::value,
-      "Incompatible view/iterator and execution space");
+  static_assert(iterators_are_accessible_from<
+                    typename ExecutionSpaceOrTeamHandleType::execution_space,
+                    IteratorTypes...>::value,
+                "Incompatible view/iterator and execution space");
 }
 
 //
@@ -182,10 +184,10 @@ struct not_openmptarget {
 #endif
 };
 
-template <class ExecutionSpace>
+template <class ExecutionSpaceOrTeamHandleType>
 KOKKOS_INLINE_FUNCTION constexpr void static_assert_is_not_openmptarget(
-    const ExecutionSpace&) {
-  static_assert(not_openmptarget<ExecutionSpace>::value,
+    const ExecutionSpaceOrTeamHandleType& /*ex_or_th*/) {
+  static_assert(not_openmptarget<ExecutionSpaceOrTeamHandleType>::value,
                 "Currently, Kokkos standard algorithms do not support custom "
                 "comparators in OpenMPTarget");
 }
@@ -194,7 +196,8 @@ KOKKOS_INLINE_FUNCTION constexpr void static_assert_is_not_openmptarget(
 // valid range
 //
 template <class IteratorType>
-void expect_valid_range(IteratorType first, IteratorType last) {
+KOKKOS_INLINE_FUNCTION void expect_valid_range(IteratorType first,
+                                               IteratorType last) {
   // this is a no-op for release
   KOKKOS_EXPECTS(last >= first);
   // avoid compiler complaining when KOKKOS_EXPECTS is no-op

diff --git a/algorithms/unit_tests/TestStdAlgorithmsCommon.cpp b/algorithms/unit_tests/TestStdAlgorithmsCommon.cpp
@@ -25,5 +25,11 @@ std::string view_tag_to_string(StridedTwoTag) { return "stride2_view"; }
 
 std::string view_tag_to_string(StridedThreeTag) { return "stride3_view"; }
 
+std::string view_tag_to_string(StridedTwoRowsTag) { return "stride2rows_view"; }
+
+std::string view_tag_to_string(StridedThreeRowsTag) {
+  return "stride3rows_view";
+}
+
 }  // namespace stdalgos
 }  // namespace Test