remove the copies between soa and aos for the particle positions

1 files changed, 18 insertions, 80 deletions

diff --git a/Source/Particles/WarpXParticleContainer.cpp b/Source/Particles/WarpXParticleContainer.cpp
index e6f729935..eab410a18 100644
--- a/Source/Particles/WarpXParticleContainer.cpp
+++ b/Source/Particles/WarpXParticleContainer.cpp
@@ -21,56 +21,6 @@ WarpXParIter::WarpXParIter (ContainerType& pc, int level)
 {
 }
 
-#if (AMREX_SPACEDIM == 2)
-void
-WarpXParIter::GetPosition (Gpu::ManagedDeviceVector<ParticleReal>& xp,
-                           Gpu::ManagedDeviceVector<ParticleReal>& yp,
-                           Gpu::ManagedDeviceVector<ParticleReal>& zp) const
-{
-    amrex::ParIter<0,0,PIdx::nattribs>::GetPosition(xp, zp);
-#ifdef WARPX_DIM_RZ
-    const auto& attribs = GetAttribs();
-    const auto& thetap = attribs[PIdx::theta];
-    yp.resize(xp.size());
-    ParticleReal* const AMREX_RESTRICT x = xp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT y = yp.dataPtr();
-    const ParticleReal* const AMREX_RESTRICT theta = thetap.dataPtr();
-    amrex::ParallelFor( xp.size(),
-        [=] AMREX_GPU_DEVICE (long i) {
-        // The x stored in the particles is actually the radius
-        y[i] = x[i]*std::sin(theta[i]);
-        x[i] = x[i]*std::cos(theta[i]);
-    });
-#else
-    yp.resize(xp.size(), std::numeric_limits<ParticleReal>::quiet_NaN());
-#endif
-}
-
-void
-WarpXParIter::SetPosition (const Gpu::ManagedDeviceVector<ParticleReal>& xp,
-                           const Gpu::ManagedDeviceVector<ParticleReal>& yp,
-                           const Gpu::ManagedDeviceVector<ParticleReal>& zp)
-{
-#ifdef WARPX_DIM_RZ
-    auto& attribs = GetAttribs();
-    auto& thetap = attribs[PIdx::theta];
-    Gpu::ManagedDeviceVector<ParticleReal> rp(xp.size());
-    const ParticleReal* const AMREX_RESTRICT x = xp.dataPtr();
-    const ParticleReal* const AMREX_RESTRICT y = yp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT r = rp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT theta = thetap.dataPtr();
-    amrex::ParallelFor( xp.size(),
-        [=] AMREX_GPU_DEVICE (long i) {
-        theta[i] = std::atan2(y[i], x[i]);
-        r[i] = std::sqrt(x[i]*x[i] + y[i]*y[i]);
-    });
-    amrex::ParIter<0,0,PIdx::nattribs>::SetPosition(rp, zp);
-#else
-    amrex::ParIter<0,0,PIdx::nattribs>::SetPosition(xp, zp);
-#endif
-}
-#endif
-
 WarpXParticleContainer::WarpXParticleContainer (AmrCore* amr_core, int ispecies)
     : ParticleContainer<0,0,PIdx::nattribs>(amr_core->GetParGDB())
     , species_id(ispecies)
@@ -107,10 +57,6 @@ WarpXParticleContainer::WarpXParticleContainer (AmrCore* amr_core, int ispecies)
     local_jx.resize(num_threads);
     local_jy.resize(num_threads);
     local_jz.resize(num_threads);
-    m_xp.resize(num_threads);
-    m_yp.resize(num_threads);
-    m_zp.resize(num_threads);
-
 }
 
 void
@@ -340,10 +286,9 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     // CPU, tiling: deposit into local_jx
     // (same for jx and jz)
 
-    ParticleReal* AMREX_RESTRICT xp = m_xp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT yp = m_yp[thread_num].dataPtr() + offset;
-
+    const auto& aos = pti.GetArrayOfStructs();
+    const ParticleType* AMREX_RESTRICT pstruct = aos().dataPtr() + offset;
+    
     // Lower corner of tile box physical domain
     // Note that this includes guard cells since it is after tilebox.ngrow
     const Dim3 lo = lbound(tilebox);
@@ -353,19 +298,19 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     if (WarpX::current_deposition_algo == CurrentDepositionAlgo::Esirkepov) {
         if        (WarpX::nox == 1){
             doEsirkepovDepositionShapeN<1>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 2){
             doEsirkepovDepositionShapeN<2>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 3){
             doEsirkepovDepositionShapeN<3>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
@@ -373,19 +318,19 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     } else {
         if        (WarpX::nox == 1){
             doDepositionShapeN<1>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
         } else if (WarpX::nox == 2){
             doDepositionShapeN<2>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
         } else if (WarpX::nox == 3){
             doDepositionShapeN<3>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                pstruct, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
@@ -479,9 +424,8 @@ WarpXParticleContainer::DepositCharge (WarpXParIter& pti, RealVector& wp,
     // GPU, no tiling: deposit directly in rho
     // CPU, tiling: deposit into local_rho
 
-    ParticleReal* AMREX_RESTRICT xp = m_xp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT yp = m_yp[thread_num].dataPtr() + offset;
+    const auto& aos = pti.GetArrayOfStructs();
+    const ParticleType* AMREX_RESTRICT pstruct = aos().dataPtr() + offset;
 
     // Lower corner of tile box physical domain
     // Note that this includes guard cells since it is after tilebox.ngrow
@@ -491,13 +435,13 @@ WarpXParticleContainer::DepositCharge (WarpXParIter& pti, RealVector& wp,
 
     BL_PROFILE_VAR_START(blp_ppc_chd);
     if        (WarpX::nox == 1){
-        doChargeDepositionShapeN<1>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<1>(pstruct, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     } else if (WarpX::nox == 2){
-        doChargeDepositionShapeN<2>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<2>(pstruct, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     } else if (WarpX::nox == 3){
-        doChargeDepositionShapeN<3>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<3>(pstruct, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     }
     BL_PROFILE_VAR_STOP(blp_ppc_chd);
@@ -536,8 +480,6 @@ WarpXParticleContainer::DepositCharge (amrex::Vector<std::unique_ptr<amrex::Mult
             const long np = pti.numParticles();
             auto& wp = pti.GetAttribs(PIdx::w);
 
-            pti.GetPosition(m_xp[thread_num], m_yp[thread_num], m_zp[thread_num]);
-
             int* AMREX_RESTRICT ion_lev;
             if (do_field_ionization){
                 ion_lev = pti.GetiAttribs(particle_icomps["ionization_level"]).dataPtr();
@@ -606,8 +548,6 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local)
             const long np = pti.numParticles();
             auto& wp = pti.GetAttribs(PIdx::w);
 
-            pti.GetPosition(m_xp[thread_num], m_yp[thread_num], m_zp[thread_num]);
-
             int* AMREX_RESTRICT ion_lev;
             if (do_field_ionization){
                 ion_lev = pti.GetiAttribs(particle_icomps["ionization_level"]).dataPtr();
@@ -809,15 +749,13 @@ WarpXParticleContainer::PushX (int lev, amrex::Real dt)
             amrex::ParallelFor( pti.numParticles(),
                 [=] AMREX_GPU_DEVICE (long i) {
                     ParticleType& p = pstructs[i]; // Particle object that gets updated
-                    ParticleReal x, y, z; // Temporary variables
 #ifndef WARPX_DIM_RZ
-                    GetPosition( x, y, z, p ); // Initialize x, y, z
-                    UpdatePosition( x, y, z, ux[i], uy[i], uz[i], dt);
-                    SetPosition( p, x, y, z ); // Update the object p
-#else
+                    UpdatePosition( p, ux[i], uy[i], uz[i], dt);
+#else                    
                     // For WARPX_DIM_RZ, the particles are still pushed in 3D Cartesian
+                    ParticleReal x, y, z; // Temporary variables
                     GetCartesianPositionFromCylindrical( x, y, z, p, theta[i] );
-                    UpdatePosition( x, y, z, ux[i], uy[i], uz[i], dt);
+                    UpdatePosition( p, ux[i], uy[i], uz[i], dt);
                     SetCylindricalPositionFromCartesian( p, theta[i], x, y, z );
 #endif
                 }