2 files changed, 50 insertions, 43 deletions

diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index 6874e2df6..7e247bd8c 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -41,40 +41,23 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     RealVector& uxp, RealVector& uyp, RealVector& uzp, RealVector& wp)
 {
     BL_PROFILE("PPC::Evolve::partition");
+
+    auto& aos = pti.GetArrayOfStructs();
     Gpu::ManagedDeviceVector<int> inexflag;
     inexflag.resize(np);
+    Gpu::ManagedDeviceVector<long> pid;
+    pid.resize(np);
 
     // Select the largest buffer first
     iMultiFab const* bmasks =
         (WarpX::n_field_gather_buffer >= WarpX::n_current_deposition_buffer) ?
         gather_masks : current_masks;
 
-    // For each particle, find whether it is in the large buffer, by looking up the mask
-    const Array4<const int>& msk = (*bmasks)[pti].array();
-    auto& aos = pti.GetArrayOfStructs();
-    ParticleType * AMREX_RESTRICT pstructs = &(aos[0]);
-    int * inexflag_ptr = inexflag.dataPtr();
-    const Geometry& geom = Geom(lev);
-    const Real prob_lo_x = geom.ProbLo(0);
-    const Real prob_lo_y = geom.ProbLo(1);
-    const Real prob_lo_z = geom.ProbLo(2);
-    const Real inv_dx = geom.InvCellSize(0);
-    const Real inv_dy = geom.InvCellSize(1);
-    const Real inv_dz = geom.InvCellSize(2);
-    const IntVect domain_small_end = geom.Domain().smallEnd();
-    amrex::ParallelFor( np,
-         [=] AMREX_GPU_DEVICE (long i) {
-
-             const IntVect& iv = findParticleIndex(pstructs[i],
-                 prob_lo_x, prob_lo_y, prob_lo_z,
-                 inv_dx, inv_dy, inv_dz, domain_small_end );
-             inexflag_ptr[i] = msk(iv);
-         }
-     );
-
-    // Partition the particles according whether they are in the large buffer or not
-    Gpu::ManagedDeviceVector<long> pid;
-    pid.resize(np);
+    // For each particle, find whether it is in the large buffer,
+    // by looking up the mask. Store the answer in `inexflag`
+    amrex::ParallelFor( np, fillBufferFlag( pti, bmasks, inexflag, Geom(lev) ) );
+
+    // Partition the particles according to whether they are in the large buffer or not
     fillWithConsecutiveIntegers( pid );
     auto sep = std::stable_partition(pid.begin(), pid.end(),
                                      [&inexflag](long id) { return inexflag[id]; });
diff --git a/Source/Particles/Sorting/SortingUtils.H b/Source/Particles/Sorting/SortingUtils.H
index 7dec21446..1133ccab5 100644
--- a/Source/Particles/Sorting/SortingUtils.H
+++ b/Source/Particles/Sorting/SortingUtils.H
@@ -1,8 +1,9 @@
 #ifndef WARPX_PARTICLES_SORTING_SORTINGUTILS_H_
 #define WARPX_PARTICLES_SORTING_SORTINGUTILS_H_
 
-#include <AMReX_Gpu.H>
+#include <WarpXParticleContainer.H>
 #include <AMReX_CudaContainers.H>
+#include <AMReX_Gpu.H>
 
 // TODO: Add documentation
 void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
@@ -16,25 +17,48 @@ void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
 }
 
 // TODO: Add documentation
-template< typename PType >
-AMREX_GPU_HOST_DEVICE AMREX_INLINE
-amrex::IntVect findParticleIndex ( PType p,
-                    const amrex::Real prob_lo_x,
-                    const amrex::Real prob_lo_y,
-                    const amrex::Real prob_lo_z,
-                    const amrex::Real inv_dx,
-                    const amrex::Real inv_dy,
-                    const amrex::Real inv_dz,
-                    const amrex::IntVect domain_small_end )
+class fillBufferFlag
 {
-    amrex::IntVect iv;
-    AMREX_D_TERM(iv[0]=static_cast<int>(floor((p.m_rdata.pos[0]-prob_lo_x)*inv_dx));,
-                 iv[1]=static_cast<int>(floor((p.m_rdata.pos[1]-prob_lo_y)*inv_dy));,
-                 iv[2]=static_cast<int>(floor((p.m_rdata.pos[2]-prob_lo_z)*inv_dz)););
+    public:
+        fillBufferFlag( WarpXParIter& pti, const amrex::iMultiFab* bmasks,
+                        amrex::Gpu::ManagedDeviceVector<int>& inexflag,
+                        const amrex::Geometry& geom ) {
+
+            // Extract simple structure that can be used directly on the GPU
+            m_particles = &(pti.GetArrayOfStructs()[0]);
+            m_buffer_mask = (*bmasks)[pti].array();
+            m_inexflag_ptr = inexflag.dataPtr();
+            m_domain_small_end = geom.Domain().smallEnd();
+            for (int idim=0; idim<AMREX_SPACEDIM; idim++) {
+                m_prob_lo[idim] = geom.ProbLo(idim);
+                m_inv_cell_size[idim] = geom.InvCellSize(idim);
+            }
+        };
+
 
-    iv += domain_small_end;
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator()( const long i ) const {
+            // Select a particle
+            auto& p = m_particles[i];
+            // Find the index of the cell where this particle is located
+            amrex::IntVect iv;
+            for (int idim=0; idim<AMREX_SPACEDIM; idim++) {
+                iv[idim] = static_cast<int>(floor(
+                    (p.m_rdata.pos[idim]-m_prob_lo[idim])*m_inv_cell_size[idim]
+                ));
+            }
+            // 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);
+        };
 
-    return iv;
+    private:
+        amrex::Real m_prob_lo[AMREX_SPACEDIM];
+        amrex::Real m_inv_cell_size[AMREX_SPACEDIM];
+        amrex::IntVect m_domain_small_end;
+        int* m_inexflag_ptr;
+        WarpXParticleContainer::ParticleType* m_particles;
+        amrex::Array4<const int> m_buffer_mask;
 };
 
 #endif // WARPX_PARTICLES_SORTING_SORTINGUTILS_H_