Minor change to unify atomic_inc with consecutive fetch

Co-authored-by: Damien L-G <dalg24+github@gmail.com>
arborx · Oct 15, 2024 · e9286b0 · e9286b0
1 parent 7a78f6b
commit e9286b0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 11 deletions.
diff --git a/src/cluster/detail/ArborX_FDBSCAN.hpp b/src/cluster/detail/ArborX_FDBSCAN.hpp
@@ -33,14 +33,12 @@ struct CountUpToN
   template <typename Query, typename Value>
   KOKKOS_FUNCTION auto operator()(Query const &query, Value const &) const
   {
-    auto i = getData(query);
-    Kokkos::atomic_inc(&_counts(i));
+    int const i = getData(query);
+    int &count = _counts(i);
+    if (Kokkos::atomic_inc_fetch(&count) >= _n)
+      return ArborX::CallbackTreeTraversalControl::early_exit;
 
-    if (_counts(i) < _n)
-      return ArborX::CallbackTreeTraversalControl::normal_continuation;
-
-    // Once count reaches threshold, terminate the traversal.
-    return ArborX::CallbackTreeTraversalControl::early_exit;
+    return ArborX::CallbackTreeTraversalControl::normal_continuation;
   }
 };
 

diff --git a/src/cluster/detail/ArborX_FDBSCANDenseBox.hpp b/src/cluster/detail/ArborX_FDBSCANDenseBox.hpp
@@ -74,16 +74,14 @@ struct CountUpToN_DenseBox
         int j = _permute(jj);
         if (distance(query_point, _primitives(j)) <= eps)
         {
-          Kokkos::atomic_inc(&count);
-          if (count >= _n)
+          if (Kokkos::atomic_inc_fetch(&count) >= _n)
             return ArborX::CallbackTreeTraversalControl::early_exit;
         }
       }
     }
     else
     {
-      Kokkos::atomic_inc(&count);
-      if (count >= _n)
+      if (Kokkos::atomic_inc_fetch(&count) >= _n)
         return ArborX::CallbackTreeTraversalControl::early_exit;
     }