Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.

cmake/dependencies/ABLASTR.cmake

@@ -178,7 +178,7 @@ set(ImpactX_openpmd_src ""
 set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git"
  CACHE STRING
  "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
-set(ImpactX_ablastr_branch "24.01"
+set(ImpactX_ablastr_branch "94ae11900131846e9f8b3704194673f7f02d8959"
  CACHE STRING
  "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")
 

examples/fodo/run_fodo_programmable.py

@@ -78,16 +78,16 @@ def my_drift(pge, pti, refpart):
 
  else:
  array = np.array
- # access AoS data such as positions and cpu/id
- aos = pti.aos()
- aos_arr = array(aos, copy=False)
 
- # access SoA data such as momentum
+ # access particle attributes
  soa = pti.soa()
  real_arrays = soa.GetRealData()
- px = array(real_arrays[0], copy=False)
- py = array(real_arrays[1], copy=False)
- pt = array(real_arrays[2], copy=False)
+ x = array(real_arrays[0], copy=False)
+ y = array(real_arrays[1], copy=False)
+ t = array(real_arrays[2], copy=False)
+ px = array(real_arrays[3], copy=False)
+ py = array(real_arrays[4], copy=False)
+ pt = array(real_arrays[5], copy=False)
 
  # length of the current slice
  slice_ds = pge.ds / pge.nslice
@@ -97,9 +97,9 @@ def my_drift(pge, pti, refpart):
  betgam2 = pt_ref**2 - 1.0
 
  # advance position and momentum (drift)
- aos_arr[:]["x"] += slice_ds * px[:]
- aos_arr[:]["y"] += slice_ds * py[:]
- aos_arr[:]["z"] += (slice_ds / betgam2) * pt[:]
+ x[:] += slice_ds * px[:]
+ y[:] += slice_ds * py[:]
+ t[:] += (slice_ds / betgam2) * pt[:]
 
 
 def my_ref_drift(pge, refpart):

src/particles/CollectLost.cpp

@@ -27,9 +27,9 @@ namespace impactx
  using DstData = ImpactXParticleContainer::ParticleTileType::ParticleTileDataType;
 
  AMREX_GPU_HOST_DEVICE
- void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept {
-  dst.m_aos[dst_ip] = src.m_aos[src_ip];
-
+ void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept
+ {
+ dst.m_idcpu[dst_ip] = src.m_idcpu[src_ip];
  for (int j = 0; j < SrcData::NAR; ++j)
  dst.m_rdata[j][dst_ip] = src.m_rdata[j][src_ip];
  for (int j = 0; j < src.m_num_runtime_real; ++j)
@@ -141,8 +141,7 @@ namespace impactx
  // move down
  int const new_index = ip - n_removed;
 
- ptile_src_data.m_aos[new_index] = ptile_src_data.m_aos[ip];
-
+ ptile_src_data.m_idcpu[new_index] = ptile_src_data.m_idcpu[ip];
  for (int j = 0; j < SrcData::NAR; ++j)
  ptile_src_data.m_rdata[j][new_index] = ptile_src_data.m_rdata[j][ip];
  for (int j = 0; j < ptile_src_data.m_num_runtime_real; ++j)

src/particles/ImpactXParticleContainer.H

@@ -18,6 +18,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,43 +36,16 @@ namespace impactx
  t ///< fixed t as the independent variable
  };
 
- /** AMReX pre-defined Real attributes
- *
- * These are the AMReX pre-defined struct indexes for the Real attributes
- * stored in an AoS in ImpactXParticleContainer. We document this here,
- * because we change the meaning of these "positions" depending on the
- * coordinate system we are currently in.
- */
- struct RealAoS
- {
- enum
- {
- x, ///< position in x [m] (at fixed s OR fixed t)
- y, ///< position in y [m] (at fixed s OR fixed t)
- t, ///< c * time-of-flight [m] (at fixed s)
- nattribs ///< the number of attributes above (always last)
- };
-
- // at fixed t, the third component represents the position z
- enum {
- z = t ///< position in z [m] (at fixed t)
- };
-
- //! named labels for fixed s
- static constexpr auto names_s = { "position_x", "position_y", "position_t" };
- //! named labels for fixed t
- static constexpr auto names_t = { "position_x", "position_y", "position_z" };
- static_assert(names_s.size() == nattribs);
- static_assert(names_t.size() == nattribs);
- };
-
- /** This struct indexes the additional Real attributes
+ /** This struct indexes the Real attributes
  * stored in an SoA in ImpactXParticleContainer
  */
  struct RealSoA
  {
  enum
  {
+ x, ///< position in x [m] (at fixed s or t)
+ y, ///< position in y [m] (at fixed s or t)
+ t, ///< time-of-flight ct [m] (at fixed s)
  px, ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t)
  py, ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t)
  pt, ///< energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s)
@@ -80,27 +54,28 @@ namespace impactx
  nattribs ///< the number of attributes above (always last)
  };
 
- // at fixed t, the third component represents the momentum in z
+ // at fixed t, the third component represents the position z, the 6th component represents the momentum in z
  enum {
+ z = t, ///< position in z [m] (at fixed t)
  pz = pt ///< momentum in z, scaled by the magnitude of the reference momentum [unitless] (at fixed t)
  };
 
  //! named labels for fixed s
- static constexpr auto names_s = { "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" };
+ static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" };
  //! named labels for fixed t
- static constexpr auto names_t = { "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" };
+ static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" };
  static_assert(names_s.size() == nattribs);
  static_assert(names_t.size() == nattribs);
  };
 
- /** This struct indexes the additional Integer attributes
+ /** This struct indexes the Integer attributes
  * stored in an SoA in ImpactXParticleContainer
  */
  struct IntSoA
  {
  enum
  {
- nattribs ///< the number of particles above (always last)
+ nattribs ///< the number of attributes above (always last)
  };
  };
 
@@ -109,46 +84,46 @@ namespace impactx
  * 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 (initialization::AmrCoreData* amr_core);
@@ -276,10 +251,6 @@ namespace impactx
  DepositCharge (std::unordered_map<int, amrex::MultiFab> & rho,
  amrex::Vector<amrex::IntVect> const & ref_ratio);
 
- /** Get the name of each Real AoS component */
- std::vector<std::string>
- RealAoS_names () const;
-
  /** Get the name of each Real SoA component */
  std::vector<std::string>
  RealSoA_names () const;
@@ -311,10 +282,6 @@ namespace impactx
 
  }; // ImpactXParticleContainer
 
- /** Get the name of each Real AoS component */
- std::vector<std::string>
- get_RealAoS_names ();
-
  /** Get the name of each Real SoA component
  *
  * @param num_real_comps number of compile-time + runtime arrays

src/particles/ImpactXParticleContainer.cpp

@@ -19,7 +19,7 @@
 #include <AMReX_AmrParGDB.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
-#include <AMReX_ParticleTile.H>
+#include <AMReX_Particle.H>
 
 #include <algorithm>
 #include <stdexcept>
@@ -38,24 +38,24 @@ namespace
 
 namespace impactx
 {
- 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 (initialization::AmrCoreData* amr_core)
- : amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
+ : amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
  {
  SetParticleSize();
  }
@@ -157,14 +157,18 @@ namespace impactx
 
  const int cpuid = amrex::ParallelDescriptor::MyProc();
 
- auto * AMREX_RESTRICT pstructs = particle_tile.GetArrayOfStructs()().dataPtr();
  auto & soa = particle_tile.GetStructOfArrays().GetRealData();
+ amrex::ParticleReal * const AMREX_RESTRICT x_arr = soa[RealSoA::x].dataPtr();
+ amrex::ParticleReal * const AMREX_RESTRICT y_arr = soa[RealSoA::y].dataPtr();
+ amrex::ParticleReal * const AMREX_RESTRICT t_arr = soa[RealSoA::t].dataPtr();
  amrex::ParticleReal * const AMREX_RESTRICT px_arr = soa[RealSoA::px].dataPtr();
  amrex::ParticleReal * const AMREX_RESTRICT py_arr = soa[RealSoA::py].dataPtr();
  amrex::ParticleReal * const AMREX_RESTRICT pt_arr = soa[RealSoA::pt].dataPtr();
  amrex::ParticleReal * const AMREX_RESTRICT qm_arr = soa[RealSoA::qm].dataPtr();
  amrex::ParticleReal * const AMREX_RESTRICT w_arr = soa[RealSoA::w ].dataPtr();
 
+ uint64_t * const AMREX_RESTRICT idcpu_arr = particle_tile.GetStructOfArrays().GetIdCPUData().dataPtr();
+
  amrex::ParticleReal const * const AMREX_RESTRICT x_ptr = x.data();
  amrex::ParticleReal const * const AMREX_RESTRICT y_ptr = y.data();
  amrex::ParticleReal const * const AMREX_RESTRICT t_ptr = t.data();
@@ -175,12 +179,12 @@ namespace impactx
  amrex::ParallelFor(np,
  [=] AMREX_GPU_DEVICE (int i) noexcept
  {
- ParticleType& p = pstructs[old_np + i];
- p.id() = pid + i;
- p.cpu() = cpuid;
- p.pos(RealAoS::x) = x_ptr[i];
- p.pos(RealAoS::y) = y_ptr[i];
- p.pos(RealAoS::t) = t_ptr[i];
+ amrex::ParticleIDWrapper{idcpu_arr[old_np+i]} = pid + i;
+ amrex::ParticleCPUWrapper{idcpu_arr[old_np+i]} = cpuid;
+
+ x_arr[old_np+i] = x_ptr[i];
+ y_arr[old_np+i] = y_ptr[i];
+ t_arr[old_np+i] = t_ptr[i];
 
  px_arr[old_np+i] = px_ptr[i];
  py_arr[old_np+i] = py_ptr[i];
@@ -240,12 +244,6 @@ namespace impactx
  >(*this);
  }
 
- std::vector<std::string>
- ImpactXParticleContainer::RealAoS_names () const
- {
- return get_RealAoS_names();
- }
-
  std::vector<std::string>
  ImpactXParticleContainer::RealSoA_names () const
  {
@@ -264,17 +262,6 @@ namespace impactx
  m_coordsystem = coord_system;
  }
 
- std::vector<std::string>
- get_RealAoS_names ()
- {
- std::vector<std::string> real_aos_names(RealAoS::names_s.size());
-
- // compile-time attributes
- std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin());
-
- return real_aos_names;
- }
-
  std::vector<std::string>
  get_RealSoA_names (int num_real_comps)
  {