Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.

Author: ax3l

cmake/dependencies/ABLASTR.cmake

@@ -166,10 +166,10 @@ set(ImpactX_openpmd_src ""
     "Local path to openPMD-api source directory (preferred if set)")
 
 # Git fetcher
-set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git"
+set(ImpactX_ablastr_repo "https://github.com/ax3l/WarpX.git"
     CACHE STRING
     "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
-set(ImpactX_ablastr_branch "f71597fca2cf9d3700e7e11533f4a85a1846a2b8"
+set(ImpactX_ablastr_branch "topic-particle-soa"
     CACHE STRING
     "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")
 

src/particles/ImpactXParticleContainer.H

@@ -17,6 +17,7 @@
 #include <AMReX_MultiFab.H>
 #include <AMReX_ParIter.H>
 #include <AMReX_Particles.H>
+#include <AMReX_ParticleTile.H>
 
 #include <AMReX_IntVect.H>
 #include <AMReX_Vector.H>
@@ -35,6 +36,7 @@ namespace impactx
      * because we change the meaning of these "positions" depending on the
      * coordinate system we are currently in.
      */
+    /*
     struct RealAoS
     {
         enum
@@ -52,6 +54,7 @@ namespace impactx
         static_assert(names_t.size() == nattribs);
         static_assert(names_s.size() == nattribs);
     };
+    */
 
     /** This struct indexes the additional Real attributes
      *  stored in an SoA in ImpactXParticleContainer
@@ -60,18 +63,21 @@ namespace impactx
     {
         enum
         {
+            x,   ///< position in x [m] (at fixed t OR fixed s)
+            y,   ///< position in y [m] (at fixed t OR fixed s)
+            z,   ///< position in z [m] (at fixed t) OR time-of-flight ct [m] (at fixed s)
             ux,  ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed t or s)
             uy,  ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed t or s)
             pt,  ///< momentum in z, scaled by the magnitude of the reference momentum [unitless] (at fixed t) OR energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s)
             m_qm, ///< charge to mass ratio, in q_e/m_e (q_e/eV) TODO: rename to qm_m
-            w,   ///< particle weight, unitless
+            w,    ///< particle weight, unitless
             nattribs ///< the number of attributes above (always last)
         };
 
         //! named labels for fixed t
-        static constexpr auto names_t = { "momentum_x", "momentum_y", "momentum_z", "qmm", "weighting" };
+        static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qmm", "weighting" };
         //! named labels for fixed s
-        static constexpr auto names_s = { "momentum_x", "momentum_y", "momentum_t", "qmm", "weighting" };
+        static constexpr auto names_s = { "position_x", "position_y", "position_ct", "momentum_x", "momentum_y", "momentum_t", "qmm", "weighting" };
         static_assert(names_t.size() == nattribs);
         static_assert(names_s.size() == nattribs);
     };
@@ -83,55 +89,64 @@ namespace impactx
     {
         enum
         {
-            nattribs ///< the number of particles above (always last)
+            id,
+            cpu,
+            nattribs ///< the number of attributes above (always last)
         };
+
+        //! named labels for fixed t
+        static constexpr auto names_t = { "id", "cpu" };
+        //! named labels for fixed s
+        static constexpr auto names_s = { "id", "cpu" };
+        static_assert(names_t.size() == nattribs);
+        static_assert(names_s.size() == nattribs);
     };
 
     /** AMReX iterator for particle boxes
      *
      * We subclass here to change the default threading strategy, which is
      * `static` in AMReX, to `dynamic` in ImpactX.
      */
-    class ParIter
-        : public amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+    class ParIterSoA
+        : public amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
-        using amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParIter;
+        using amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>::ParIterSoA;
 
-        ParIter (ContainerType& pc, int level);
+        ParIterSoA (ContainerType& pc, int level);
 
-        ParIter (ContainerType& pc, int level, amrex::MFItInfo& info);
+        ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
     };
 
     /** Const AMReX iterator for particle boxes - data is read only.
      *
      * We subclass here to change the default threading strategy, which is
      * `static` in AMReX, to `dynamic` in ImpactX.
      */
-    class ParConstIter
-        : public amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+    class ParConstIterSoA
+        : public amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
-        using amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParConstIter;
+        using amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>::ParConstIterSoA;
 
-        ParConstIter (ContainerType& pc, int level);
+        ParConstIterSoA (ContainerType& pc, int level);
 
-        ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info);
+        ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
     };
 
     /** Beam Particles in ImpactX
      *
      * This class stores particles, distributed over MPI ranks.
      */
     class ImpactXParticleContainer
-        : public amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+        : public amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
         //! amrex iterator for particle boxes
-        using iterator = impactx::ParIter;
+        using iterator = impactx::ParIterSoA;
 
         //! amrex constant iterator for particle boxes (read-only)
-        using const_iterator = impactx::ParConstIter;
+        using const_iterator = impactx::ParConstIterSoA;
 
         //! Construct a new particle container
         ImpactXParticleContainer (amrex::AmrCore* amr_core);

src/particles/ImpactXParticleContainer.cpp

@@ -17,7 +17,6 @@
 #include <AMReX_AmrParGDB.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
-#include <AMReX_ParticleTile.H>
 
 #include <stdexcept>
 
@@ -31,24 +30,24 @@ namespace impactx
         return do_dynamic;
     }
 
-    ParIter::ParIter (ContainerType& pc, int level)
-        : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParIterSoA::ParIterSoA (ContainerType& pc, int level)
+        : amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
                    amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {}
 
-    ParIter::ParIter (ContainerType& pc, int level, amrex::MFItInfo& info)
-        : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParIterSoA::ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info)
+        : amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               info.SetDynamic(do_omp_dynamic())) {}
 
-    ParConstIter::ParConstIter (ContainerType& pc, int level)
-        : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level)
+        : amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {}
 
-    ParConstIter::ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info)
-        : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info)
+        : amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               info.SetDynamic(do_omp_dynamic())) {}
 
     ImpactXParticleContainer::ImpactXParticleContainer (amrex::AmrCore* amr_core)
-        : amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
+        : amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
     {
         SetParticleSize();
     }
@@ -105,35 +104,26 @@ namespace impactx
         reserveData();
         resizeData();
 
+        // write Real attributes (SoA) to particle initialized zero
         auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0);
 
         /* Create a temporary tile to obtain data from simulation. This data
          * is then copied to the permanent tile which is stored on the particle
          * (particle_tile).
          */
-        using PinnedTile = amrex::ParticleTile<
-            amrex::Particle<NStructReal, NStructInt>,
-            NArrayReal, NArrayInt,
-            amrex::PinnedArenaAllocator
-        >;
+        using PinnedTile = typename ContainerLike<amrex::PinnedArenaAllocator>::ParticleTileType;
         PinnedTile pinned_tile;
         pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps());
 
         for (int i = 0; i < np; i++)
         {
-            ParticleType p;
-            p.id() = ParticleType::NextID();
-            p.cpu() = amrex::ParallelDescriptor::MyProc();
-            p.pos(0) = x[i];
-            p.pos(1) = y[i];
-            p.pos(2) = z[i];
-            // write position, creating cpu id, and particle id
-            pinned_tile.push_back(p);
+            pinned_tile.push_back_int(IntSoA::id, ParticleType::NextID());
+            pinned_tile.push_back_int(IntSoA::cpu, amrex::ParallelDescriptor::MyProc());
         }
 
-        // write Real attributes (SoA) to particle initialized zero
-        DefineAndReturnParticleTile(0, 0, 0);
-
+        pinned_tile.push_back_real(RealSoA::x, x);
+        pinned_tile.push_back_real(RealSoA::y, y);
+        pinned_tile.push_back_real(RealSoA::z, z);
         pinned_tile.push_back_real(RealSoA::ux, px);
         pinned_tile.push_back_real(RealSoA::uy, py);
         pinned_tile.push_back_real(RealSoA::pt, pz);
@@ -145,6 +135,7 @@ namespace impactx
          */
         auto old_np = particle_tile.numParticles();
         auto new_np = old_np + pinned_tile.numParticles();
+        amrex::AllPrint() << "old_np=" << old_np << " new_np=" << new_np << "\n";
         particle_tile.resize(new_np);
         amrex::copyParticles(
                 particle_tile, pinned_tile, 0, old_np, pinned_tile.numParticles());

src/particles/diagnostics/DiagnosticOutput.cpp

@@ -17,6 +17,7 @@
 #include <AMReX_ParmParse.H>  // for ParmParse
 #include <AMReX_REAL.H>       // for ParticleReal
 #include <AMReX_Print.H>      // for PrintToFile
+#include <AMReX_ParticleTile.H>     // for constructor of SoAParticle
 
 
 namespace impactx::diagnostics
@@ -46,7 +47,8 @@ namespace impactx::diagnostics
         }
 
         // create a host-side particle buffer
-        auto tmp = pc.make_alike<amrex::PinnedArenaAllocator>();
+        using ParticleType = amrex::SoAParticle<RealSoA::nattribs,IntSoA::nattribs>;
+        auto tmp = pc.template make_alike<amrex::PinnedArenaAllocator>();
 
         // copy device-to-host
         bool const local = true;
@@ -56,34 +58,34 @@ namespace impactx::diagnostics
         int const nLevel = tmp.finestLevel();
         for (int lev = 0; lev <= nLevel; ++lev) {
             // loop over all particle boxes
-            using ParIt = typename decltype(tmp)::ParConstIterType;
-            for (ParIt pti(tmp, lev); pti.isValid(); ++pti) {
+            using MyPinnedParIter = amrex::ParIterSoA<RealSoA::nattribs,IntSoA::nattribs, amrex::PinnedArenaAllocator>;
+            for (MyPinnedParIter pti(tmp, lev); pti.isValid(); ++pti) {
                 const int np = pti.numParticles();
 
                 // preparing access to particle data: AoS
-                using PType = ImpactXParticleContainer::ParticleType;
-                auto const &aos = pti.GetArrayOfStructs();
-                PType const *const AMREX_RESTRICT aos_ptr = aos().dataPtr();
+                //using PType = ImpactXParticleContainer::ParticleType;
+                //auto const &aos = pti.GetArrayOfStructs();
+                //PType const *const AMREX_RESTRICT aos_ptr = aos().dataPtr();
 
                 // preparing access to particle data: SoA of Reals
-                auto &soa_real = pti.GetStructOfArrays().GetRealData();
-                amrex::ParticleReal const *const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr();
-                amrex::ParticleReal const *const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr();
-                amrex::ParticleReal const *const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr();
+                auto const& particle_attributes = pti.GetStructOfArrays();
+                auto const& part_px = particle_attributes.GetRealData(RealSoA::ux);
+                auto const& part_py = particle_attributes.GetRealData(RealSoA::uy);
+                auto const& part_pt = particle_attributes.GetRealData(RealSoA::pt);
+
+                // Preparing the constructor for the new SoA Particle Type
+                auto ptd = pti.GetParticleTile().getParticleTileData();
 
                 if (otype == OutputType::PrintParticles) {
                     // print out particles (this hack works only on CPU and on GPUs with
                     // unified memory access)
                     for (int i = 0; i < np; ++i) {
-
-                        // access AoS data such as positions and cpu/id
-                        PType const &p = aos_ptr[i];
+                        ParticleType p(ptd, i);
                         amrex::ParticleReal const x = p.pos(0);
                         amrex::ParticleReal const y = p.pos(1);
                         amrex::ParticleReal const t = p.pos(2);
                         uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu());
 
-                        // access SoA Real data
                         amrex::ParticleReal const px = part_px[i];
                         amrex::ParticleReal const py = part_py[i];
                         amrex::ParticleReal const pt = part_pt[i];
@@ -119,14 +121,11 @@ namespace impactx::diagnostics
                     // print out particles (this hack works only on CPU and on GPUs with
                     // unified memory access)
                     for (int i = 0; i < np; ++i) {
-
-                        // access AoS data such as positions and cpu/id
-                        PType const &p = aos_ptr[i];
+                        ParticleType p(ptd, i);
                         amrex::ParticleReal const x = p.pos(0);
                         amrex::ParticleReal const y = p.pos(1);
                         uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu());
 
-                        // access SoA Real data
                         amrex::ParticleReal const px = part_px[i];
                         amrex::ParticleReal const py = part_py[i];