Perform partition in smaller buffer

2 files changed, 16 insertions, 17 deletions

diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index d64b9fa25..138f28105 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -58,7 +58,7 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
         gather_masks : current_masks;
     // - For each particle, find whether it is in the larger buffer,
     //   by looking up the mask. Store the answer in `inexflag`.
-    amrex::ParallelFor( np, fillBufferFlag(pti, bmasks, inexflag, Geom(lev)) );
+    amrex::ParallelFor( np, fillBufferFlag(pti, bmasks, inexflag, Geom(lev), 0) );
     // - Find the indices that reorder particles so that the last particles
     //   are in the larger buffer
     fillWithConsecutiveIntegers( pid );
@@ -67,6 +67,7 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     // reorder the particles, and `sep` is the position in the array that
     // separates the particles that deposit/gather on the fine patch (first part)
     // and the particles that deposit/gather in the buffers (last part)
+    long n_fine = iteratorDistance(pid.begin(), sep); // Number of particles on fine patch
 
     // Second, among particles that are in the larger buffer, partition
     // particles into the smaller buffer
@@ -79,29 +80,25 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     } else {
         int n_buf;
         if (bmasks == gather_masks) {
-            nfine_gather = iteratorDistance(pid.begin(), sep);
+            nfine_gather = n_fine;
             bmasks = current_masks;
             n_buf = WarpX::n_current_deposition_buffer;
         } else {
-            nfine_current = iteratorDistance(pid.begin(), sep);
+            nfine_current = n_fine;
             bmasks = gather_masks;
             n_buf = WarpX::n_field_gather_buffer;
         }
         if (n_buf > 0)
         {
-            const auto& msk2 = (*bmasks)[pti];
-            for (auto it = sep; it != pid.end(); ++it) {
-                const long id = *it;
-                const IntVect& iv = Index(aos[id], lev);
-                inexflag[id] = msk2(iv);
-            }
-
-            auto sep2 = std::stable_partition(sep, pid.end(),
-                                              [&inexflag](long id) {return inexflag[id];});
+            // - For each particle in the large buffer, find whether it is in
+            // the smaller buffer, by looking up the mask. Store the answer in `inexflag`.
+            amrex::ParallelFor( np - n_fine,
+               fillBufferFlag(pti, bmasks, inexflag, Geom(lev), n_fine) );
+            auto sep2 = stablePartition( sep, pid.end(), inexflag );
             if (bmasks == gather_masks) {
-                nfine_gather = std::distance(pid.begin(), sep2);
+                nfine_gather = iteratorDistance(pid.begin(), sep2);
             } else {
-                nfine_current = std::distance(pid.begin(), sep2);
+                nfine_current = iteratorDistance(pid.begin(), sep2);
             }
         }
     }
diff --git a/Source/Particles/Sorting/SortingUtils.H b/Source/Particles/Sorting/SortingUtils.H
index 7d53a352e..a17fe85db 100644
--- a/Source/Particles/Sorting/SortingUtils.H
+++ b/Source/Particles/Sorting/SortingUtils.H
@@ -61,7 +61,7 @@ class fillBufferFlag
     public:
         fillBufferFlag( WarpXParIter& pti, const amrex::iMultiFab* bmasks,
                         amrex::Gpu::ManagedDeviceVector<int>& inexflag,
-                        const amrex::Geometry& geom ) {
+                        const amrex::Geometry& geom, long start_index=0 ) {
 
             // Extract simple structure that can be used directly on the GPU
             m_particles = &(pti.GetArrayOfStructs()[0]);
@@ -72,13 +72,14 @@ class fillBufferFlag
                 m_prob_lo[idim] = geom.ProbLo(idim);
                 m_inv_cell_size[idim] = geom.InvCellSize(idim);
             }
+            m_start_index = start_index;
         };
 
 
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         void operator()( const long i ) const {
             // Select a particle
-            auto& p = m_particles[i];
+            auto& p = m_particles[i+m_start_index];
             // Find the index of the cell where this particle is located
             amrex::IntVect iv;
             for (int idim=0; idim<AMREX_SPACEDIM; idim++) {
@@ -88,7 +89,7 @@ class fillBufferFlag
             }
             // Find the value of the buffer flag in this cell and
             // store it at the corresponding particle position in the array `inexflag`
-            m_inexflag_ptr[i] = m_buffer_mask(iv);
+            m_inexflag_ptr[i+m_start_index] = m_buffer_mask(iv);
         };
 
     private:
@@ -98,6 +99,7 @@ class fillBufferFlag
         int* m_inexflag_ptr;
         WarpXParticleContainer::ParticleType* m_particles;
         amrex::Array4<const int> m_buffer_mask;
+        long m_start_index;
 };
 
 #endif // WARPX_PARTICLES_SORTING_SORTINGUTILS_H_