use thrust::device_vector to store temporary particle quantities

7 files changed, 119 insertions, 56 deletions

diff --git a/Source/LaserParticleContainer.cpp b/Source/LaserParticleContainer.cpp
index 23ba2457c..bc8789ae3 100644
--- a/Source/LaserParticleContainer.cpp
+++ b/Source/LaserParticleContainer.cpp
@@ -312,7 +312,8 @@ LaserParticleContainer::Evolve (int lev,
 #pragma omp parallel
 #endif
     {
-	RealVector xp, yp, zp, giv, plane_Xp, plane_Yp, amplitude_E;
+	thrust::device_vector<Real> xp, yp, zp, giv;
+        RealVector plane_Xp, plane_Yp, amplitude_E;
         FArrayBox local_rho, local_jx, local_jy, local_jz;
 
         for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
@@ -395,7 +396,10 @@ LaserParticleContainer::Evolve (int lev,
                 const long nz = rholen[1]-1-2*ngRho;
 #endif
             	warpx_charge_deposition(data_ptr, &np,
-					xp.data(), yp.data(), zp.data(), wp.data(),
+					thrust::raw_pointer_cast(xp.data()),
+                                        thrust::raw_pointer_cast(yp.data()),
+                                        thrust::raw_pointer_cast(zp.data()),
+                                        wp.data(),
                                         &this->charge,
 					&xyzmin[0], &xyzmin[1], &xyzmin[2],
                                         &dx[0], &dx[1], &dx[2], &nx, &ny, &nz,
@@ -504,9 +508,13 @@ LaserParticleContainer::Evolve (int lev,
                 jx_ptr, &ngJ, jxntot,
                 jy_ptr, &ngJ, jyntot,
                 jz_ptr, &ngJ, jzntot,
-                &np, xp.data(), yp.data(), zp.data(),
+                &np,
+                thrust::raw_pointer_cast(xp.data()),
+                thrust::raw_pointer_cast(yp.data()),
+                thrust::raw_pointer_cast(zp.data()),
                 uxp.data(), uyp.data(), uzp.data(),
-                giv.data(), wp.data(), &this->charge,
+                thrust::raw_pointer_cast(giv.data()),
+                wp.data(), &this->charge,
                 &xyzmin[0], &xyzmin[1], &xyzmin[2],
                 &dt, &dx[0], &dx[1], &dx[2],
                 &WarpX::nox,&WarpX::noy,&WarpX::noz,
diff --git a/Source/PhysicalParticleContainer.H b/Source/PhysicalParticleContainer.H
index 873382f45..30de2e8c1 100644
--- a/Source/PhysicalParticleContainer.H
+++ b/Source/PhysicalParticleContainer.H
@@ -61,10 +61,10 @@ public:
                          amrex::Real dt) override;
 
     virtual void PushPX(WarpXParIter& pti,
-	                RealVector& xp,
-                        RealVector& yp,
-                        RealVector& zp,
-                        RealVector& giv,
+	                thrust::device_vector<amrex::Real>& xp,
+                        thrust::device_vector<amrex::Real>& yp,
+                        thrust::device_vector<amrex::Real>& zp,
+                        thrust::device_vector<amrex::Real>& giv,
                         amrex::Real dt);
 
     virtual void PushP (int lev, amrex::Real dt,
diff --git a/Source/PhysicalParticleContainer.cpp b/Source/PhysicalParticleContainer.cpp
index d9fd55afa..58ee695b9 100644
--- a/Source/PhysicalParticleContainer.cpp
+++ b/Source/PhysicalParticleContainer.cpp
@@ -607,7 +607,7 @@ PhysicalParticleContainer::FieldGather (int lev,
 #pragma omp parallel
 #endif
     {
-	RealVector xp, yp, zp;
+	thrust::device_vector<Real> xp, yp, zp;
 
 	for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
 	{
@@ -656,7 +656,10 @@ PhysicalParticleContainer::FieldGather (int lev,
 	    const int l_lower_order_in_v = warpx_l_lower_order_in_v();
             long lvect_fieldgathe = 64;
 	    warpx_geteb_energy_conserving(
-	       &np, xp.data(), yp.data(), zp.data(),
+	       &np,
+               thrust::raw_pointer_cast(xp.data()),
+               thrust::raw_pointer_cast(yp.data()),
+               thrust::raw_pointer_cast(zp.data()),
 	       Exp.data(),Eyp.data(),Ezp.data(),
 	       Bxp.data(),Byp.data(),Bzp.data(),
                ixyzmin,
@@ -722,7 +725,7 @@ PhysicalParticleContainer::Evolve (int lev,
 #pragma omp parallel
 #endif
     {
-	RealVector xp, yp, zp, giv;
+        thrust::device_vector<Real> xp, yp, zp, giv;
 
         std::unique_ptr<FArrayBox> local_rho(new FArrayBox());
         std::unique_ptr<FArrayBox> local_jx(new FArrayBox());
@@ -983,7 +986,10 @@ PhysicalParticleContainer::Evolve (int lev,
                     const long nz = rholen[1]-1-2*ngRho;
 #endif
                     warpx_charge_deposition(data_ptr, &np_current,
-                                            xp.data(), yp.data(), zp.data(), wp.data(),
+                                            thrust::raw_pointer_cast(xp.data()),
+                                            thrust::raw_pointer_cast(yp.data()),
+                                            thrust::raw_pointer_cast(zp.data()),
+                                            wp.data(),
                                             &this->charge,
                                             &xyzmin[0], &xyzmin[1], &xyzmin[2],
                                             &dx[0], &dx[1], &dx[2], &nx, &ny, &nz,
@@ -1028,9 +1034,9 @@ PhysicalParticleContainer::Evolve (int lev,
                     
                     long ncrse = np - nfine_current;
                     warpx_charge_deposition(data_ptr, &ncrse,
-                                            xp.data() + nfine_current,
-                                            yp.data() + nfine_current,
-                                            zp.data() + nfine_current,
+                                            thrust::raw_pointer_cast(xp.data() + nfine_current),
+                                            thrust::raw_pointer_cast(yp.data() + nfine_current),
+                                            thrust::raw_pointer_cast(zp.data() + nfine_current),
                                             wp.data() + nfine_current,
                                             &this->charge,
                                             &cxyzmin_tile[0], &cxyzmin_tile[1], &cxyzmin_tile[2],
@@ -1067,7 +1073,10 @@ PhysicalParticleContainer::Evolve (int lev,
                 BL_PROFILE_VAR_START(blp_pxr_fg);
 
                 warpx_geteb_energy_conserving(
-                    &np_gather, xp.data(), yp.data(), zp.data(),
+                    &np_gather,
+                    thrust::raw_pointer_cast(xp.data()),
+                    thrust::raw_pointer_cast(yp.data()),
+                    thrust::raw_pointer_cast(zp.data()),
                     Exp.data(),Eyp.data(),Ezp.data(),
                     Bxp.data(),Byp.data(),Bzp.data(),
                     ixyzmin_grid,
@@ -1161,7 +1170,10 @@ PhysicalParticleContainer::Evolve (int lev,
                     
                     long ncrse = np - nfine_gather;
                     warpx_geteb_energy_conserving(
-                        &ncrse, xp.data()+nfine_gather, yp.data()+nfine_gather, zp.data()+nfine_gather,
+                        &ncrse,
+                        thrust::raw_pointer_cast(xp.data()+nfine_gather),
+                        thrust::raw_pointer_cast(yp.data()+nfine_gather),
+                        thrust::raw_pointer_cast(zp.data()+nfine_gather),
                         Exp.data()+nfine_gather, Eyp.data()+nfine_gather, Ezp.data()+nfine_gather,
                         Bxp.data()+nfine_gather, Byp.data()+nfine_gather, Bzp.data()+nfine_gather,
                         cixyzmin_grid,
@@ -1226,9 +1238,13 @@ PhysicalParticleContainer::Evolve (int lev,
                         jx_ptr, &ngJ, jxntot,
                         jy_ptr, &ngJ, jyntot,
                         jz_ptr, &ngJ, jzntot,
-                        &np_current, xp.data(), yp.data(), zp.data(),
+                        &np_current,
+                        thrust::raw_pointer_cast(xp.data()),
+                        thrust::raw_pointer_cast(yp.data()),
+                        thrust::raw_pointer_cast(zp.data()),
                         uxp.data(), uyp.data(), uzp.data(),
-                        giv.data(), wp.data(), &this->charge,
+                        thrust::raw_pointer_cast(giv.data()),
+                        wp.data(), &this->charge,
                         &xyzmin[0], &xyzmin[1], &xyzmin[2],
                         &dt, &dx[0], &dx[1], &dx[2],
                         &WarpX::nox,&WarpX::noy,&WarpX::noz,
@@ -1296,9 +1312,15 @@ PhysicalParticleContainer::Evolve (int lev,
                         jx_ptr, &ngJ, jxntot,
                         jy_ptr, &ngJ, jyntot,
                         jz_ptr, &ngJ, jzntot,
-                        &ncrse, xp.data()+nfine_current, yp.data()+nfine_current, zp.data()+nfine_current,
-                        uxp.data()+nfine_current, uyp.data()+nfine_current, uzp.data()+nfine_current,
-                        giv.data()+nfine_current, wp.data()+nfine_current, &this->charge,
+                        &ncrse,
+                        thrust::raw_pointer_cast(xp.data()+nfine_current),
+                        thrust::raw_pointer_cast(yp.data()+nfine_current),
+                        thrust::raw_pointer_cast(zp.data()+nfine_current),
+                        uxp.data()+nfine_current,
+                        uyp.data()+nfine_current,
+                        uzp.data()+nfine_current,
+                        thrust::raw_pointer_cast(giv.data()+nfine_current),
+                        wp.data()+nfine_current, &this->charge,
                         &cxyzmin_tile[0], &cxyzmin_tile[1], &cxyzmin_tile[2],
                         &dt, &cdx[0], &cdx[1], &cdx[2],
                         &WarpX::nox,&WarpX::noy,&WarpX::noz,
@@ -1355,8 +1377,10 @@ PhysicalParticleContainer::Evolve (int lev,
 
 void
 PhysicalParticleContainer::PushPX(WarpXParIter& pti,
-	                          RealVector& xp, RealVector& yp, RealVector& zp,
-                                  RealVector& giv,
+	                          thrust::device_vector<Real>& xp,
+                                  thrust::device_vector<Real>& yp,
+                                  thrust::device_vector<Real>& zp,
+                                  thrust::device_vector<Real>& giv,
                                   Real dt)
 {
 
@@ -1388,8 +1412,12 @@ PhysicalParticleContainer::PushPX(WarpXParIter& pti,
 
 #endif
 
-    warpx_particle_pusher(&np, xp.data(), yp.data(), zp.data(),
-                          uxp.data(), uyp.data(), uzp.data(), giv.data(),
+    warpx_particle_pusher(&np,
+                          thrust::raw_pointer_cast(xp.data()),
+                          thrust::raw_pointer_cast(yp.data()),
+                          thrust::raw_pointer_cast(zp.data()),
+                          uxp.data(), uyp.data(), uzp.data(),
+                          thrust::raw_pointer_cast(giv.data()),
                           Exp.dataPtr(), Eyp.dataPtr(), Ezp.dataPtr(),
                           Bxp.dataPtr(), Byp.dataPtr(), Bzp.dataPtr(),
                           &this->charge, &this->mass, &dt,
@@ -1410,7 +1438,7 @@ PhysicalParticleContainer::PushP (int lev, Real dt,
 #pragma omp parallel
 #endif
     {
-	RealVector xp, yp, zp, giv;
+	thrust::device_vector<Real> xp, yp, zp, giv;
 
         for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
 	{
@@ -1459,7 +1487,10 @@ PhysicalParticleContainer::PushP (int lev, Real dt,
             const int l_lower_order_in_v = true;
             long lvect_fieldgathe = 64;
             warpx_geteb_energy_conserving(
-                &np, xp.data(), yp.data(), zp.data(),
+                &np,
+                thrust::raw_pointer_cast(xp.data()),
+                thrust::raw_pointer_cast(yp.data()),
+                thrust::raw_pointer_cast(zp.data()),
                 Exp.data(),Eyp.data(),Ezp.data(),
                 Bxp.data(),Byp.data(),Bzp.data(),
                 ixyzmin_grid,
@@ -1475,8 +1506,12 @@ PhysicalParticleContainer::PushP (int lev, Real dt,
                 &ll4symtry, &l_lower_order_in_v,
                 &lvect_fieldgathe, &WarpX::field_gathering_algo);
 
-            warpx_particle_pusher_momenta(&np, xp.data(), yp.data(), zp.data(),
-                                          uxp.data(), uyp.data(), uzp.data(), giv.data(),
+            warpx_particle_pusher_momenta(&np,
+                                          thrust::raw_pointer_cast(xp.data()),
+                                          thrust::raw_pointer_cast(yp.data()),
+                                          thrust::raw_pointer_cast(zp.data()),
+                                          uxp.data(), uyp.data(), uzp.data(),
+                                          thrust::raw_pointer_cast(giv.data()),
                                           Exp.dataPtr(), Eyp.dataPtr(), Ezp.dataPtr(),
                                           Bxp.dataPtr(), Byp.dataPtr(), Bzp.dataPtr(),
                                           &this->charge, &this->mass, &dt,
diff --git a/Source/RigidInjectedParticleContainer.H b/Source/RigidInjectedParticleContainer.H
index ea405180e..cd53db71f 100644
--- a/Source/RigidInjectedParticleContainer.H
+++ b/Source/RigidInjectedParticleContainer.H
@@ -43,10 +43,10 @@ public:
                          amrex::Real dt) override;
 
     virtual void PushPX(WarpXParIter& pti,
-	                RealVector& xp,
-                        RealVector& yp,
-                        RealVector& zp,
-                        RealVector& giv,
+	                thrust::device_vector<amrex::Real>& xp,
+                        thrust::device_vector<amrex::Real>& yp,
+                        thrust::device_vector<amrex::Real>& zp,
+                        thrust::device_vector<amrex::Real>& giv,
                         amrex::Real dt) override;
 
     virtual void PushP (int lev, amrex::Real dt,
diff --git a/Source/RigidInjectedParticleContainer.cpp b/Source/RigidInjectedParticleContainer.cpp
index 7b6b48326..1a85b13ef 100644
--- a/Source/RigidInjectedParticleContainer.cpp
+++ b/Source/RigidInjectedParticleContainer.cpp
@@ -73,7 +73,7 @@ RigidInjectedParticleContainer::RemapParticles()
                 // Note that the particles are already in the boosted frame.
                 // This value is saved to advance the particles not injected yet
 
-                RealVector xp, yp, zp;
+                thrust::device_vector<Real> xp, yp, zp;
 
                 for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
                 {
@@ -134,7 +134,7 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
 #pragma omp parallel
 #endif
     {
-        RealVector xp, yp, zp;
+        thrust::device_vector<Real> xp, yp, zp;
 
         for (WarpXParIter pti(*this, 0); pti.isValid(); ++pti)
         {
@@ -205,8 +205,10 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
 
 void
 RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
-	                               RealVector& xp, RealVector& yp, RealVector& zp,
-                                       RealVector& giv,
+	                               thrust::device_vector<Real>& xp,
+                                       thrust::device_vector<Real>& yp,
+                                       thrust::device_vector<Real>& zp,
+                                       thrust::device_vector<Real>& giv,
                                        Real dt)
 {
 
@@ -239,7 +241,7 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
 #endif
 
     // Save the position and momenta, making copies
-    RealVector xp_save, yp_save, zp_save, uxp_save, uyp_save, uzp_save;
+    thrust::device_vector<Real> xp_save, yp_save, zp_save, uxp_save, uyp_save, uzp_save;
 
     if (!done_injecting_lev) {
         xp_save = xp;
@@ -265,8 +267,12 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
         }
     }
 
-    warpx_particle_pusher(&np, xp.data(), yp.data(), zp.data(),
-                          uxp.data(), uyp.data(), uzp.data(), giv.data(),
+    warpx_particle_pusher(&np,
+                          thrust::raw_pointer_cast(xp.data()),
+                          thrust::raw_pointer_cast(yp.data()),
+                          thrust::raw_pointer_cast(zp.data()),
+                          uxp.data(), uyp.data(), uzp.data(),
+                          thrust::raw_pointer_cast(giv.data()),
                           Exp.dataPtr(), Eyp.dataPtr(), Ezp.dataPtr(),
                           Bxp.dataPtr(), Byp.dataPtr(), Bzp.dataPtr(),
                           &this->charge, &this->mass, &dt,
@@ -355,7 +361,7 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
 #pragma omp parallel
 #endif
     {
-        RealVector xp, yp, zp, giv;
+        thrust::device_vector<Real> xp, yp, zp, giv;
 
         for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
         {
@@ -404,7 +410,10 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
             const int l_lower_order_in_v = true;
             long lvect_fieldgathe = 64;
             warpx_geteb_energy_conserving(
-                &np, xp.data(), yp.data(), zp.data(),
+                &np,
+                thrust::raw_pointer_cast(xp.data()),
+                thrust::raw_pointer_cast(yp.data()),
+                thrust::raw_pointer_cast(zp.data()),
                 Exp.data(),Eyp.data(),Ezp.data(),
                 Bxp.data(),Byp.data(),Bzp.data(),
                 ixyzmin_grid,
@@ -425,8 +434,12 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
             auto uyp_save = uyp;
             auto uzp_save = uzp;
 
-            warpx_particle_pusher_momenta(&np, xp.data(), yp.data(), zp.data(),
-                                          uxp.data(), uyp.data(), uzp.data(), giv.data(),
+            warpx_particle_pusher_momenta(&np,
+                                          thrust::raw_pointer_cast(xp.data()),
+                                          thrust::raw_pointer_cast(yp.data()),
+                                          thrust::raw_pointer_cast(zp.data()),
+                                          uxp.data(), uyp.data(), uzp.data(),
+                                          thrust::raw_pointer_cast(giv.data()),
                                           Exp.dataPtr(), Eyp.dataPtr(), Ezp.dataPtr(),
                                           Bxp.dataPtr(), Byp.dataPtr(), Bzp.dataPtr(),
                                           &this->charge, &this->mass, &dt,
diff --git a/Source/WarpXParticleContainer.H b/Source/WarpXParticleContainer.H
index a806f6211..b1bab9aef 100644
--- a/Source/WarpXParticleContainer.H
+++ b/Source/WarpXParticleContainer.H
@@ -36,12 +36,12 @@ public:
     WarpXParIter (ContainerType& pc, int level);
 
 #if (AMREX_SPACEDIM == 2)
-    void GetPosition (RealVector& x,
-                      RealVector& y,
-                      RealVector& z) const;
-    void SetPosition (const RealVector& x,
-                      const RealVector& y,
-                      const RealVector& z);
+    void GetPosition (thrust::device_vector<Real>& x,
+                      thrust::device_vector<Real>& y,
+                      thrust::device_vector<Real>& z) const;
+    void SetPosition (const thrust::device_vector<Real>& x,
+                      const thrust::device_vector<Real>& y,
+                      const thrust::device_vector<Real>& z);
 #endif
 
     const std::array<RealVector, PIdx::nattribs>& GetAttribs () const { 
diff --git a/Source/WarpXParticleContainer.cpp b/Source/WarpXParticleContainer.cpp
index 3526ba9c8..252f10691 100644
--- a/Source/WarpXParticleContainer.cpp
+++ b/Source/WarpXParticleContainer.cpp
@@ -250,7 +250,7 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local)
 #pragma omp parallel
 #endif
     {
-        RealVector xp, yp, zp;
+        thrust::device_vector<Real> xp, yp, zp;
         FArrayBox local_rho;
 
         for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
@@ -300,7 +300,10 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local)
             long lvect = 8;
             
             warpx_charge_deposition(data_ptr,
-                                    &np, xp.data(), yp.data(), zp.data(), wp.data(),
+                                    &np,
+                                    thrust::raw_pointer_cast(xp.data()),
+                                    thrust::raw_pointer_cast(yp.data()),
+                                    thrust::raw_pointer_cast(zp.data()), wp.data(),
                                     &this->charge, &xyzmin[0], &xyzmin[1], &xyzmin[2], 
                                     &dx[0], &dx[1], &dx[2], &nx, &ny, &nz,
                                     &nxg, &nyg, &nzg, &WarpX::nox,&WarpX::noy,&WarpX::noz,
@@ -471,7 +474,7 @@ WarpXParticleContainer::PushX (int lev, Real dt)
 #pragma omp parallel
 #endif
     {
-        RealVector xp, yp, zp, giv;
+        thrust::device_vector<Real> xp, yp, zp, giv;
 
         for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
         {
@@ -498,8 +501,12 @@ WarpXParticleContainer::PushX (int lev, Real dt)
             // Particle Push
             //
             BL_PROFILE_VAR_START(blp_pxr_pp);
-            warpx_particle_pusher_positions(&np, xp.data(), yp.data(), zp.data(),
-                                            uxp.data(), uyp.data(), uzp.data(), giv.data(), &dt);
+            warpx_particle_pusher_positions(&np,
+                                            thrust::raw_pointer_cast(xp.data()),
+                                            thrust::raw_pointer_cast(yp.data()),
+                                            thrust::raw_pointer_cast(zp.data()),
+                                            uxp.data(), uyp.data(), uzp.data(),
+                                            thrust::raw_pointer_cast(giv.data()), &dt);
             BL_PROFILE_VAR_STOP(blp_pxr_pp);
             
             //