SoA: InitRandom/RandomPerBox/OnePerCell

Support more init methods on SoA particle containers.

Src/Particle/AMReX_ParticleContainer.H

@@ -108,17 +108,35 @@ struct ParticleLocData
 *
 * 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}};
+* ParticleInitType<SoAParticle<3, 2>, 3, 2> pdata = {{1.5, 2.5, 3.5, 4.5}, {7, 9}};
 */
-template<int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
-struct ParticleInitType
+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 ParticleInitTypeLegacy
 {
  std::array<double, NStructReal> real_struct_data;
  std::array<int, NStructInt > int_struct_data;
  std::array<double, NArrayReal > real_array_data;
  std::array<int, NArrayInt > int_array_data;
 };
 
+template <typename T_ParticleType, int NArrayReal=0, int NArrayInt=0>
+struct ParticleInitType
+ : public std::conditional<
+ T_ParticleType::is_soa_particle,
+ ParticleInitTypeSoA<NArrayReal, NArrayInt>,
+ ParticleInitTypeLegacy<T_ParticleType::NReal, T_ParticleType::NInt, NArrayReal, NArrayInt>
+ >::type
+{
+ using ParticleType = T_ParticleType;
+};
+
 template <bool is_const, typename T_ParticleType, int NArrayReal, int NArrayInt,
  template<class> class Allocator>
 
@@ -177,7 +195,7 @@ 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 = ParticleInitType<ParticleType, NArrayReal, NArrayInt>;
 
  //! A single level worth of particles is indexed (grid id, tile id)
  //! for both SoA and AoS data.

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<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]);
  }
  }
 

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];
  }