From 54d0623298b357f9489c01e99ac1645ffea8aefe Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 26 Apr 2023 19:12:13 -0700 Subject: [PATCH] SoA: InitRandom/RandomPerBox/OnePerCell Support more init methods on SoA particle containers. --- Src/Particle/AMReX_ParticleContainer.H | 24 ++- Src/Particle/AMReX_ParticleInit.H | 189 +++++++++++++--------- Tests/Particles/AssignDensity/main.cpp | 6 +- Tests/Particles/InitRandom/main.cpp | 6 +- Tests/Particles/ParticleIterator/main.cpp | 15 +- 5 files changed, 151 insertions(+), 89 deletions(-) diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H index abffbefc860..a9353b46674 100644 --- a/Src/Particle/AMReX_ParticleContainer.H +++ b/Src/Particle/AMReX_ParticleContainer.H @@ -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, 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 +template +struct ParticleInitTypeSoA +{ + std::array real_array_data; + std::array int_array_data; +}; +template struct ParticleInitType { std::array real_struct_data; @@ -177,7 +189,13 @@ public: using ParticleContainerType = ParticleContainer_impl; using ParticleTileType = ParticleTile; - using ParticleInitData = ParticleInitType; + 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, + // legacy particle: init w/o positions and id/cpu, only implicitly listed in define + ParticleInitType + >::type; //! A single level worth of particles is indexed (grid id, tile id) //! for both SoA and AoS data. diff --git a/Src/Particle/AMReX_ParticleInit.H b/Src/Particle/AMReX_ParticleInit.H index 28427de32fb..15d51a6fdbd 100644 --- a/Src/Particle/AMReX_ParticleInit.H +++ b/Src/Particle/AMReX_ParticleInit.H @@ -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(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]; @@ -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(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]); - } } } } @@ -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); @@ -1362,44 +1362,71 @@ ParticleContainer_impl for (Long jcnt = 0; jcnt < icount_per_box; jcnt++) { for (Long kcnt = 0; kcnt < icount_per_box; kcnt++) { - p.pos(0) = static_cast(grid_box.lo(0) + (dist(mt) + icnt) / icount_per_box * grid_box.length(0)); - p.pos(1) = static_cast(grid_box.lo(1) + (dist(mt) + jcnt) / icount_per_box * grid_box.length(1)); - p.pos(2) = static_cast(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(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(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 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(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 ind(pld.m_grid, pld.m_tile); + else { + // the real struct data + for (int i = 0; i < NStructReal; i++) { + p.rdata(i) = static_cast(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(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 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(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]); + } } } } } @@ -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 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(grid_box.lo(0) + (x_off + cell[0]-beg[0])*dx[0]);, - p.pos(1) = static_cast(grid_box.lo(1) + (y_off + cell[1]-beg[1])*dx[1]);, - p.pos(2) = static_cast(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(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(pdata.real_struct_data[i]); + if constexpr(!ParticleType::is_soa_particle) { + for (int n = 0; n < NStructReal; n++) { + ptd.rdata(n)[i] = static_cast(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{}(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(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(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]); } } diff --git a/Tests/Particles/AssignDensity/main.cpp b/Tests/Particles/AssignDensity/main.cpp index 146781ecb6b..7656faf8e20 100644 --- a/Tests/Particles/AssignDensity/main.cpp +++ b/Tests/Particles/AssignDensity/main.cpp @@ -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); } @@ -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); diff --git a/Tests/Particles/InitRandom/main.cpp b/Tests/Particles/InitRandom/main.cpp index c3c9f730b53..2d7cec84c62 100644 --- a/Tests/Particles/InitRandom/main.cpp +++ b/Tests/Particles/InitRandom/main.cpp @@ -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 MyPC; MyPC myPC(geom, dmap, ba, ref_ratio); @@ -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}}; diff --git a/Tests/Particles/ParticleIterator/main.cpp b/Tests/Particles/ParticleIterator/main.cpp index 074759de0c5..4647f9274b2 100644 --- a/Tests/Particles/ParticleIterator/main.cpp +++ b/Tests/Particles/ParticleIterator/main.cpp @@ -1,10 +1,11 @@ +#include "AMReX_Vector.H" +#include "AMReX_FabArray.H" +#include "AMReX_Particles.H" + #include #include #include -#include -#include "AMReX_FabArray.H" -#include "AMReX_Particles.H" using namespace amrex; @@ -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); @@ -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; @@ -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"; } }