Skip to content

Commit

Permalink
SoA-only Particle Container (#1444)
Browse files Browse the repository at this point in the history 
* made all real particle attributes SoA for performance; particle initialization functions run on the GPU

* bugfix for uniform particle distribution with more than 1 particle per cell

---------

Co-authored-by: David Gardner <[email protected]>
  • Loading branch information
@debog @gardner48
debog and gardner48 authored Feb 16, 2024
1 parent 0a1e223 commit a580170
Show file tree
Hide file tree
Showing 3 changed files with 354 additions and 270 deletions.
124 changes: 68 additions & 56 deletions Source/Particles/ERFPC.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,23 +6,35 @@
#include <string>
#include <AMReX_Particles.H>

struct ERFParticlesRealIdx
struct ERFParticlesIntIdxAoS
{
enum {
vx = 0,
vy,
vz,
mass,
k = 0,
ncomps
};
};

struct ERFParticlesIntIdx
struct ERFParticlesRealIdxAoS
{
enum {
ncomps = 0
};
};

struct ERFParticlesIntIdxSoA
{
enum {
i = 0,
j,
k,
ncomps = 0
};
};

struct ERFParticlesRealIdxSoA
{
enum {
vx = 0,
vy,
vz,
mass,
ncomps
};
};
Expand Down Expand Up @@ -53,7 +65,7 @@ struct ERFParticlesAssignor
amrex::IntVect iv(
AMREX_D_DECL( int(amrex::Math::floor((p.pos(0)-plo[0])*dxi[0])),
int(amrex::Math::floor((p.pos(1)-plo[1])*dxi[1])),
p.idata(ERFParticlesIntIdx::k) ) );
p.idata(ERFParticlesIntIdxAoS::k) ) );
iv[0] += domain.smallEnd()[0];
iv[1] += domain.smallEnd()[1];
return iv;
Expand All @@ -62,42 +74,39 @@ struct ERFParticlesAssignor

template <typename P>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void update_location_idata (P& p,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi,
const amrex::Array4<amrex::Real const>& height_arr)
void update_location_idata ( P& a_p,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& a_plo,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& a_dxi,
const amrex::Array4<amrex::Real const>& a_height_arr )
{
amrex::IntVect iv( int(amrex::Math::floor((p.pos(0)-plo[0])*dxi[0])),
int(amrex::Math::floor((p.pos(1)-plo[1])*dxi[1])),
p.idata(ERFParticlesIntIdx::k) );

p.idata(ERFParticlesIntIdx::i) = iv[0];
p.idata(ERFParticlesIntIdx::j) = iv[1];

if (height_arr) {
amrex::Real lx = (p.pos(0)-plo[0])*dxi[0] - static_cast<amrex::Real>(iv[0]);
amrex::Real ly = (p.pos(1)-plo[1])*dxi[1] - static_cast<amrex::Real>(iv[1]);
auto zlo = height_arr(iv[0] ,iv[1] ,iv[2] ) * (1.0-lx) * (1.0-ly) +
height_arr(iv[0]+1,iv[1] ,iv[2] ) * lx * (1.0-ly) +
height_arr(iv[0] ,iv[1]+1,iv[2] ) * (1.0-lx) * ly +
height_arr(iv[0]+1,iv[1]+1,iv[2] ) * lx * ly;
auto zhi = height_arr(iv[0] ,iv[1] ,iv[2]+1) * (1.0-lx) * (1.0-ly) +
height_arr(iv[0]+1,iv[1] ,iv[2]+1) * lx * (1.0-ly) +
height_arr(iv[0] ,iv[1]+1,iv[2]+1) * (1.0-lx) * ly +
height_arr(iv[0]+1,iv[1]+1,iv[2]+1) * lx * ly;

if (p.pos(2) > zhi) {
p.idata(ERFParticlesIntIdx::k) += 1;
} else if (p.pos(2) <= zlo) {
p.idata(ERFParticlesIntIdx::k) -= 1;
amrex::IntVect iv( int(amrex::Math::floor((a_p.pos(0)-a_plo[0])*a_dxi[0])),
int(amrex::Math::floor((a_p.pos(1)-a_plo[1])*a_dxi[1])),
a_p.idata(ERFParticlesIntIdxAoS::k) );

if (a_height_arr) {
amrex::Real lx = (a_p.pos(0)-a_plo[0])*a_dxi[0] - static_cast<amrex::Real>(iv[0]);
amrex::Real ly = (a_p.pos(1)-a_plo[1])*a_dxi[1] - static_cast<amrex::Real>(iv[1]);
auto zlo = a_height_arr(iv[0] ,iv[1] ,iv[2] ) * (1.0-lx) * (1.0-ly) +
a_height_arr(iv[0]+1,iv[1] ,iv[2] ) * lx * (1.0-ly) +
a_height_arr(iv[0] ,iv[1]+1,iv[2] ) * (1.0-lx) * ly +
a_height_arr(iv[0]+1,iv[1]+1,iv[2] ) * lx * ly;
auto zhi = a_height_arr(iv[0] ,iv[1] ,iv[2]+1) * (1.0-lx) * (1.0-ly) +
a_height_arr(iv[0]+1,iv[1] ,iv[2]+1) * lx * (1.0-ly) +
a_height_arr(iv[0] ,iv[1]+1,iv[2]+1) * (1.0-lx) * ly +
a_height_arr(iv[0]+1,iv[1]+1,iv[2]+1) * lx * ly;

if (a_p.pos(2) > zhi) {
a_p.idata(ERFParticlesIntIdxAoS::k) += 1;
} else if (a_p.pos(2) <= zlo) {
a_p.idata(ERFParticlesIntIdxAoS::k) -= 1;
}
}
}

class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
ERFParticlesIntIdx::ncomps,
0,
0,
class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdxAoS::ncomps, // AoS real attributes
ERFParticlesIntIdxAoS::ncomps, // AoS integer attributes
ERFParticlesRealIdxSoA::ncomps, // SoA real attributes
ERFParticlesIntIdxSoA::ncomps, // SoA integer attributes
amrex::DefaultAllocator,
ERFParticlesAssignor >
{
Expand All @@ -106,10 +115,10 @@ class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
/*! Constructor */
ERFPC ( amrex::ParGDBBase* a_gdb,
const std::string& a_name = "particles" )
: amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
ERFParticlesIntIdx::ncomps,
0,
0,
: amrex::ParticleContainer< ERFParticlesRealIdxAoS::ncomps, // AoS real attributes
ERFParticlesIntIdxAoS::ncomps, // AoS integer attributes
ERFParticlesRealIdxSoA::ncomps, // SoA real attributes
ERFParticlesIntIdxSoA::ncomps, // SoA integer attributes
amrex::DefaultAllocator,
ERFParticlesAssignor> (a_gdb)
{
Expand All @@ -123,10 +132,10 @@ class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
const amrex::DistributionMapping& a_dmap,
const amrex::BoxArray& a_ba,
const std::string& a_name = "particles" )
: amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
ERFParticlesIntIdx::ncomps,
0,
0,
: amrex::ParticleContainer< ERFParticlesRealIdxAoS::ncomps, // AoS real attributes
ERFParticlesIntIdxAoS::ncomps, // AoS real attributes
ERFParticlesRealIdxSoA::ncomps, // SoA real attributes
ERFParticlesIntIdxSoA::ncomps, // SoA integer attributes
amrex::DefaultAllocator,
ERFParticlesAssignor> ( a_geom, a_dmap, a_ba )
{
Expand Down Expand Up @@ -175,6 +184,16 @@ class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
m_advect_w_gravity = a_flag;
}

// the following functions should ideally be private or protected, but need to be
// public due to CUDA extended lambda capture rules

/*! Default initialization for tracer particles for WoA case (ref: AA) */
void initializeParticlesDefaultTracersWoA (const std::unique_ptr<amrex::MultiFab>& a_ptr=nullptr);
/*! Default initialization for hydro particles (ref: AA) */
void initializeParticlesDefaultHydro (const std::unique_ptr<amrex::MultiFab>& a_ptr = nullptr);
/*! Default particle initialization */
void initializeParticlesUniformDistribution (const std::unique_ptr<amrex::MultiFab>& a_ptr = nullptr);

protected:

bool m_advect_w_flow; /*!< advect with flow velocity */
Expand All @@ -188,17 +207,10 @@ class ERFPC : public amrex::ParticleContainer< ERFParticlesRealIdx::ncomps,
/*! read inputs from file */
virtual void readInputs ();

/*! Default particle initialization */
void initializeParticlesUniformDistribution (const std::unique_ptr<amrex::MultiFab>& a_ptr = nullptr);

private:

/*! Default particle initialization */
void initializeParticlesDefault (const std::unique_ptr<amrex::MultiFab>& a_ptr = nullptr);
/*! Default initialization for tracer particles for WoA case (ref: AA) */
void initializeParticlesDefaultTracersWoA (const std::unique_ptr<amrex::MultiFab>& a_ptr=nullptr);
/*! Default initialization for hydro particles (ref: AA) */
void initializeParticlesDefaultHydro (const std::unique_ptr<amrex::MultiFab>& a_ptr = nullptr);

};

Expand Down
37 changes: 21 additions & 16 deletions Source/Particles/ERFPCEvolve.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,8 @@ void ERFPC::AdvectWithFlow ( MultiFab* a_umac,

const auto strttime = amrex::second();
const Geometry& geom = m_gdb->Geom(a_lev);
const Box& domain = geom.Domain();
const auto plo = geom.ProbLoArray();
const auto dxi = geom.InvCellSizeArray();
const auto dx = geom.CellSizeArray();

Vector<std::unique_ptr<MultiFab> > raii_umac(AMREX_SPACEDIM);
Vector<MultiFab*> umac_pointer(AMREX_SPACEDIM);
Expand Down Expand Up @@ -83,8 +81,15 @@ void ERFPC::AdvectWithFlow ( MultiFab* a_umac,
int grid = pti.index();
auto& ptile = ParticlesAt(a_lev, pti);
auto& aos = ptile.GetArrayOfStructs();
auto& soa = ptile.GetStructOfArrays();
const int n = aos.numParticles();
auto *p_pbox = aos().data();

Array<ParticleReal*,AMREX_SPACEDIM> v_ptr;
v_ptr[0] = soa.GetRealData(ERFParticlesRealIdxSoA::vx).data();
v_ptr[1] = soa.GetRealData(ERFParticlesRealIdxSoA::vy).data();
v_ptr[2] = soa.GetRealData(ERFParticlesRealIdxSoA::vz).data();

const FArrayBox* fab[AMREX_SPACEDIM] = { AMREX_D_DECL(&((*umac_pointer[0])[grid]),
&((*umac_pointer[1])[grid]),
&((*umac_pointer[2])[grid])) };
Expand All @@ -102,28 +107,27 @@ void ERFPC::AdvectWithFlow ( MultiFab* a_umac,
{
ParticleType& p = p_pbox[i];
if (p.id() <= 0) { return; }

ParticleReal v[AMREX_SPACEDIM];
if (use_terrain) {
mac_interpolate_mapped_z(p, plo, dxi, umacarr, zheight, v);
} else {
mac_interpolate(p, plo, dxi, umacarr, v);
}
if (ipass == 0)
{

if (ipass == 0) {
for (int dim=0; dim < AMREX_SPACEDIM; dim++)
{
p.rdata(dim) = p.pos(dim);
v_ptr[dim][i] = p.pos(dim);
p.pos(dim) += static_cast<ParticleReal>(ParticleReal(0.5)*a_dt*v[dim]);
}
// Update z-coordinate carried by the particle
update_location_idata(p,plo,dxi,zheight);
}
else
{
} else {
for (int dim=0; dim < AMREX_SPACEDIM; dim++)
{
p.pos(dim) = p.rdata(dim) + static_cast<ParticleReal>(a_dt*v[dim]);
p.rdata(dim) = v[dim];
p.pos(dim) = v_ptr[dim][i] + static_cast<ParticleReal>(a_dt*v[dim]);
v_ptr[dim][i] = v[dim];
}
// Update z-coordinate carried by the particle
update_location_idata(p,plo,dxi,zheight);
Expand Down Expand Up @@ -158,9 +162,7 @@ void ERFPC::AdvectWithGravity ( int a_lev,

const auto strttime = amrex::second();
const Geometry& geom = m_gdb->Geom(a_lev);
const Box& domain = geom.Domain();
const auto plo = geom.ProbLoArray();
const auto dx = geom.CellSizeArray();
const auto dxi = geom.InvCellSizeArray();

#ifdef AMREX_USE_OMP
Expand All @@ -171,9 +173,12 @@ void ERFPC::AdvectWithGravity ( int a_lev,
int grid = pti.index();
auto& ptile = ParticlesAt(a_lev, pti);
auto& aos = ptile.GetArrayOfStructs();
auto& soa = ptile.GetStructOfArrays();
const int n = aos.numParticles();
auto *p_pbox = aos().data();

auto vz_ptr = soa.GetRealData(ERFParticlesRealIdxSoA::vz).data();

bool use_terrain = (a_z_height != nullptr);
auto zheight = use_terrain ? (*a_z_height)[grid].array() : Array4<Real>{};

Expand All @@ -182,7 +187,7 @@ void ERFPC::AdvectWithGravity ( int a_lev,
ParticleType& p = p_pbox[i];
if (p.id() <= 0) { return; }

ParticleReal v = p.rdata(ERFParticlesRealIdx::vz);
ParticleReal v = vz_ptr[i];

// Define acceleration to be (gravity minus drag) where drag is defined
// such the particles will reach a terminal velocity of 5.0 (totally arbitrary)
Expand All @@ -193,13 +198,13 @@ void ERFPC::AdvectWithGravity ( int a_lev,
ParticleReal half_dt = 0.5 * a_dt;

// Update the particle velocity over first half of step (a_dt/2)
p.rdata(ERFParticlesRealIdx::vz) -= (grav - drag) * half_dt;
vz_ptr[i] -= (grav - drag) * half_dt;

// Update the particle position over (a_dt)
p.pos(2) += static_cast<ParticleReal>(ParticleReal(0.5)*a_dt*p.rdata(ERFParticlesRealIdx::vz));
p.pos(2) += static_cast<ParticleReal>(ParticleReal(0.5)*a_dt*vz_ptr[i]);

// Update the particle velocity over second half of step (a_dt/2)
p.rdata(ERFParticlesRealIdx::vz) -= (grav - drag) * half_dt;
vz_ptr[i] -= (grav - drag) * half_dt;

// also update z-coordinate here
update_location_idata(p,plo,dxi,zheight);
Expand Down
Loading

0 comments on commit a580170

Please sign in to comment.