diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H index 64c44be84ac..c6123e1d7f8 100644 --- a/Src/Particle/AMReX_ParticleContainer.H +++ b/Src/Particle/AMReX_ParticleContainer.H @@ -108,10 +108,17 @@ struct ParticleLocData * * 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}}; +* ParticleInitType, 3, 2> pdata = {{1.5, 2.5, 3.5, 4.5}, {7, 9}}; */ -template -struct ParticleInitType +template +struct ParticleInitTypeSoA +{ + std::array real_array_data; + std::array int_array_data; +}; +template +struct ParticleInitTypeLegacy { std::array real_struct_data; std::array int_struct_data; @@ -119,6 +126,17 @@ struct ParticleInitType std::array int_array_data; }; +template +struct ParticleInitType + : public std::conditional< + T_ParticleType::is_soa_particle, + ParticleInitTypeSoA, + ParticleInitTypeLegacy + >::type +{ + using ParticleType = T_ParticleType; +}; + template class Allocator> @@ -177,7 +195,7 @@ public: using ParticleContainerType = ParticleContainer_impl; using ParticleTileType = ParticleTile; - using ParticleInitData = ParticleInitType; + using ParticleInitData = ParticleInitType; //! 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 dbb354800ef..3899b2a5393 100644 --- a/Src/Particle/AMReX_ParticleInit.H +++ b/Src/Particle/AMReX_ParticleInit.H @@ -1354,48 +1354,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/Src/Particle/AMReX_ParticleTile.H b/Src/Particle/AMReX_ParticleTile.H index ba68812e5bf..77a3ddd7eb4 100644 --- a/Src/Particle/AMReX_ParticleTile.H +++ b/Src/Particle/AMReX_ParticleTile.H @@ -68,6 +68,16 @@ struct ParticleTileData return this->m_rdata[dir][index]; } + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + auto& cpuid (const int index) & + { + if constexpr(!ParticleType::is_soa_particle) { + return this->m_aos[index].id(); + } else { + return this->m_idata[0][index]; + } + } + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE auto id (const int index) const & { @@ -82,7 +92,7 @@ struct ParticleTileData auto& id (const int index) & { if constexpr(!ParticleType::is_soa_particle) { - return this->m_aos[index].id(); + return this->m_aos[index].id(); } else { return this->m_idata[0][index]; }