Skip to content

Commit

Permalink
SoA: InitRandom/RandomPerBox/OnePerCell
Browse files Browse the repository at this point in the history 
Support more init methods on SoA particle containers.
  • Loading branch information
@ax3l
ax3l committed May 30, 2023
1 parent 07f87b0 commit 54d0623
Show file tree
Hide file tree
Showing 5 changed files with 151 additions and 89 deletions.
24 changes: 21 additions & 3 deletions Src/Particle/AMReX_ParticleContainer.H
View file
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
View file
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) + icnt) / icount_per_box * grid_box.length(0));
p.pos(1) = static_cast<ParticleReal>(grid_box.lo(1) + (dist(mt) + jcnt) / icount_per_box * grid_box.length(1));
p.pos(2) = static_cast<ParticleReal>(grid_box.lo(2) + (dist(mt) + kcnt) / icount_per_box * grid_box.length(2));
// the position data
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
View file
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/InitRandom/main.cpp
View file
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;
constexpr int NInt = 2 + 4;

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
15 changes: 8 additions & 7 deletions Tests/Particles/ParticleIterator/main.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
#include "AMReX_Vector.H"
#include "AMReX_FabArray.H"
#include "AMReX_Particles.H"

#include <iostream>
#include <map>
#include <vector>

#include <AMReX_Vector.H>
#include "AMReX_FabArray.H"
#include "AMReX_Particles.H"

using namespace amrex;

Expand All @@ -19,7 +20,7 @@ int main(int argc, char* argv[])
int coord = 0;

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 Down Expand Up @@ -49,10 +50,10 @@ int main(int argc, char* argv[])
for (int lev = 0; lev < nlevs; lev++)
dmap[lev].define(ba[lev]);

typedef ParticleContainer<1+BL_SPACEDIM> MyParticleContainer;
using MyParticleContainer = ParticleContainer<1+AMREX_SPACEDIM>;
MyParticleContainer MyPC(geom, dmap, ba, rr);

MyParticleContainer::ParticleInitData pdata = {{1.0},{},{},{}};
MyParticleContainer::ParticleInitData pdata = {{1.0, AMREX_D_DECL(1.0, 2.0, 3.0)},{},{},{}};
MyPC.InitOnePerCell(0.5, 0.5, 0.5, pdata);
MyParticleContainer::do_tiling = true;

Expand All @@ -61,7 +62,7 @@ int main(int argc, char* argv[])
#ifdef AMREX_USE_OMP
#pragma omp parallel
#endif
for (ParIter<1+BL_SPACEDIM> mfi(MyPC, 0); mfi.isValid(); ++mfi) {
for (ParIter<1+AMREX_SPACEDIM> mfi(MyPC, 0); mfi.isValid(); ++mfi) {
amrex::AllPrintToFile("particle_iterator_out") << mfi.index() << " " << mfi.tileIndex() << "\n";
}
}
Expand Down

0 comments on commit 54d0623

Please sign in to comment.