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
ax3l committed Apr 19, 2023
1 parent cdbac8b commit c4268f4
Show file tree
Hide file tree
Showing 8 changed files with 127 additions and 143 deletions.
4 changes: 2 additions & 2 deletions cmake/dependencies/ABLASTR.cmake
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)")

Expand Down
49 changes: 32 additions & 17 deletions src/particles/ImpactXParticleContainer.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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);
};
Expand All @@ -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);
Expand Down
43 changes: 17 additions & 26 deletions src/particles/ImpactXParticleContainer.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
#include <AMReX_AmrParGDB.H>
#include <AMReX_ParallelDescriptor.H>
#include <AMReX_ParmParse.H>
#include <AMReX_ParticleTile.H>

#include <stdexcept>

Expand All @@ -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();
}
Expand Down Expand Up @@ -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);
Expand All @@ -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());
Expand Down
34 changes: 14 additions & 20 deletions src/particles/diagnostics/DiagnosticOutput.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand All @@ -56,34 +58,29 @@ 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();

// 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];
Expand Down Expand Up @@ -119,14 +116,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];

Expand Down
Loading

0 comments on commit c4268f4

Please sign in to comment.