Re-order all particle quantities when doing MR partition (#2888)

* Re-order all particle quantities when doing partition

* Fix compilation

1 files changed, 10 insertions, 27 deletions

diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index 0fba22559..18bebed18 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -48,16 +48,13 @@ using namespace amrex;
  *       in the deposition buffers or in the interior of the fine patch
  * \param gather_masks indicates, for each cell, whether that cell is
  *       in the gather buffers or in the interior of the fine patch
- * \param uxp, uyp, uzp, wp references to the particle momenta and weight
- *         (modified by this function)
  */
 void
 PhysicalParticleContainer::PartitionParticlesInBuffers(
     long& nfine_current, long& nfine_gather, long const np,
     WarpXParIter& pti, int const lev,
     iMultiFab const* current_masks,
-    iMultiFab const* gather_masks,
-    RealVector& uxp, RealVector& uyp, RealVector& uzp, RealVector& wp)
+    iMultiFab const* gather_masks )
 {
     WARPX_PROFILE("PhysicalParticleContainer::PartitionParticlesInBuffers");
 
@@ -134,31 +131,17 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     // Reorder the actual particle array, using the `pid` indices
     if (nfine_current != np || nfine_gather != np)
     {
-        // Temporary array for particle AoS
-        ParticleVector particle_tmp;
-        particle_tmp.resize(np);
-
-        // Copy particle AoS
-        auto& aos = pti.GetArrayOfStructs();
-        amrex::ParallelFor( np,
-            copyAndReorder<ParticleType>( aos(), particle_tmp, pid ) );
-        std::swap(aos(), particle_tmp);
-
-        // Temporary array for particle individual attributes
-        RealVector tmp;
-        tmp.resize(np);
+        // Prepare temporary particle tile to copy to
+        ParticleTileType ptile_tmp;
+        ptile_tmp.define(NumRuntimeRealComps(), NumRuntimeIntComps());
+        ptile_tmp.resize(np);
 
-        // Copy individual attributes
-        amrex::ParallelFor( np, copyAndReorder<Real>( wp, tmp, pid ) );
-        std::swap(wp, tmp);
-        amrex::ParallelFor( np, copyAndReorder<Real>( uxp, tmp, pid ) );
-        std::swap(uxp, tmp);
-        amrex::ParallelFor( np, copyAndReorder<Real>( uyp, tmp, pid ) );
-        std::swap(uyp, tmp);
-        amrex::ParallelFor( np, copyAndReorder<Real>( uzp, tmp, pid ) );
-        std::swap(uzp, tmp);
+        // Copy and re-order the data of the current particle tile
+        ParticleTileType& ptile = pti.GetParticleTile();
+        amrex::gatherParticles(ptile_tmp, ptile, np, pid.dataPtr());
+        ptile.swap(ptile_tmp);
 
-        // Make sure that the temporary arrays are not destroyed before
+        // Make sure that the temporary particle tile is not destroyed before
         // the GPU kernels finish running
         Gpu::streamSynchronize();
     }