Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Start SoA in WarpX #3417

Closed
wants to merge 1 commit into from
Closed

[WIP] Start SoA in WarpX #3417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ 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:
FieldProbeParticleContainer (amrex::AmrCore* amr_core);
Expand All @@ -46,7 +46,7 @@ public:
//! amrex iterator for our number of attributes
using iterator = amrex::ParIter<0, 0, 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
4 changes: 2 additions & 2 deletions Source/Diagnostics/ReducedDiags/FieldProbeParticleContainer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,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 @@ -91,7 +91,7 @@ FieldProbeParticleContainer::AddNParticles (int lev,
* (particle_tile).
*/

using PinnedTile = ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
using PinnedTile = ParticleTile<FieldProbeParticleContainer, NArrayReal, NArrayInt,
amrex::PinnedArenaAllocator>;
PinnedTile pinned_tile;
pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps());
Expand Down
25 changes: 14 additions & 11 deletions Source/Particles/NamedComponentParticleContainer.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,10 @@
struct PIdx
{
enum {
w = 0, ///< weight
x = 0,
y,
z,
w, ///< weight
ux, uy, uz,
#ifdef WARPX_DIM_RZ
theta, ///< RZ needs all three position components
Expand All @@ -43,11 +46,11 @@ struct PIdx
*/
template <template<class> class T_Allocator=amrex::DefaultAllocator>
class NamedComponentParticleContainer :
public amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>
public amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator>
{
public:
/** Construct an empty NamedComponentParticleContainer **/
NamedComponentParticleContainer () : amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>() {}
NamedComponentParticleContainer () : amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator>() {}

/** Construct a NamedComponentParticleContainer from an AmrParGDB object
*
Expand All @@ -59,7 +62,7 @@ public:
* AMR hierarchy. Usually, this is generated by an AmrCore or AmrLevel object.
*/
NamedComponentParticleContainer (amrex::AmrParGDB* amr_pgdb)
: amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>(amr_pgdb) {
: amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator>(amr_pgdb) {
// build up the map of string names to particle component numbers
particle_comps["w"] = PIdx::w;
particle_comps["ux"] = PIdx::ux;
Expand All @@ -83,12 +86,12 @@ public:
* @param p_ricomps name-to-index map for run-time integer components
*/
NamedComponentParticleContainer(
amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator> && pc,
amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator> && pc,
std::map<std::string, int> p_comps,
std::map<std::string, int> p_icomps,
std::map<std::string, int> p_rcomps,
std::map<std::string, int> p_ricomps)
: amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>(std::move(pc)),
: amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator>(std::move(pc)),
particle_comps(p_comps),
particle_icomps(p_icomps),
particle_runtime_comps(p_rcomps),
Expand All @@ -111,7 +114,7 @@ public:
NamedComponentParticleContainer<NewAllocator>
make_alike () const {
auto tmp = NamedComponentParticleContainer<NewAllocator>(
amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::template make_alike<NewAllocator>(),
amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, T_Allocator>::template make_alike<NewAllocator>(),
particle_comps,
particle_icomps,
particle_runtime_comps,
Expand All @@ -120,10 +123,10 @@ public:
return tmp;
}

using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::NumRealComps;
using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::NumIntComps;
using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::AddRealComp;
using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::AddIntComp;
using amrex::ParticleContainerPureSoA<PIdx::nattribs,0,T_Allocator>::NumRealComps;
using amrex::ParticleContainerPureSoA<PIdx::nattribs,0,T_Allocator>::NumIntComps;
using amrex::ParticleContainerPureSoA<PIdx::nattribs,0,T_Allocator>::AddRealComp;
using amrex::ParticleContainerPureSoA<PIdx::nattribs,0,T_Allocator>::AddIntComp;

/** Allocate a new run-time real component
*
Expand Down
120 changes: 76 additions & 44 deletions Source/Particles/Pusher/GetAndSetPosition.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,10 +51,19 @@ void get_particle_position (const WarpXParticleContainer::SuperParticleType& p,
*/
struct GetParticlePosition
{
using PType = WarpXParticleContainer::ParticleType;
using RType = amrex::ParticleReal;

const PType* AMREX_RESTRICT m_structs = nullptr;
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
RType* AMREX_RESTRICT m_x = nullptr;
RType* AMREX_RESTRICT m_z = nullptr;
#elif defined(WARPX_DIM_3D)
const RType* AMREX_RESTRICT m_x = nullptr;
const RType* AMREX_RESTRICT m_y = nullptr;
const RType* AMREX_RESTRICT m_z = nullptr;
#elif defined(WARPX_DIM_1D_Z)
RType* AMREX_RESTRICT m_z = nullptr;;
#endif

#if defined(WARPX_DIM_RZ)
const RType* m_theta = nullptr;
#elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
Expand All @@ -76,10 +85,19 @@ struct GetParticlePosition
template <typename ptiType>
GetParticlePosition (const ptiType& a_pti, int a_offset = 0) noexcept
{
const auto& aos = a_pti.GetArrayOfStructs();
m_structs = aos().dataPtr() + a_offset;
#if defined(WARPX_DIM_RZ)
const auto& soa = a_pti.GetStructOfArrays();

#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
m_x = soa.GetRealData(0).dataPtr() + a_offset;
m_z = soa.GetRealData(1).dataPtr() + a_offset;
#elif defined(WARPX_DIM_3D)
m_x = soa.GetRealData(0).dataPtr() + a_offset;
m_y = soa.GetRealData(1).dataPtr() + a_offset;
m_z = soa.GetRealData(2).dataPtr() + a_offset;
#elif defined(WARPX_DIM_1D_Z)
m_z = soa.GetRealData(0).dataPtr() + a_offset;
#endif
#if defined(WARPX_DIM_RZ)
m_theta = soa.GetRealData(PIdx::theta).dataPtr() + a_offset;
#endif
}
Expand All @@ -90,24 +108,23 @@ struct GetParticlePosition
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (const int i, RType& x, RType& y, RType& z) const noexcept
{
const PType& p = m_structs[i];
#ifdef WARPX_DIM_RZ
RType r = p.pos(0);
RType r = m_x[i];
x = r*std::cos(m_theta[i]);
y = r*std::sin(m_theta[i]);
z = p.pos(1);
z = m_z[i];
#elif WARPX_DIM_3D
x = p.pos(0);
y = p.pos(1);
z = p.pos(2);
x = m_x[i];
y = m_y[i];
z = m_z[i];
#elif WARPX_DIM_XZ
x = p.pos(0);
x = m_x[i];
y = m_y_default;
z = p.pos(1);
z = m_z[i];
#else
x = m_x_default;
y = m_y_default;
z = p.pos(0);
z = m_z[i];
#endif
}

Expand All @@ -119,23 +136,22 @@ struct GetParticlePosition
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void AsStored (const int i, RType& x, RType& y, RType& z) const noexcept
{
const PType& p = m_structs[i];
#ifdef WARPX_DIM_RZ
x = p.pos(0);
x = m_x[i];
y = m_theta[i];
z = p.pos(1);
z = m_z[i];
#elif WARPX_DIM_3D
x = p.pos(0);
y = p.pos(1);
z = p.pos(2);
x = m_x[i];
y = m_y[i];
z = m_z[i];
#elif WARPX_DIM_XZ
x = p.pos(0);
x = m_x[i];
y = m_y_default;
z = p.pos(1);
z = m_z[i];
#else
x = m_x_default;
y = m_y_default;
z = p.pos(0);
z = m_z[i];
#endif
}
};
Expand All @@ -148,21 +164,37 @@ struct GetParticlePosition
*/
struct SetParticlePosition
{
using PType = WarpXParticleContainer::ParticleType;
using RType = amrex::ParticleReal;

PType* AMREX_RESTRICT m_structs;
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
RType* AMREX_RESTRICT m_x;
RType* AMREX_RESTRICT m_z;
#elif defined(WARPX_DIM_3D)
RType* AMREX_RESTRICT m_x;
RType* AMREX_RESTRICT m_y;
RType* AMREX_RESTRICT m_z;
#elif defined(WARPX_DIM_1D_Z)
RType* AMREX_RESTRICT m_z;
#endif
#if defined(WARPX_DIM_RZ)
RType* AMREX_RESTRICT m_theta;
#endif

template <typename ptiType>
SetParticlePosition (const ptiType& a_pti, int a_offset = 0) noexcept
{
auto& aos = a_pti.GetArrayOfStructs();
m_structs = aos().dataPtr() + a_offset;
#if defined(WARPX_DIM_RZ)
auto& soa = a_pti.GetStructOfArrays();
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
m_x = soa.GetRealData(0).dataPtr() + a_offset;
m_z = soa.GetRealData(1).dataPtr() + a_offset;
#elif defined(WARPX_DIM_3D)
m_x = soa.GetRealData(0).dataPtr() + a_offset;
m_y = soa.GetRealData(1).dataPtr() + a_offset;
m_z = soa.GetRealData(2).dataPtr() + a_offset;
#elif defined(WARPX_DIM_1D_Z)
m_z = soa.GetRealData(0).dataPtr() + a_offset;
#endif
#if defined(WARPX_DIM_RZ)
m_theta = soa.GetRealData(PIdx::theta).dataPtr() + a_offset;
#endif
}
Expand All @@ -180,17 +212,17 @@ struct SetParticlePosition
#endif
#ifdef WARPX_DIM_RZ
m_theta[i] = std::atan2(y, x);
m_structs[i].pos(0) = std::sqrt(x*x + y*y);
m_structs[i].pos(1) = z;
m_x[i] = std::sqrt(x*x + y*y);
m_z[i] = z;
#elif WARPX_DIM_3D
m_structs[i].pos(0) = x;
m_structs[i].pos(1) = y;
m_structs[i].pos(2) = z;
m_x[i] = x;
m_y[i] = y;
m_z[i] = z;
#elif WARPX_DIM_XZ
m_structs[i].pos(0) = x;
m_structs[i].pos(1) = z;
m_x[i] = x;
m_z[i] = z;
#else
m_structs[i].pos(0) = z;
m_z[i] = z;
#endif
}

Expand All @@ -208,18 +240,18 @@ struct SetParticlePosition
amrex::ignore_unused(x,y);
#endif
#ifdef WARPX_DIM_RZ
m_structs[i].pos(0) = x;
m_x[i] = x;
m_theta[i] = y;
m_structs[i].pos(1) = z;
m_z[i] = z;
#elif WARPX_DIM_3D
m_structs[i].pos(0) = x;
m_structs[i].pos(1) = y;
m_structs[i].pos(2) = z;
m_x[i] = x;
m_y[i] = y;
m_z[i] = z;
#elif WARPX_DIM_XZ
m_structs[i].pos(0) = x;
m_structs[i].pos(1) = z;
m_x[i] = x;
m_z[i] = z;
#else
m_structs[i].pos(0) = z;
m_z[i] = z;
#endif
}
};
Expand Down
6 changes: 3 additions & 3 deletions Source/Particles/WarpXParticleContainer.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ public:
* particle container classes (that store a collection of particles) derive. Derived
* classes can be used for plasma particles, photon particles, or non-physical
* particles (e.g., for the laser antenna).
* It derives from amrex::ParticleContainer<0,0,PIdx::nattribs>, where the
* It derives from amrex::ParticleContainerPureSoA<PIdx::nattribs>, where the
* template arguments stand for the number of int and amrex::Real SoA and AoS
* data in amrex::Particle.
* - AoS amrex::Real: x, y, z (default), 0 additional (first template
Expand Down Expand Up @@ -144,8 +144,8 @@ public:
* class.
*/
virtual void DefaultInitializeRuntimeAttributes (
amrex::ParticleTile<NStructReal, NStructInt, NArrayReal,
NArrayInt,amrex::PinnedArenaAllocator>& pinned_tile,
amrex::ParticleTile<WarpXParticleContainer, NArrayReal,
NArrayInt, amrex::PinnedArenaAllocator>& pinned_tile,
const int n_external_attr_real,
const int n_external_attr_int,
const amrex::RandomEngine& engine) = 0;
Expand Down
25 changes: 14 additions & 11 deletions Source/ablastr/particles/ParticleMoments.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ namespace particles {
amrex::ParticleReal, amrex::ParticleReal>
MinAndMaxPositions (T_PC const & pc)
{
using PType = typename T_PC::SuperParticleType;
using ConstParticleTileDataType = typename T_PC::ParticleTileType::ConstParticleTileDataType;

// Get min and max for the local rank
amrex::ReduceOps<
Expand All @@ -46,11 +46,11 @@ namespace particles {
amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
>(
pc,
[=] AMREX_GPU_DEVICE(PType const & p) noexcept
[=] AMREX_GPU_DEVICE(const ConstParticleTileDataType& ptd, const int i) noexcept
{
amrex::ParticleReal const x = p.pos(0);
amrex::ParticleReal const y = p.pos(1);
amrex::ParticleReal const z = p.pos(2);
const amrex::ParticleReal x = ptd.rdata(0)[i];
const amrex::ParticleReal y = ptd.rdata(1)[i];
const amrex::ParticleReal z = ptd.rdata(2)[i];

return amrex::makeTuple(x, y, z, x, y, z);
},
Expand Down Expand Up @@ -90,7 +90,8 @@ namespace particles {
amrex::ParticleReal, amrex::ParticleReal>
MeanAndStdPositions (T_PC const & pc)
{
using PType = typename T_PC::SuperParticleType;

using ConstParticleTileDataType = typename T_PC::ParticleTileType::ConstParticleTileDataType;

amrex::ReduceOps<
amrex::ReduceOpSum, amrex::ReduceOpSum, amrex::ReduceOpSum,
Expand All @@ -103,12 +104,14 @@ namespace particles {
amrex::ParticleReal>
>(
pc,
[=] AMREX_GPU_DEVICE(const PType& p) noexcept
[=] AMREX_GPU_DEVICE(const ConstParticleTileDataType& ptd, const int i) noexcept
{
amrex::ParticleReal const x = p.pos(0);
amrex::ParticleReal const y = p.pos(1);
amrex::ParticleReal const z = p.pos(2);
amrex::ParticleReal const w = p.rdata(T_RealSoAWeight);

const amrex::ParticleReal x = ptd.rdata(0)[i];
const amrex::ParticleReal y = ptd.rdata(1)[i];
const amrex::ParticleReal z = ptd.rdata(2)[i];

const amrex::ParticleReal w = ptd.rdata(T_RealSoAWeight)[i];

return amrex::makeTuple(x, x*x, y, y*y, z, z*z, w);
},
Expand Down