Skip to content

Commit

Permalink
First change to impactx
Browse files Browse the repository at this point in the history 
Adapting to SoA structures

Link the CMake file to Axel Warpx branch

Avoid branch conflicts

Correction of ParIter and ConstParIter

Adding corrections to make_alike function call

Update Diagnostics

Correction of the ParIter loop

Still adaptating ImpactX to SoAParticle

Update Init

Note:
  AMReX-Codes/amrex#2878 (comment)

Adapted the PushSingleParticle function to SoAParticle
  • Loading branch information
@Thierry992 @ax3l
Thierry992 authored and ax3l committed Apr 4, 2023
1 parent 3dbe247 commit 56a47f8
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 82 deletions.
8 changes: 4 additions & 4 deletions cmake/dependencies/ABLASTR.cmake
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,18 +131,18 @@ set(ImpactX_ablastr_src ""
"Local path to ABLASTR 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/Thierry992/WarpX.git"
CACHE STRING
"Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
set(ImpactX_ablastr_branch "22.10"
set(ImpactX_ablastr_branch "soa-particle"
CACHE STRING
"Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")

# AMReX is transitively pulled through ABLASTR
set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
set(ImpactX_amrex_repo "https://github.com/Thierry992/amrex.git"
CACHE STRING
"Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
set(ImpactX_amrex_branch "ed1ecd62acb3fd7d39b8a23aa4e9ad09669741bb"
set(ImpactX_amrex_branch "particle_soa_refactor"
CACHE STRING
"Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")

Expand Down
53 changes: 30 additions & 23 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,15 +36,16 @@ namespace impactx
* 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 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)
};
};

//struct RealAoS
//{
// 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)
// };
//};

/** This struct indexes the additional Real attributes
* stored in an SoA in ImpactXParticleContainer
Expand All @@ -52,6 +54,9 @@ namespace impactx
{
enum
{
x,
y,
z,
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)
Expand All @@ -68,7 +73,9 @@ namespace impactx
{
enum
{
nattribs ///< the number of particles above (always last)
id,
cpu,
nattribs ///< the number of attributes above (always last)
};
};

Expand All @@ -77,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 (amrex::AmrCore* amr_core);
Expand Down
40 changes: 17 additions & 23 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,32 +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<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 @@ -142,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
22 changes: 14 additions & 8 deletions src/particles/Push.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ namespace detail
struct PushSingleParticle
{
using PType = ImpactXParticleContainer::ParticleType;
using ParticleTileType = amrex::ParticleTile<amrex::SoAParticle<RealSoA::nattribs, IntSoA::nattribs>,RealSoA::nattribs, IntSoA::nattribs>;
using ParticleTileDataType = amrex::ParticleTileData<ParticleTileType::StorageParticleType,RealSoA::nattribs, IntSoA::nattribs>;

/** Constructor taking in pointers to particle data
*
Expand All @@ -46,12 +48,12 @@ namespace detail
* @param ref_part the struct containing the reference particle
*/
PushSingleParticle (T_Element element,
PType* AMREX_RESTRICT aos_ptr,
ParticleTileDataType ptd,
amrex::ParticleReal* AMREX_RESTRICT part_px,
amrex::ParticleReal* AMREX_RESTRICT part_py,
amrex::ParticleReal* AMREX_RESTRICT part_pt,
RefPart ref_part)
: m_element(element), m_aos_ptr(aos_ptr),
: m_element(element), m_ptd(ptd),
m_part_px(part_px), m_part_py(part_py), m_part_pt(part_pt),
m_ref_part(ref_part)
{
Expand All @@ -70,8 +72,8 @@ namespace detail
void
operator() (long i) const
{
// access AoS data such as positions and cpu/id
PType& AMREX_RESTRICT p = m_aos_ptr[i];

PType p(m_ptd,i);

// access SoA Real data
amrex::ParticleReal & AMREX_RESTRICT px = m_part_px[i];
Expand All @@ -85,7 +87,7 @@ namespace detail

private:
T_Element const m_element;
PType* const AMREX_RESTRICT m_aos_ptr;
ParticleTileDataType const m_ptd;
amrex::ParticleReal* const AMREX_RESTRICT m_part_px;
amrex::ParticleReal* const AMREX_RESTRICT m_part_py;
amrex::ParticleReal* const AMREX_RESTRICT m_part_pt;
Expand Down Expand Up @@ -129,8 +131,12 @@ namespace detail

// preparing access to particle data: AoS
using PType = ImpactXParticleContainer::ParticleType;
auto& aos = pti.GetArrayOfStructs();
PType* AMREX_RESTRICT aos_ptr = aos().dataPtr();
using ParticleTile = amrex::ParticleTile<amrex::SoAParticle<RealSoA::nattribs, IntSoA::nattribs>,RealSoA::nattribs, IntSoA::nattribs>;
using ParticleTileDataType = amrex::ParticleTileData<ParticleTile::StorageParticleType,RealSoA::nattribs, IntSoA::nattribs>;

ParticleTileDataType tile_data = pti.GetParticleTile().getParticleTileData();
//auto& aos = pti.GetArrayOfStructs();
//PType* AMREX_RESTRICT aos_ptr = aos().dataPtr();

// preparing access to particle data: SoA of Reals
auto& soa_real = pti.GetStructOfArrays().GetRealData();
Expand All @@ -143,7 +149,7 @@ namespace detail
[=, &ref_part](auto element) {
// push beam particles relative to reference particle
detail::PushSingleParticle<decltype(element)> const pushSingleParticle(
element, aos_ptr, part_px, part_py, part_pt, ref_part);
element, tile_data, part_px, part_py, part_pt, ref_part);
// loop over beam particles in the box
amrex::ParallelFor(np, pushSingleParticle);

Expand Down
35 changes: 17 additions & 18 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,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];
Expand Down Expand Up @@ -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];

Expand Down
Loading

0 comments on commit 56a47f8

Please sign in to comment.