Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.
ax3l · Dec 22, 2023 · c94c0a5 · c94c0a5
1 parent 9ce7082
commit c94c0a5
Show file tree

Hide file tree

Showing 37 changed files with 430 additions and 506 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/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,8 @@ 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++] = idcpu[ip];
+                        dvp[idx++] = 0;
                         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,8 @@ struct FieldProbePIdx
 {
     enum
     {
-        Ex = 0, Ey, Ez,
+        x, y, z,
+        Ex, Ey, Ez,
         Bx, By, Bz,
         S, //!< the Poynting vector
 #ifdef WARPX_DIM_RZ
@@ -40,9 +41,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 +58,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
@@ -126,6 +110,9 @@ FieldProbeParticleContainer::AddNParticles (int lev,
     pinned_tile.push_back_real(FieldProbePIdx::theta, np, 0.0);
 #endif
 
+    pinned_tile.push_back_real(FieldProbePIdx::x, x);
+    pinned_tile.push_back_real(FieldProbePIdx::y, y);
+    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);
   unsigned long GetTotalNumParticles () {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
    *

diff --git a/Source/Diagnostics/WarpXOpenPMD.cpp b/Source/Diagnostics/WarpXOpenPMD.cpp
@@ -18,11 +18,9 @@
 #include "WarpX.H"
 #include "OpenPMDHelpFunction.H"
 
-#include <ablastr/particles/IndexHandling.H>
 #include <ablastr/warn_manager/WarnManager.H>
 
 #include <AMReX.H>
-#include <AMReX_ArrayOfStructs.H>
 #include <AMReX_BLassert.H>
 #include <AMReX_Box.H>
 #include <AMReX_Config.H>
@@ -722,77 +720,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
 
             contributed_particles = true;
 
-            // get position and particle ID from aos
-            // note: this implementation iterates the AoS 4x...
-            // if we flush late as we do now, we can also copy out the data in one go
-            const auto &aos = pti.GetArrayOfStructs();  // size =  numParticlesOnTile
-            {
-                // Save positions
-#if defined(WARPX_DIM_RZ)
-                {
-                   const std::shared_ptr<amrex::ParticleReal> z(
-                           new amrex::ParticleReal[numParticleOnTile],
-                           [](amrex::ParticleReal const *p) { delete[] p; }
-                   );
-                   for (auto i = 0; i < numParticleOnTile; i++) {
-                       z.get()[i] = aos[i].pos(1);  // {0: "r", 1: "z"}
-                   }
-                   std::string const positionComponent = "z";
-                   currSpecies["position"]["z"].storeChunk(z, {offset}, {numParticleOnTile64});
-                }
-
-                //   reconstruct x and y from polar coordinates r, theta
-                auto const& soa = pti.GetStructOfArrays();
-                amrex::ParticleReal const* theta = soa.GetRealData(PIdx::theta).dataPtr();
-                WARPX_ALWAYS_ASSERT_WITH_MESSAGE(theta != nullptr, "openPMD: invalid theta pointer.");
-                WARPX_ALWAYS_ASSERT_WITH_MESSAGE(int(soa.GetRealData(PIdx::theta).size()) == numParticleOnTile,
-                                                 "openPMD: theta and tile size do not match");
-                {
-                    const std::shared_ptr< amrex::ParticleReal > x(
-                            new amrex::ParticleReal[numParticleOnTile],
-                            [](amrex::ParticleReal const *p){ delete[] p; }
-                    );
-                    const std::shared_ptr< amrex::ParticleReal > y(
-                            new amrex::ParticleReal[numParticleOnTile],
-                            [](amrex::ParticleReal const *p){ delete[] p; }
-                    );
-                    for (auto i=0; i<numParticleOnTile; i++) {
-                        auto const r = aos[i].pos(0);  // {0: "r", 1: "z"}
-                        x.get()[i] = r * std::cos(theta[i]);
-                        y.get()[i] = r * std::sin(theta[i]);
-                    }
-                    currSpecies["position"]["x"].storeChunk(x, {offset}, {numParticleOnTile64});
-                    currSpecies["position"]["y"].storeChunk(y, {offset}, {numParticleOnTile64});
-                }
-#else
-                auto const positionComponents = detail::getParticlePositionComponentLabels();
-                for (auto currDim = 0; currDim < AMREX_SPACEDIM; currDim++) {
-                    const std::shared_ptr<amrex::ParticleReal> curr(
-                            new amrex::ParticleReal[numParticleOnTile],
-                            [](amrex::ParticleReal const *p) { delete[] p; }
-                    );
-                    for (auto i = 0; i < numParticleOnTile; i++) {
-                        curr.get()[i] = aos[i].pos(currDim);
-                    }
-                    std::string const positionComponent = positionComponents[currDim];
-                    currSpecies["position"][positionComponent].storeChunk(curr, {offset},
-                                                                          {numParticleOnTile64});
-                }
-#endif
-
-                // save particle ID after converting it to a globally unique ID
-                const std::shared_ptr<uint64_t> ids(
-                        new uint64_t[numParticleOnTile],
-                        [](uint64_t const *p) { delete[] p; }
-                );
-                for (auto i = 0; i < numParticleOnTile; i++) {
-                    ids.get()[i] = ablastr::particles::localIDtoGlobal(static_cast<int>(aos[i].id()), static_cast<int>(aos[i].cpu()));
-                }
-                const auto *const scalar = openPMD::RecordComponent::SCALAR;
-                currSpecies["id"][scalar].storeChunk(ids, {offset}, {numParticleOnTile64});
-
-            }
-            //  save "extra" particle properties in AoS and SoA
+            //  save particle properties
             SaveRealProperty(pti,
                              currSpecies,
                              offset,
@@ -949,49 +877,30 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
 
 {
   auto const numParticleOnTile = pti.numParticles();
-  auto const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
-  auto const& aos = pti.GetArrayOfStructs();  // size =  numParticlesOnTile
+  uint64_t const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
   auto const& soa = pti.GetStructOfArrays();
-  // first we concatenate the AoS into contiguous arrays
-  {
-    // note: WarpX does not yet use extra AoS Real attributes
-    for( auto idx=0; idx<ParticleIter::ContainerType::NStructReal; idx++ ) {  // lgtm [cpp/constant-comparison]
-        if( write_real_comp[idx] ) {
-            // handle scalar and non-scalar records by name
-            const auto [record_name, component_name] = detail::name2openPMD(real_comp_names[idx]);
-            auto currRecord = currSpecies[record_name];
-            auto currRecordComp = currRecord[component_name];
-
-            const std::shared_ptr< amrex::ParticleReal > d(
-                new amrex::ParticleReal[numParticleOnTile],
-                [](amrex::ParticleReal const *p){ delete[] p; }
-            );
-
-            for( auto kk=0; kk<numParticleOnTile; kk++ ) {
-                d.get()[kk] = aos[kk].rdata(idx);
-            }
-
-            currRecordComp.storeChunk(d,
-                {offset}, {numParticleOnTile64});
-        }
-    }
-  }
 
   auto const getComponentRecord = [&currSpecies](std::string const comp_name) {
     // handle scalar and non-scalar records by name
     const auto [record_name, component_name] = detail::name2openPMD(comp_name);
     return currSpecies[record_name][component_name];
   };
 
+    // here we the save the SoA properties (idcpu)
+    {
+        // todo: add support to not write the particle index
+        getComponentRecord("id").storeChunkRaw(
+        soa.GetIdCPUData().data(), {offset}, {numParticleOnTile64});
+    }
+
   // here we the save the SoA properties (real)
   {
     auto const real_counter = std::min(write_real_comp.size(), real_comp_names.size());
     for (auto idx=0; idx<real_counter; idx++) {
-        auto ii = ParticleIter::ContainerType::NStructReal + idx;  // jump over extra AoS names
-        if (write_real_comp[ii]) {
-            getComponentRecord(real_comp_names[ii]).storeChunkRaw(
-                soa.GetRealData(idx).data(), {offset}, {numParticleOnTile64});
-        }
+      if (write_real_comp[idx]) {
+        getComponentRecord(real_comp_names[idx]).storeChunkRaw(
+          soa.GetRealData(idx).data(), {offset}, {numParticleOnTile64});
+      }
     }
   }
   // and now SoA int properties

diff --git a/Source/EmbeddedBoundary/ParticleBoundaryProcess.H b/Source/EmbeddedBoundary/ParticleBoundaryProcess.H
@@ -25,12 +25,11 @@ struct NoOp {
 struct Absorb {
     template <typename PData>
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-    void operator() (const PData& ptd, int i,
+    void operator() (PData& ptd, int i,
                      const amrex::RealVect& /*pos*/, const amrex::RealVect& /*normal*/,
                      amrex::RandomEngine const& /*engine*/) const noexcept
     {
-        auto& p = ptd.m_aos[i];
-        p.id() = -p.id();
+        ptd.id(i) = -ptd.id(i);
     }
 };
 }

diff --git a/Source/EmbeddedBoundary/ParticleScraper.H b/Source/EmbeddedBoundary/ParticleScraper.H
@@ -170,13 +170,12 @@ scrapeParticles (PC& pc, const amrex::Vector<const amrex::MultiFab*>& distance_t
             auto& tile = pti.GetParticleTile();
             auto ptd = tile.getParticleTileData();
             const auto np = tile.numParticles();
-            amrex::Particle<0,0> * const particles = tile.GetArrayOfStructs()().data();
             auto phi = (*distance_to_eb[lev])[pti].array();  // signed distance function
             amrex::ParallelForRNG( np,
             [=] AMREX_GPU_DEVICE (const int ip, amrex::RandomEngine const& engine) noexcept
             {
                 // skip particles that are already flagged for removal
-                if (particles[ip].id() < 0) return;
+                if (ptd.id(ip) < 0) return;
 
                 amrex::ParticleReal xp, yp, zp;
                 getPosition(ip, xp, yp, zp);