Add version of ParticleTile using runtime-only 2D array

Src/Particle/AMReX_NeighborParticles.H

@@ -9,6 +9,7 @@
 #include <AMReX_NeighborList.H>
 #include <AMReX_OpenMP.H>
 #include <AMReX_ParticleTile.H>
+#include <AMReX_ParticleTileRT.H>
 
 namespace amrex {
 

Src/Particle/AMReX_ParIter.H

@@ -18,6 +18,9 @@ struct Particle;
 template <int NArrayReal, int NArrayInt>
 struct SoAParticle;
 
+template <class RType, class IType>
+struct RTSoAParticle;
+
 struct DefaultAssignor;
 
 // for backwards compatibility
@@ -87,11 +90,11 @@ public:
 
     [[nodiscard]] SoARef GetStructOfArrays () const { return GetParticleTile().GetStructOfArrays(); }
 
-    [[nodiscard]] int numParticles () const { return GetParticleTile().numParticles(); }
+    [[nodiscard]] auto numParticles () const { return GetParticleTile().numParticles(); }
 
-    [[nodiscard]] int numRealParticles () const { return GetParticleTile().numRealParticles(); }
+    [[nodiscard]] auto numRealParticles () const { return GetParticleTile().numRealParticles(); }
 
-    [[nodiscard]] int numNeighborParticles () const { return GetParticleTile().numNeighborParticles(); }
+    [[nodiscard]] auto numNeighborParticles () const { return GetParticleTile().numNeighborParticles(); }
 
     [[nodiscard]] Long capacity () const { return GetParticleTile().capacity(); }
 
@@ -227,13 +230,19 @@ template <int T_NArrayReal, int T_NArrayInt, template<class> class Allocator=Def
 using ParConstIterSoA = ParConstIter_impl<SoAParticle<T_NArrayReal, T_NArrayInt>, T_NArrayReal, T_NArrayInt, Allocator, CellAssignor>;
 
 /// \ingroup amrex_particles
+template <class RType=ParticleReal, class IType=int, class CellAssignor=DefaultAssignor>
+using ParConstIterRTSoA = ParConstIter_impl<RTSoAParticle<RType, IType>, 0, 0, PolymorphicArenaAllocator, CellAssignor>;
+
 template <int T_NStructReal, int T_NStructInt=0, int T_NArrayReal=0, int T_NArrayInt=0,
           template<class> class Allocator=DefaultAllocator, class CellAssignor=DefaultAssignor>
 using ParIter = ParIter_impl<Particle<T_NStructReal, T_NStructInt>, T_NArrayReal, T_NArrayInt, Allocator, CellAssignor>;
 
 template <int T_NArrayReal, int T_NArrayInt, template<class> class Allocator=DefaultAllocator, class CellAssignor=DefaultAssignor>
 using ParIterSoA = ParIter_impl<SoAParticle<T_NArrayReal, T_NArrayInt>, T_NArrayReal, T_NArrayInt, Allocator, CellAssignor>;
 
+template <class RType=ParticleReal, class IType=int, class CellAssignor=DefaultAssignor>
+using ParIterRTSoA = ParIter_impl<RTSoAParticle<RType, IType>, 0, 0, PolymorphicArenaAllocator, CellAssignor>;
+
 }
 
 #endif

Src/Particle/AMReX_Particle.H

@@ -389,6 +389,7 @@ struct SoAParticleBase
     static constexpr int NReal=0;
     static constexpr int NInt=0;
     static constexpr bool is_soa_particle = true;
+    static constexpr bool is_rtsoa_particle = false;
 };
 
 /**
@@ -403,6 +404,7 @@ struct alignas(sizeof(double)) Particle
     : ParticleBase<ParticleReal,T_NReal,T_NInt>
 {
     static constexpr bool is_soa_particle = false;
+    static constexpr bool is_rtsoa_particle = false;
     using StorageParticleType = Particle;
     using ConstType = Particle const;
 

Src/Particle/AMReX_ParticleContainer.H

@@ -23,6 +23,7 @@
 #include <AMReX_ArrayOfStructs.H>
 #include <AMReX_Particle.H>
 #include <AMReX_ParticleTile.H>
+#include <AMReX_ParticleTileRT.H>
 #include <AMReX_TypeTraits.H>
 #include <AMReX_GpuContainers.H>
 #include <AMReX_ParticleUtil.H>
@@ -180,8 +181,14 @@ public:
     RealDescriptor ParticleRealDescriptor = FPC::Native64RealDescriptor();
 #endif
 
+    static constexpr bool is_rtsoa_pc = ParticleType::is_rtsoa_particle;
+
     using ParticleContainerType = ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>;
-    using ParticleTileType = ParticleTile<ParticleType, NArrayReal, NArrayInt, Allocator>;
+    using ParticleTileType = std::conditional_t<
+        is_rtsoa_pc,
+        ParticleTileRT<typename ParticleType::RealType, typename ParticleType::IntType>,
+        ParticleTile<ParticleType, NArrayReal, NArrayInt, Allocator>
+    >;
     using ParticleInitData = ParticleInitType<NStructReal, NStructInt, NArrayReal, NArrayInt>;
 
     //! A single level worth of particles is indexed (grid id, tile id)
@@ -202,7 +209,7 @@ public:
     static constexpr bool has_polymorphic_allocator =
         IsPolymorphicArenaAllocator<Allocator<RealType>>::value;
 
-    static_assert(!ParticleType::is_soa_particle || T_NArrayReal >= AMREX_SPACEDIM,
+    static_assert(!ParticleType::is_soa_particle || T_NArrayReal >= AMREX_SPACEDIM || is_rtsoa_pc,
                   "Pure SoA particle containers require at least AMREX_SPACEDIM real components for positions");
 
     //! \brief Default constructor - construct an empty particle container that has no concept
@@ -1467,12 +1474,13 @@ protected:
                                int lev_min = 0, int lev_max = -1, int local_grid=-1) const;
 
 public:
+    using FlagsVector = PODVector<int, Allocator<int>>;
+
     void
     WriteParticles (int level, std::ofstream& ofs, int fnum,
                     Vector<int>& which, Vector<int>& count, Vector<Long>& where,
                     const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
-                    const Vector<std::map<std::pair<int, int>,IntVector>>& particle_io_flags, bool is_checkpoint) const;
-
+                    const Vector<std::map<std::pair<int, int>,FlagsVector>>& particle_io_flags, bool is_checkpoint) const;
 #ifdef AMREX_USE_HDF5
 #include "AMReX_ParticlesHDF5.H"
 #endif
@@ -1564,6 +1572,9 @@ using ParticleContainer = ParticleContainer_impl<Particle<T_NStructReal, T_NStru
 template <int T_NArrayReal, int T_NArrayInt, template<class> class Allocator=DefaultAllocator, class CellAssignor=DefaultAssignor>
 using ParticleContainerPureSoA = ParticleContainer_impl<SoAParticle<T_NArrayReal, T_NArrayInt>, T_NArrayReal, T_NArrayInt, Allocator, CellAssignor>;
 
+template <class RType=ParticleReal, class IType=int, class CellAssignor=DefaultAssignor>
+using ParticleContainerRTSoA = ParticleContainer_impl<RTSoAParticle<RType, IType>, 0, 0, PolymorphicArenaAllocator, CellAssignor>;
+
 }
 
 #include "AMReX_ParticleInit.H"

Src/Particle/AMReX_ParticleContainerI.H

@@ -1143,7 +1143,7 @@ void
 ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::
 copyParticles (const PCType& other, bool local)
 {
-    using PData = ConstParticleTileData<typename ParticleType::StorageParticleType, NArrayReal, NArrayInt>;
+    using PData = typename ParticleTileType::ConstParticleTileDataType;
     copyParticles(other, [] AMREX_GPU_HOST_DEVICE (const PData& /*data*/, int /*i*/) { return 1; }, local);
 }
 
@@ -1154,7 +1154,7 @@ void
 ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::
 addParticles (const PCType& other, bool local)
 {
-    using PData = ConstParticleTileData<typename ParticleType::StorageParticleType, NArrayReal, NArrayInt>;
+    using PData = typename ParticleTileType::ConstParticleTileDataType;
     addParticles(other, [] AMREX_GPU_HOST_DEVICE (const PData& /*data*/, int /*i*/) { return 1; }, local);
 }
 
@@ -1217,7 +1217,8 @@ addParticles (const PCType& other, F const& f, bool local)
             auto dst_index = ptile.numParticles();
             ptile.resize(dst_index + np);
 
-            auto count = amrex::filterParticles(ptile, ptile_other, f, 0, dst_index, np);
+            auto count = amrex::filterParticles(ptile, ptile_other, f,
+                static_cast<decltype(np)>(0), dst_index, np);
 
             ptile.resize(dst_index + count);
         }
@@ -1265,85 +1266,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     const size_t np_total = np + ptile.numNeighborParticles();
 
     if (memEfficientSort) {
-#if defined(AMREX_USE_CUDA) && defined(_WIN32)
-        if           (!ParticleType::is_soa_particle) {
-#else
-        if constexpr (!ParticleType::is_soa_particle) {
-#endif
-            static_assert(sizeof(ParticleType)%4 == 0 && sizeof(uint32_t) == 4);
-            using tmp_t = std::conditional_t<sizeof(ParticleType)%8 == 0,
-                                             uint64_t, uint32_t>;
-            constexpr std::size_t nchunks = sizeof(ParticleType) / sizeof(tmp_t);
-            Gpu::DeviceVector<tmp_t> tmp(np);
-            auto* ptmp = tmp.data();
-            auto* paos = (tmp_t*)(ptile.getParticleTileData().m_aos);
-            for (std::size_t ichunk = 0; ichunk < nchunks; ++ichunk) {
-                // Do not need to reorder neighbor particles
-                AMREX_HOST_DEVICE_FOR_1D(np, i,
-                {
-                    ptmp[i] = paos[permutations[i]*nchunks+ichunk];
-                });
-                AMREX_HOST_DEVICE_FOR_1D(np, i,
-                {
-                    paos[i*nchunks+ichunk] = ptmp[i];
-                });
-            }
-            Gpu::streamSynchronize();
-        } else {
-            typename SoA::IdCPU tmp_idcpu;
-            if constexpr (has_polymorphic_allocator) {
-                tmp_idcpu.setArena(arena());
-            }
-            tmp_idcpu.resize(np_total);
-            auto src = ptile.GetStructOfArrays().GetIdCPUData().data();
-            uint64_t* dst = tmp_idcpu.data();
-            AMREX_HOST_DEVICE_FOR_1D( np_total, i,
-            {
-                dst[i] = i < np ? src[permutations[i]] : src[i];
-            });
-
-            Gpu::streamSynchronize();
-
-            ptile.GetStructOfArrays().GetIdCPUData().swap(tmp_idcpu);
-        }
-
-        { // Create a scope for the temporary vector below
-            RealVector tmp_real;
-            if constexpr (has_polymorphic_allocator) {
-                tmp_real.setArena(arena());
-            }
-            tmp_real.resize(np_total);
-            for (int comp = 0; comp < NArrayReal + m_num_runtime_real; ++comp) {
-                auto src = ptile.GetStructOfArrays().GetRealData(comp).data();
-                ParticleReal* dst = tmp_real.data();
-                AMREX_HOST_DEVICE_FOR_1D( np_total, i,
-                {
-                    dst[i] = i < np ? src[permutations[i]] : src[i];
-                });
-
-                Gpu::streamSynchronize();
-
-                ptile.GetStructOfArrays().GetRealData(comp).swap(tmp_real);
-            }
-        }
-
-        IntVector tmp_int;
-        if constexpr (has_polymorphic_allocator) {
-            tmp_int.setArena(arena());
-        }
-        tmp_int.resize(np_total);
-        for (int comp = 0; comp < NArrayInt + m_num_runtime_int; ++comp) {
-            auto src = ptile.GetStructOfArrays().GetIntData(comp).data();
-            int* dst = tmp_int.data();
-            AMREX_HOST_DEVICE_FOR_1D( np_total , i,
-            {
-                dst[i] = i < np ? src[permutations[i]] : src[i];
-            });
-
-            Gpu::streamSynchronize();
-
-            ptile.GetStructOfArrays().GetIntData(comp).swap(tmp_int);
-        }
+        amrex::ReorderParticles(ptile, permutations);
     } else {
         ParticleTileType ptile_tmp;
         ptile_tmp.define(m_num_runtime_real, m_num_runtime_int,
@@ -1838,7 +1761,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
                     ++pindex;
                 }
 
-                auto tot_npart = particle_tile->numTotalParticles();
+                Long tot_npart = particle_tile->numTotalParticles();
 
                 if (last != npart - 1) {
                     pindex = last + 1;
@@ -1881,9 +1804,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
 
             for (int i = 0; i < num_threads; ++i) {
 
-                auto to_copy = src_ptile[i].numParticles();
-                auto old_size = dst_ptile.numTotalParticles();
-                auto new_size = old_size + to_copy;
+                Long to_copy = src_ptile[i].numParticles();
+                Long old_size = dst_ptile.numTotalParticles();
+                Long new_size = old_size + to_copy;
                 dst_ptile.resize(new_size);
 
                 auto dst_ptd = dst_ptile.getParticleTileData();

Src/Particle/AMReX_ParticleIO.H

@@ -586,7 +586,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
                   Vector<int>& which, Vector<int>& count, Vector<Long>& where,
                   const Vector<int>& write_real_comp,
                   const Vector<int>& write_int_comp,
-                  const Vector<std::map<std::pair<int, int>, IntVector>>& particle_io_flags,
+                  const Vector<std::map<std::pair<int, int>, FlagsVector>>& particle_io_flags,
                   bool is_checkpoint) const
 {
     BL_PROFILE("ParticleContainer::WriteParticles()");

Src/Particle/AMReX_ParticleMesh.H

@@ -37,6 +37,21 @@ auto call_f (F const& f,
         return f(p, i, fabarr);
     }
 }
+
+template <typename F, typename T, template<class, class> class PTDType, class RType, class IType>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+auto call_f (F const& f,
+             const PTDType<RType, IType>& p,
+             const int i, Array4<T> const& fabarr,
+             GpuArray<Real,AMREX_SPACEDIM> const& plo,
+             GpuArray<Real,AMREX_SPACEDIM> const& dxi) noexcept
+{
+    if constexpr (IsCallable<F, decltype(p), int, decltype(fabarr), decltype(plo), decltype(dxi)>::value) {
+        return f(p, i, fabarr, plo, dxi);
+    } else {
+        return f(p, i, fabarr);
+    }
+}
 }
 /// \endcond
 

Src/Particle/AMReX_ParticleReduce.H

@@ -34,6 +34,15 @@ auto call_f (F const& f,
         return f(p, i);
     }
 }
+
+template <typename F, template<class, class> class PTDType, class RType, class IType>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+auto call_f (F const& f,
+             const PTDType<RType, IType>& p,
+             const int i) noexcept
+{
+    return f(p, i);
+}
 }
 /// \endcond
 

Src/Particle/AMReX_ParticleTile.H

@@ -329,6 +329,7 @@ struct alignas(sizeof(double)) ConstSoAParticle : SoAParticleBase
     using StorageParticleType = SoAParticleBase;
     using ConstPTD = ConstParticleTileData<SoAParticleBase, NArrayReal, NArrayInt>;
     static constexpr bool is_soa_particle = true;
+    static constexpr bool is_rtsoa_particle = false;
     static constexpr bool is_constsoa_particle = true;
 
     using RealType = ParticleReal;
@@ -392,6 +393,7 @@ struct alignas(sizeof(double)) SoAParticle : SoAParticleBase
     using StorageParticleType = SoAParticleBase;
     using PTD = ParticleTileData<SoAParticleBase, NArrayReal, NArrayInt>;
     static constexpr bool is_soa_particle = true;
+    static constexpr bool is_rtsoa_particle = false;
     static constexpr bool is_constsoa_particle = false;
 
     using ConstType = ConstSoAParticle<T_NArrayReal, T_NArrayInt>;