From 8575478f370f3db2525a8d5444e6949776e9da29 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Sep 2019 08:51:00 -0700
Subject: Implemented stable partition

---
 Source/Particles/Sorting/SortingUtils.H | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

(limited to 'Source/Particles/Sorting/SortingUtils.H')

diff --git a/Source/Particles/Sorting/SortingUtils.H b/Source/Particles/Sorting/SortingUtils.H
index 1133ccab5..ede59b53b 100644
--- a/Source/Particles/Sorting/SortingUtils.H
+++ b/Source/Particles/Sorting/SortingUtils.H
@@ -6,7 +6,8 @@
 #include <AMReX_Gpu.H>
 
 // TODO: Add documentation
-void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
+void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v )
+{
 #ifdef AMREX_USE_GPU
     // On GPU: Use thrust
     thrust::sequence( v.begin(), v.end() );
@@ -16,6 +17,31 @@ void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
 #endif
 }
 
+// TODO: Add documentation
+template< typename ForwardIterator >
+ForwardIterator stablePartition(ForwardIterator index_begin,
+                               ForwardIterator index_end,
+                               amrex::Gpu::ManagedDeviceVector<int>& predicate)
+{
+#ifdef AMREX_USE_GPU
+    // On GPU: Use thrust
+    int* AMREX_RESTRICT predicate_ptr = predicate.dataPtr();
+    ForwardIterator sep = thrust::stable_partition(
+                    thrust::cuda::par(Cuda::The_ThrustCachedAllocator()),
+                    index_begin, index_end,
+                    AMREX_GPU_HOST_DEVICE
+                    [predicate_ptr](long i) { return predicate_ptr[i]; }
+                );
+#else
+    // On CPU: Use std library
+    ForwardIterator sep = std::stable_partition(
+                    index_begin, index_end,
+                    [&predicate](long i) { return predicate[i]; }
+                );
+#endif
+    return sep;
+}
+
 // TODO: Add documentation
 class fillBufferFlag
 {
-- 
cgit v1.2.3