use single

tcclevenger · Sep 7, 2023 · dd5c624 · dd5c624
1 parent c4c9ed5
commit dd5c624
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 19 deletions.
diff --git a/algorithms/src/std_algorithms/impl/Kokkos_CopyIf.hpp b/algorithms/src/std_algorithms/impl/Kokkos_CopyIf.hpp
@@ -125,18 +125,18 @@ KOKKOS_FUNCTION OutputIterator copy_if_team_impl(
   // FIXME: there is no parallel_scan overload that accepts TeamThreadRange and
   // return_value, so temporarily serial implementation is used here
   const std::size_t num_elements = Kokkos::Experimental::distance(first, last);
-  std::size_t count              = {};
-  if (teamHandle.team_rank() == 0) {
+  std::size_t count              = 0;
+  Kokkos::View<std::size_t, typename TeamHandleType::execution_space> countView(
+      &count);
+  Kokkos::single(Kokkos::PerTeam(teamHandle), [=]() {
     for (std::size_t i = 0; i < num_elements; ++i) {
       const auto& myval = first[i];
       if (pred(myval)) {
-        d_first[count++] = myval;
+        d_first[countView()++] = myval;
       }
     }
-  }
-
-  // no need for barrier because calling broadcast implicitly blocks
-  teamHandle.team_broadcast(count, 0);
+  });
+  teamHandle.team_barrier();
 
   return d_first + count;
 }

diff --git a/algorithms/src/std_algorithms/impl/Kokkos_RemoveAllVariants.hpp b/algorithms/src/std_algorithms/impl/Kokkos_RemoveAllVariants.hpp
@@ -158,17 +158,19 @@ remove_if_team_impl(const TeamHandleType& teamHandle, IteratorType first,
     const std::size_t num_elements =
         ::Kokkos::Experimental::distance(first, last);
 
-    if (teamHandle.team_rank() == 0) {
-      if (remove_count > 0) {
-        std::size_t count = 0;
+    if (remove_count > 0) {
+      std::size_t count = 0;
+      Kokkos::View<std::size_t, typename TeamHandleType::execution_space>
+          countView(&count);
+      Kokkos::single(Kokkos::PerTeam(teamHandle), [=]() {
         for (std::size_t i = 0; i < num_elements; ++i) {
           if (!pred(first[i])) {
-            first[count++] = std::move(first[i]);
+            first[countView()++] = std::move(first[i]);
           }
         }
-      }
+      });
+      teamHandle.team_barrier();
     }
-    teamHandle.team_barrier();
 
     return first + num_elements - remove_count;
   }

diff --git a/algorithms/src/std_algorithms/impl/Kokkos_UniqueCopy.hpp b/algorithms/src/std_algorithms/impl/Kokkos_UniqueCopy.hpp
@@ -145,20 +145,21 @@ KOKKOS_FUNCTION OutputIterator unique_copy_team_impl(
     // if this can be done in parallel
 
     std::size_t count = 0;
-    if (teamHandle.team_rank() == 0) {
+    Kokkos::View<std::size_t, typename TeamHandleType::execution_space>
+        countView(&count);
+    Kokkos::single(Kokkos::PerTeam(teamHandle), [=]() {
       for (std::size_t i = 0; i < num_elements - 1; ++i) {
         const auto& val_i   = first[i];
         const auto& val_ip1 = first[i + 1];
         if (!pred(val_i, val_ip1)) {
-          d_first[count++] = val_i;
+          d_first[countView()++] = val_i;
         }
       }
 
       // we need to copy the last element always
-      d_first[count++] = first[num_elements - 1];
-    }
-
-    teamHandle.team_broadcast(count, 0);
+      d_first[countView()++] = first[num_elements - 1];
+    });
+    teamHandle.team_barrier();
 
     // return the correct iterator: we need +1 here because we need to
     // return iterator to the element past the last element copied