Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.
ax3l · Jan 24, 2024 · ede6114 · ede6114
1 parent 47e8ace
commit ede6114
Show file tree

Hide file tree

Showing 44 changed files with 507 additions and 543 deletions.
diff --git a/Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.H b/Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.H
@@ -143,19 +143,19 @@ struct LorentzTransformParticles
         const amrex::ParticleReal uzp = uz_old_p * weight_old
                                       + uz_new_p * weight_new;
 #if defined (WARPX_DIM_3D)
-        dst.m_aos[i_dst].pos(0) = xp;
-        dst.m_aos[i_dst].pos(1) = yp;
-        dst.m_aos[i_dst].pos(2) = zp;
+        dst.m_rdata[PIdx::x][i_dst] = xp;
+        dst.m_rdata[PIdx::y][i_dst] = yp;
+        dst.m_rdata[PIdx::z][i_dst] = zp;
 #elif defined (WARPX_DIM_RZ)
-        dst.m_aos[i_dst].pos(0) = std::sqrt(xp*xp + yp*yp);
-        dst.m_aos[i_dst].pos(1) = zp;
+        dst.m_rdata[PIdx::x][i_dst] = std::sqrt(xp*xp + yp*yp);
+        dst.m_rdata[PIdx::z][i_dst] = zp;
         dst.m_rdata[PIdx::theta][i_dst] = std::atan2(yp, xp);
 #elif defined (WARPX_DIM_XZ)
-        dst.m_aos[i_dst].pos(0) = xp;
-        dst.m_aos[i_dst].pos(1) = zp;
+        dst.m_rdata[PIdx::x][i_dst] = xp;
+        dst.m_rdata[PIdx::z][i_dst] = zp;
         amrex::ignore_unused(yp);
 #elif defined (WARPX_DIM_1D_Z)
-        dst.m_aos[i_dst].pos(0) = zp;
+        dst.m_rdata[PIdx::z][i_dst] = zp;
         amrex::ignore_unused(xp, yp);
 #else
         amrex::ignore_unused(xp, yp, zp);

diff --git a/Source/Diagnostics/FlushFormats/FlushFormatAscent.cpp b/Source/Diagnostics/FlushFormats/FlushFormatAscent.cpp
@@ -94,9 +94,6 @@ FlushFormatAscent::WriteParticles(const amrex::Vector<ParticleDiag>& particle_di
         // get names of real comps
         std::map<std::string, int> real_comps_map = pc->getParticleComps();
 
-        // WarpXParticleContainer compile-time extra AoS attributes (Real): 0
-        // WarpXParticleContainer compile-time extra AoS attributes (int): 0
-
         // WarpXParticleContainer compile-time extra SoA attributes (Real): PIdx::nattribs
         // not an efficient search, but N is small...
         for(int j = 0; j < PIdx::nattribs; ++j)

diff --git a/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp b/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp
@@ -178,8 +178,8 @@ FlushFormatCheckpoint::CheckpointParticles (
         Vector<std::string> real_names;
         Vector<std::string> int_names;
 
+        // note: positions skipped here, since we reconstruct a plotfile SoA from them
         real_names.push_back("weight");
-
         real_names.push_back("momentum_x");
         real_names.push_back("momentum_y");
         real_names.push_back("momentum_z");

diff --git a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
@@ -355,8 +355,8 @@ FlushFormatPlotfile::WriteParticles(const std::string& dir,
         Vector<int> int_flags;
         Vector<int> real_flags;
 
+        // note: positions skipped here, since we reconstruct a plotfile SoA from them
         real_names.push_back("weight");
-
         real_names.push_back("momentum_x");
         real_names.push_back("momentum_y");
         real_names.push_back("momentum_z");
@@ -366,13 +366,14 @@ FlushFormatPlotfile::WriteParticles(const std::string& dir,
 #endif
 
         // get the names of the real comps
-        real_names.resize(tmp.NumRealComps());
+        real_names.resize(tmp.NumRealComps()-AMREX_SPACEDIM);  // skip positions for pure SoA plotfiles
         auto runtime_rnames = tmp.getParticleRuntimeComps();
         for (auto const& x : runtime_rnames) { real_names[x.second+PIdx::nattribs] = x.first; }
 
         // plot any "extra" fields by default
         real_flags = part_diag.m_plot_flags;
         real_flags.resize(tmp.NumRealComps(), 1);
+        real_flags.erase(real_flags.begin(), real_flags.begin()+AMREX_SPACEDIM);  // skip positions for pure SoA plotfiles
 
         // and the names
         int_names.resize(tmp.NumIntComps());

diff --git a/Source/Diagnostics/ReducedDiags/FieldProbe.cpp b/Source/Diagnostics/ReducedDiags/FieldProbe.cpp
@@ -431,8 +431,6 @@ void FieldProbe::ComputeDiags (int step)
         {
             const auto getPosition = GetParticlePosition<FieldProbePIdx>(pti);
             auto setPosition = SetParticlePosition<FieldProbePIdx>(pti);
-            const auto& aos = pti.GetArrayOfStructs();
-            const auto* AMREX_RESTRICT m_structs = aos().dataPtr();
 
             auto const np = pti.numParticles();
             if (update_particles_moving_window)
@@ -482,6 +480,8 @@ void FieldProbe::ComputeDiags (int step)
                 ParticleReal* const AMREX_RESTRICT part_Bz = attribs[FieldProbePIdx::Bz].dataPtr();
                 ParticleReal* const AMREX_RESTRICT part_S = attribs[FieldProbePIdx::S].dataPtr();
 
+                auto * const AMREX_RESTRICT idcpu = pti.GetStructOfArrays().GetIdCPUData().data();
+
                 const auto &xyzmin = WarpX::LowerCorner(box, lev, 0._rt);
                 const std::array<Real, 3> &dx = WarpX::CellSize(lev);
 
@@ -556,7 +556,7 @@ void FieldProbe::ComputeDiags (int step)
                         amrex::ParticleReal xp, yp, zp;
                         getPosition(ip, xp, yp, zp);
                         long idx = ip*noutputs;
-                        dvp[idx++] = m_structs[ip].id();
+                        dvp[idx++] = amrex::ParticleIDWrapper{idcpu[ip]};  // all particles created on IO cpu
                         dvp[idx++] = xp;
                         dvp[idx++] = yp;
                         dvp[idx++] = zp;

diff --git a/Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.H b/Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.H
@@ -24,7 +24,14 @@ struct FieldProbePIdx
 {
     enum
     {
-        Ex = 0, Ey, Ez,
+#if !defined (WARPX_DIM_1D_Z)
+        x,
+#endif
+#if defined (WARPX_DIM_3D)
+        y,
+#endif
+        z,
+        Ex, Ey, Ez,
         Bx, By, Bz,
         S, //!< the Poynting vector
 #ifdef WARPX_DIM_RZ
@@ -40,9 +47,14 @@ struct FieldProbePIdx
  * nattribs tells the particle container to allot 7 SOA values.
  */
 class FieldProbeParticleContainer
-    : public amrex::ParticleContainer<0, 0, FieldProbePIdx::nattribs>
+    : public amrex::ParticleContainerPureSoA<FieldProbePIdx::nattribs, 0>
 {
 public:
+    static constexpr int NStructReal = 0;
+    static constexpr int NStructInt = 0;
+    static constexpr int NReal = FieldProbePIdx::nattribs;
+    static constexpr int NInt = 0;
+
     FieldProbeParticleContainer (amrex::AmrCore* amr_core);
     ~FieldProbeParticleContainer() override = default;
 
@@ -52,9 +64,9 @@ public:
     FieldProbeParticleContainer& operator= ( FieldProbeParticleContainer&& )       = default;
 
     //! amrex iterator for our number of attributes
-    using iterator = amrex::ParIter<0, 0, FieldProbePIdx::nattribs, 0>;
+    using iterator = amrex::ParIterSoA<FieldProbePIdx::nattribs, 0>;
     //! amrex iterator for our number of attributes (read-only)
-    using const_iterator = amrex::ParConstIter<0, 0, FieldProbePIdx::nattribs, 0>;
+    using const_iterator = amrex::ParConstIterSoA<FieldProbePIdx::nattribs, 0>;
 
     //! similar to WarpXParticleContainer::AddNParticles but does not include u(x,y,z)
     void AddNParticles (int lev, amrex::Vector<amrex::ParticleReal> const & x, amrex::Vector<amrex::ParticleReal> const & y, amrex::Vector<amrex::ParticleReal> const & z);

diff --git a/Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.cpp b/Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.cpp
@@ -59,7 +59,7 @@
 using namespace amrex;
 
 FieldProbeParticleContainer::FieldProbeParticleContainer (AmrCore* amr_core)
-    : ParticleContainer<0, 0, FieldProbePIdx::nattribs>(amr_core->GetParGDB())
+    : ParticleContainerPureSoA<FieldProbePIdx::nattribs, 0>(amr_core->GetParGDB())
 {
     SetParticleSize();
 }
@@ -89,33 +89,17 @@ FieldProbeParticleContainer::AddNParticles (int lev,
      * is then coppied to the permament tile which is stored on the particle
      * (particle_tile).
      */
+    using PinnedTile = typename ContainerLike<amrex::PinnedArenaAllocator>::ParticleTileType;
 
-    using PinnedTile = ParticleTile<amrex::Particle<NStructReal, NStructInt>,
-                                    NArrayReal, NArrayInt,
-                                    amrex::PinnedArenaAllocator>;
     PinnedTile pinned_tile;
     pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps());
 
     for (int i = 0; i < np; i++)
     {
-        ParticleType p;
-        p.id() = ParticleType::NextID();
-        p.cpu() = ParallelDescriptor::MyProc();
-#if defined(WARPX_DIM_3D)
-        p.pos(0) = x[i];
-        p.pos(1) = y[i];
-        p.pos(2) = z[i];
-#elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
-        amrex::ignore_unused(y);
-        p.pos(0) = x[i];
-        p.pos(1) = z[i];
-#elif defined(WARPX_DIM_1D_Z)
-        amrex::ignore_unused(x, y);
-        p.pos(0) = z[i];
-#endif
-
-        // write position, cpu id, and particle id to particle
-        pinned_tile.push_back(p);
+        auto & idcpu_data = pinned_tile.GetStructOfArrays().GetIdCPUData();
+        idcpu_data.push_back(0);
+        amrex::ParticleIDWrapper{idcpu_data.back()} = ParticleType::NextID();
+        amrex::ParticleCPUWrapper(idcpu_data.back()) = ParallelDescriptor::MyProc();
     }
 
     // write Real attributes (SoA) to particle initialized zero
@@ -125,7 +109,13 @@ FieldProbeParticleContainer::AddNParticles (int lev,
 #ifdef WARPX_DIM_RZ
     pinned_tile.push_back_real(FieldProbePIdx::theta, np, 0.0);
 #endif
-
+#if !defined (WARPX_DIM_1D_Z)
+    pinned_tile.push_back_real(FieldProbePIdx::x, x);
+#endif
+#if defined (WARPX_DIM_3D)
+    pinned_tile.push_back_real(FieldProbePIdx::y, y);
+#endif
+    pinned_tile.push_back_real(FieldProbePIdx::z, z);
     pinned_tile.push_back_real(FieldProbePIdx::Ex, np, 0.0);
     pinned_tile.push_back_real(FieldProbePIdx::Ey, np, 0.0);
     pinned_tile.push_back_real(FieldProbePIdx::Ez, np, 0.0);

diff --git a/Source/Diagnostics/ReducedDiags/LoadBalanceCosts.cpp b/Source/Diagnostics/ReducedDiags/LoadBalanceCosts.cpp
@@ -56,8 +56,7 @@ namespace
             auto const & plev  = pc.GetParticles(lev);
 
             auto const & ptile = plev.at(box_index);
-            auto const & aos   = ptile.GetArrayOfStructs();
-            auto const np = aos.numParticles();
+            auto const np = ptile.numParticles();
             num_macro_particles += np;
         }
 

diff --git a/Source/Diagnostics/WarpXOpenPMD.H b/Source/Diagnostics/WarpXOpenPMD.H
@@ -41,7 +41,7 @@ class WarpXParticleCounter
 {
 public:
   using ParticleContainer = typename WarpXParticleContainer::ContainerLike<amrex::PinnedArenaAllocator>;
-  using ParticleIter = typename amrex::ParIter<0, 0, PIdx::nattribs, 0, amrex::PinnedArenaAllocator>;
+  using ParticleIter = typename amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>;
 
   WarpXParticleCounter (ParticleContainer* pc);
   [[nodiscard]] unsigned long GetTotalNumParticles () const {return m_Total;}
@@ -77,7 +77,7 @@ class WarpXOpenPMDPlot
 {
 public:
   using ParticleContainer = typename WarpXParticleContainer::ContainerLike<amrex::PinnedArenaAllocator>;
-  using ParticleIter = typename amrex::ParConstIter<0, 0, PIdx::nattribs, 0, amrex::PinnedArenaAllocator>;
+  using ParticleIter = typename amrex::ParConstIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>;
 
   /** Initialize openPMD I/O routines
    *