fix single as per Christian's suggestion

masterleinad · Sep 7, 2023 · eb8ee28 · eb8ee28
1 parent f4d7ea5
commit eb8ee28
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 36 deletions.
diff --git a/algorithms/src/std_algorithms/impl/Kokkos_CopyIf.hpp b/algorithms/src/std_algorithms/impl/Kokkos_CopyIf.hpp
@@ -126,17 +126,18 @@ KOKKOS_FUNCTION OutputIterator copy_if_team_impl(
   // return_value, so temporarily serial implementation is used here
   const std::size_t num_elements = Kokkos::Experimental::distance(first, last);
   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[countView()++] = myval;
-      }
-    }
-  });
-  teamHandle.team_barrier();
+  Kokkos::single(
+      Kokkos::PerTeam(teamHandle),
+      [=](std::size_t& lcount) {
+        lcount = 0;
+        for (std::size_t i = 0; i < num_elements; ++i) {
+          const auto& myval = first[i];
+          if (pred(myval)) {
+            d_first[lcount++] = myval;
+          }
+        }
+      },
+      count);
 
   return d_first + count;
 }

diff --git a/algorithms/src/std_algorithms/impl/Kokkos_RemoveAllVariants.hpp b/algorithms/src/std_algorithms/impl/Kokkos_RemoveAllVariants.hpp
@@ -160,16 +160,17 @@ remove_if_team_impl(const TeamHandleType& teamHandle, IteratorType first,
 
     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[countView()++] = std::move(first[i]);
-          }
-        }
-      });
-      teamHandle.team_barrier();
+      Kokkos::single(
+          Kokkos::PerTeam(teamHandle),
+          [=](std::size_t& lcount) {
+            lcount = 0;
+            for (std::size_t i = 0; i < num_elements; ++i) {
+              if (!pred(first[i])) {
+                first[lcount++] = std::move(first[i]);
+              }
+            }
+          },
+          count);
     }
 
     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,21 +145,21 @@ KOKKOS_FUNCTION OutputIterator unique_copy_team_impl(
     // if this can be done in parallel
 
     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 - 1; ++i) {
-        const auto& val_i   = first[i];
-        const auto& val_ip1 = first[i + 1];
-        if (!pred(val_i, val_ip1)) {
-          d_first[countView()++] = val_i;
-        }
-      }
-
-      // we need to copy the last element always
-      d_first[countView()++] = first[num_elements - 1];
-    });
-    teamHandle.team_barrier();
+    Kokkos::single(
+        Kokkos::PerTeam(teamHandle),
+        [=](std::size_t& lcount) {
+          lcount = 0;
+          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[lcount++] = val_i;
+            }
+          }
+          // we need to copy the last element always
+          d_first[lcount++] = first[num_elements - 1];
+        },
+        count);
 
     // return the correct iterator: we need +1 here because we need to
     // return iterator to the element past the last element copied