Skip to content

Commit

Permalink
Update Particle Container to Pure SoA
Browse files Browse the repository at this point in the history
Transition particle containers to pure SoA layouts.
  • Loading branch information
ax3l committed Jan 13, 2024
1 parent 82f7f1e commit c8b0111
Show file tree
Hide file tree
Showing 44 changed files with 505 additions and 542 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
3 changes: 0 additions & 3 deletions Source/Diagnostics/FlushFormats/FlushFormatAscent.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
2 changes: 1 addition & 1 deletion Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
6 changes: 3 additions & 3 deletions Source/Diagnostics/ReducedDiags/FieldProbe.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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;
Expand Down
20 changes: 16 additions & 4 deletions Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.H
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;

Expand All @@ -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);
Expand Down
36 changes: 13 additions & 23 deletions Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down Expand Up @@ -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
Expand All @@ -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);
Expand Down
3 changes: 1 addition & 2 deletions Source/Diagnostics/ReducedDiags/LoadBalanceCosts.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
4 changes: 2 additions & 2 deletions Source/Diagnostics/WarpXOpenPMD.H
Original file line number Diff line number Diff line change
Expand Up @@ -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;}
Expand Down Expand Up @@ -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
*
Expand Down
Loading

0 comments on commit c8b0111

Please sign in to comment.