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.

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

SoA: InitRandom/OnePerCell #3280

Open
wants to merge 1 commit into
base: development
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
24 changes: 21 additions & 3 deletions Src/Particle/AMReX_ParticleContainer.H
Original file line number Diff line number Diff line change
Expand Up @@ -106,11 +106,23 @@ struct ParticleLocData
* a given component, then the extra values will be set to zero. If more components
* are specified, it is a compile-time error.
*
* Note that the position attributes, first AMREX_SPACEDIM in ParticleReal, and the IDs, first
* two ints, are skipped here.
* Note that the counts in the type of the SoAParticle definition are explicit about these default-required
* attributes.
*
* Example usage:
*
* ParticleInitType<0, 2, 4, 1> pdata = {{}, {7, 9}, {1.5, 2.5, 3.5, 4.5}, {11}};
* ParticleInitType<Particle<0, 2>, 4, 1> pdata = {{}, {7, 9}, {1.5, 2.5, 3.5, 4.5}, {11}};
* ParticleInitTypeSoA<4, 3> pdata = {{1.5, 2.5, 3.5, 4.5}, {7, 9, 11}}; // 3D: SoAParticle<7, 5>
*/
template<int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
template <int NArrayReal, int NArrayInt>
struct ParticleInitTypeSoA
{
std::array<double, NArrayReal > real_array_data;
std::array<int, NArrayInt > int_array_data;
};
template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
struct ParticleInitType
{
std::array<double, NStructReal> real_struct_data;
Expand Down Expand Up @@ -177,7 +189,13 @@ public:

using ParticleContainerType = ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>;
using ParticleTileType = ParticleTile<ParticleType, NArrayReal, NArrayInt, Allocator>;
using ParticleInitData = ParticleInitType<NStructReal, NStructInt, NArrayReal, NArrayInt>;
using ParticleInitData = typename std::conditional<
T_ParticleType::is_soa_particle,
// SoA Particle: init w/o positions and id/cpu, but explicitly listed in define
ParticleInitTypeSoA<NArrayReal - AMREX_SPACEDIM, NArrayInt - 2>,
// legacy particle: init w/o positions and id/cpu, only implicitly listed in define
ParticleInitType<T_ParticleType::NReal, T_ParticleType::NInt, NArrayReal, NArrayInt>
>::type;

//! A single level worth of particles is indexed (grid id, tile id)
//! for both SoA and AoS data.
Expand Down
189 changes: 116 additions & 73 deletions Src/Particle/AMReX_ParticleInit.H
Original file line number Diff line number Diff line change
Expand Up @@ -1063,8 +1063,27 @@ InitRandom (Long icount,

if (who == MyProc) {

if constexpr(!ParticleType::is_soa_particle)
{
if constexpr(ParticleType::is_soa_particle) {
for (int i = 0; i < AMREX_SPACEDIM; i++) {
host_real_attribs[pld.m_lev][ind][i].push_back(pos[j*AMREX_SPACEDIM+i]);
}

host_int_attribs[pld.m_lev][ind][0].push_back(ParticleType::NextID());
host_int_attribs[pld.m_lev][ind][1].push_back(MyProc);

host_particles[pld.m_lev][ind];

// add the real...
for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
}

// ... and int array data
for (int i = 2; i < NArrayInt; i++) {
host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
}
}
else {
ParticleType p;
for (int i = 0; i < AMREX_SPACEDIM; i++) {
p.pos(i) = pos[j*AMREX_SPACEDIM + i];
Expand Down Expand Up @@ -1094,25 +1113,6 @@ InitRandom (Long icount,
for (int i = 0; i < NArrayInt; i++) {
host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
}
} else {
for (int i = 0; i < AMREX_SPACEDIM; i++) {
host_real_attribs[pld.m_lev][ind][i].push_back(pos[j*AMREX_SPACEDIM+i]);
}

host_int_attribs[pld.m_lev][ind][0].push_back(ParticleType::NextID());
host_int_attribs[pld.m_lev][ind][1].push_back(MyProc);

host_particles[pld.m_lev][ind];

// add the real...
for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
}

// ... and int array data
for (int i = 2; i < NArrayInt; i++) {
host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
}
}
}
}
Expand All @@ -1127,11 +1127,11 @@ InitRandom (Long icount,
auto& dst_tile = GetParticles(host_lev)[std::make_pair(grid,tile)];
auto old_size = dst_tile.GetArrayOfStructs().size();
auto new_size = old_size;
if constexpr(!ParticleType::is_soa_particle)
if constexpr(ParticleType::is_soa_particle)
{
new_size += src_tile.size();
} else {
new_size += host_real_attribs[host_lev][std::make_pair(grid,tile)][0].size();
} else {
new_size += src_tile.size();
}
dst_tile.resize(new_size);

Expand Down Expand Up @@ -1362,44 +1362,71 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
for (Long jcnt = 0; jcnt < icount_per_box; jcnt++) {
for (Long kcnt = 0; kcnt < icount_per_box; kcnt++)
{
p.pos(0) = static_cast<ParticleReal>(grid_box.lo(0) + (dist(mt) + double(icnt)) / double(icount_per_box) * grid_box.length(0));
p.pos(1) = static_cast<ParticleReal>(grid_box.lo(1) + (dist(mt) + double(jcnt)) / double(icount_per_box) * grid_box.length(1));
p.pos(2) = static_cast<ParticleReal>(grid_box.lo(2) + (dist(mt) + double(kcnt)) / double(icount_per_box) * grid_box.length(2));
// the position data
Copy link
Member Author

@ax3l ax3l May 30, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think InitRandomPerBox needs a bigger rewrite.

for (int d = 0; d < AMREX_SPACEDIM; d++) {
p.pos(d) = static_cast<ParticleReal>(grid_box.lo(d) + (dist(mt) + icnt) / icount_per_box * grid_box.length(d));
}

for (int i = 0; i < AMREX_SPACEDIM; i++)
AMREX_ASSERT(p.pos(i) < grid_box.hi(i));
for (int d = 0; d < AMREX_SPACEDIM; d++)
AMREX_ASSERT(p.pos(d) < grid_box.hi(d));

// the real struct data
for (int i = 0; i < NStructReal; i++) {
p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
}
if constexpr(ParticleType::is_soa_particle) {
if (!Where(p, pld)) {
amrex::Abort("ParticleContainer::InitRandom(): invalid particle");
}
AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
std::pair<int, int> ind(pld.m_grid, pld.m_tile);

// the int struct data
p.id() = ParticleType::NextID();
p.cpu() = ParallelDescriptor::MyProc();
// IDs
p.id() = ParticleType::NextID();
p.cpu() = ParallelDescriptor::MyProc();

for (int i = 0; i < NStructInt; i++) {
p.idata(i) = pdata.int_struct_data[i];
}
// add the real (after position)
for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
}

// locate the particle
if (!Where(p, pld)) {
amrex::Abort("ParticleContainer::InitRandomPerBox(): invalid particle");
// add the int array data (after id, cpu)
for (int i = 2; i < NArrayInt; i++) {
m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
}

// add
m_particles[pld.m_lev][ind].push_back(p);
}
AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
std::pair<int, int> ind(pld.m_grid, pld.m_tile);
else {
// the real struct data
for (int i = 0; i < NStructReal; i++) {
p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
}

// add the struct
m_particles[pld.m_lev][ind].push_back(p);
// the int struct data
p.id() = ParticleType::NextID();
p.cpu() = ParallelDescriptor::MyProc();

// add the real...
for (int i = 0; i < NArrayReal; i++) {
m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
}
for (int i = 0; i < NStructInt; i++) {
p.idata(i) = pdata.int_struct_data[i];
}

// ... and int array data
for (int i = 0; i < NArrayInt; i++) {
m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
// locate the particle
if (!Where(p, pld)) {
amrex::Abort("ParticleContainer::InitRandomPerBox(): invalid particle");
}
AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
std::pair<int, int> ind(pld.m_grid, pld.m_tile);

// add the struct
m_particles[pld.m_lev][ind].push_back(p);

// add the real...
for (int i = 0; i < NArrayReal; i++) {
m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
}

// ... and int array data
for (int i = 0; i < NArrayInt; i++) {
m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
}
}

} } }
Expand Down Expand Up @@ -1438,48 +1465,64 @@ InitOnePerCell (Real x_off, Real y_off, Real z_off, const ParticleInitData& pdat

const Real* dx = geom.CellSize();

ParticleType p;

// We'll generate the particles in parallel -- but no tiling of the grid here.
for (MFIter mfi(*m_dummy_mf[0], false); mfi.isValid(); ++mfi) {
Box grid = ParticleBoxArray(0)[mfi.index()];
auto ind = std::make_pair(mfi.index(), mfi.LocalTileIndex());
RealBox grid_box (grid,dx,geom.ProbLo());

// tile for one particle
ParticleTile<ParticleType, NArrayReal, NArrayInt, amrex::PinnedArenaAllocator> ptile_tmp;
ptile_tmp.resize(1);
auto ptd = ptile_tmp.getParticleTileData();

for (IntVect beg = grid.smallEnd(), end=grid.bigEnd(), cell = grid.smallEnd(); cell <= end; grid.next(cell))
{
// the real struct data
AMREX_D_TERM(p.pos(0) = static_cast<ParticleReal>(grid_box.lo(0) + (x_off + cell[0]-beg[0])*dx[0]);,
p.pos(1) = static_cast<ParticleReal>(grid_box.lo(1) + (y_off + cell[1]-beg[1])*dx[1]);,
p.pos(2) = static_cast<ParticleReal>(grid_box.lo(2) + (z_off + cell[2]-beg[2])*dx[2]););
// particle index
constexpr int i = 0;

// the position data
for (int d = 0; d < AMREX_SPACEDIM; d++) {
ptile_tmp.pos(i, d) = static_cast<ParticleReal>(grid_box.lo(d) + (x_off + cell[d]-beg[d])*dx[d]);
}

for (int d = 0; d < AMREX_SPACEDIM; ++d) {
AMREX_ASSERT(p.pos(d) < grid_box.hi(d));
AMREX_ASSERT(ptile_tmp.pos(i, d) < grid_box.hi(d));
}

for (int i = 0; i < NStructReal; i++) {
p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
if constexpr(!ParticleType::is_soa_particle) {
for (int n = 0; n < NStructReal; n++) {
ptd.rdata(n)[i] = static_cast<ParticleReal>(pdata.real_struct_data[n]);
}
}

// the int struct data
p.id() = ParticleType::NextID();
p.cpu() = ParallelDescriptor::MyProc();

for (int i = 0; i < NStructInt; i++) {
p.idata(i) = pdata.int_struct_data[i];
if constexpr(ParticleType::is_soa_particle) {
ptd.idata(0)[i] = ParticleType::NextID();
ptd.idata(1)[i] = ParallelDescriptor::MyProc();
}
else {
auto& p = make_particle<ParticleType>{}(ptd, i);
p.id() = ParticleType::NextID();
p.cpu() = ParallelDescriptor::MyProc();
}

// add the struct
ptile_tmp.push_back(p);
if constexpr(!ParticleType::is_soa_particle) {
for (int n = 0; n < NStructInt; n++) {
ptd.idata(n)[i] = pdata.int_struct_data[n];
}
}

// add the real...
for (int i = 0; i < NArrayReal; i++) {
ptile_tmp.push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
int n_min_real = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // jump over position
for (int n = n_min_real; n < NArrayReal; n++) {
ptile_tmp.push_back_real(n, static_cast<ParticleReal>(pdata.real_array_data[n]));
}

// ... and int array data
for (int i = 0; i < NArrayInt; i++) {
ptile_tmp.push_back_int(i, pdata.int_array_data[i]);
int n_min_int = ParticleType::is_soa_particle ? 2 : 0; // jump over cpuid
for (int n = n_min_int; n < NArrayInt; n++) {
ptile_tmp.push_back_int(n, pdata.int_array_data[n]);
}
}

Expand Down
6 changes: 3 additions & 3 deletions Tests/Particles/AssignDensity/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ void test_assign_density(TestParams& parms)
{

RealBox real_box;
for (int n = 0; n < BL_SPACEDIM; n++) {
for (int n = 0; n < AMREX_SPACEDIM; n++) {
real_box.setLo(n, 0.0);
real_box.setHi(n, 1.0);
}
Expand All @@ -48,10 +48,10 @@ void test_assign_density(TestParams& parms)
MultiFab density(ba, dmap, 1, 0);
density.setVal(0.0);

MultiFab partMF(ba, dmap, 1 + BL_SPACEDIM, 1);
MultiFab partMF(ba, dmap, 1 + AMREX_SPACEDIM, 1);
partMF.setVal(0.0);

typedef ParticleContainer<1 + BL_SPACEDIM> MyParticleContainer;
typedef ParticleContainer<1 + AMREX_SPACEDIM> MyParticleContainer;
MyParticleContainer myPC(geom, dmap, ba);
myPC.SetVerbose(false);

Expand Down
6 changes: 3 additions & 3 deletions Tests/Particles/CheckpointRestartSOA/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,8 @@ void test ()
}

// Add some particles
constexpr int NReal = 12;
constexpr int NInt = 4;
constexpr int NReal = AMREX_SPACEDIM + 12;
constexpr int NInt = 2 + 4;

typedef ParticleContainerPureSoA<NReal, NInt> MyPC;
MyPC myPC(geom, dmap, ba, ref_ratio);
Expand All @@ -108,7 +108,7 @@ void test ()
int num_particles = nppc * AMREX_D_TERM(ncells, * ncells, * ncells);
bool serialize = false;
int iseed = 451;
MyPC::ParticleInitData pdata = {{}, {},
MyPC::ParticleInitData pdata = {
{1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0},
{5, 14, 15, 16}};

Expand Down
6 changes: 3 additions & 3 deletions Tests/Particles/InitRandom/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@ void testSOA ()
}

// Add some particles
constexpr int NReal = 12;
constexpr int NInt = 4;
constexpr int NReal = AMREX_SPACEDIM + 12;

Check notice

Code scanning / CodeQL

Unused local variable

Variable NReal is not used.
constexpr int NInt = 2 + 4;

Check notice

Code scanning / CodeQL

Unused local variable

Variable NInt is not used.

typedef ParticleContainerPureSoA<NReal, NInt> MyPC;
MyPC myPC(geom, dmap, ba, ref_ratio);
Expand All @@ -70,7 +70,7 @@ void testSOA ()
int num_particles = nppc * AMREX_D_TERM(ncells, * ncells, * ncells);
bool serialize = false;
int iseed = 451;
MyPC::ParticleInitData pdata = {{}, {},
MyPC::ParticleInitData pdata = {
{1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0},
{5, 14, 15, 16}};

Expand Down
Loading