SoA: InitRandom/RandomPerBox/OnePerCell

Support more init methods on SoA particle containers.
AMReX-Codes · Apr 27, 2023 · 3ca6008 · 3ca6008
1 parent a6e89ec
commit 3ca6008
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 25 deletions.
diff --git a/Src/Particle/AMReX_ParticleContainer.H b/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.

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

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