Reserve space in particle redistribution

Src/Particle/AMReX_ArrayOfStructs.H

@@ -108,6 +108,11 @@ public:
     [[nodiscard]] typename ParticleVector::const_iterator end () const { return m_data.end(); }
     [[nodiscard]] typename ParticleVector::const_iterator cend () const { return m_data.cend(); }
 
+    void collectVectors (Vector<ParticleVector*>& pvs)
+    {
+        pvs.push_back(&m_data);
+    }
+
     int m_num_neighbor_particles{0};
 
 private:

Src/Particle/AMReX_ParticleContainer.H

@@ -1275,6 +1275,8 @@ public:
     void RedistributeGPU (int lev_min = 0, int lev_max = -1, int nGrow = 0, int local=0,
                           bool remove_negative=true);
 
+    void ReserveForRedistribute (ParticleCopyPlan const& plan);
+
     Long superParticleSize() const { return superparticle_size; }
 
     void AddRealComp (std::string const & name, int communicate=1)

Src/Particle/AMReX_ParticleContainerI.H

@@ -1587,6 +1587,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     {
         plan.buildMPIFinish(BufferMap());
         communicateParticlesStart(*this, plan, snd_buffer, rcv_buffer);
+        this->ReserveForRedistribute(plan);
         unpackBuffer(*this, plan, snd_buffer, RedistributeUnpackPolicy());
         communicateParticlesFinish(plan);
         unpackRemotes(*this, plan, rcv_buffer, RedistributeUnpackPolicy());
@@ -1609,6 +1610,8 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
             communicateParticlesStart(*this, plan, pinned_snd_buffer, pinned_rcv_buffer);
         }
 
+        this->ReserveForRedistribute(plan);
+
         rcv_buffer.resize(pinned_rcv_buffer.size());
         unpackBuffer(*this, plan, snd_buffer, RedistributeUnpackPolicy());
         communicateParticlesFinish(plan);
@@ -1623,6 +1626,42 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
 #endif
 }
 
+template <typename ParticleType, int NArrayReal, int NArrayInt,
+          template<class> class Allocator, class CellAssignor>
+void
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>
+::ReserveForRedistribute (ParticleCopyPlan const& plan)
+{
+    BL_PROFILE("ParticleContainer::ReserveForRedistribute()");
+
+    std::map<ParticleTileType*, int> addsizes;
+
+    for (int lev = 0; lev < this->BufferMap().numLevels(); ++lev) {
+        for (MFIter mfi = this->MakeMFIter(lev); mfi.isValid(); ++mfi) {
+            int gid = mfi.index();
+            int tid = mfi.LocalTileIndex();
+            auto& tile = this->DefineAndReturnParticleTile(lev, gid, tid);
+            int num_copies = plan.m_box_counts_h[this->BufferMap().gridAndLevToBucket(gid, lev)];
+            if (num_copies > 0) {
+                addsizes[&tile] += num_copies;
+            }
+        }
+    }
+
+    if (plan.m_nrcvs > 0) {
+        for (int i = 0, N = int(plan.m_rcv_box_counts.size()); i < N; ++i) {
+            int copy_size = plan.m_rcv_box_counts[i];
+            int lev = plan.m_rcv_box_levs[i];
+            int gid = plan.m_rcv_box_ids[i];
+            int tid = 0; // xxxxx why is tid 0?
+            auto& tile = this->DefineAndReturnParticleTile(lev, gid, tid);
+            addsizes[&tile] += copy_size;
+        }
+    }
+
+    ParticleTileType::reserve(addsizes);
+}
+
 //
 // The CPU implementation of Redistribute
 //

Src/Particle/AMReX_ParticleTile.H

@@ -735,7 +735,7 @@ struct ParticleTile
         ParticleType::is_soa_particle,
         ThisParticleTileHasNoAoS,
         ArrayOfStructs<ParticleType, Allocator>>;
-    //using ParticleVector = typename AoS::ParticleVector;
+    using ParticleVector = typename AoS::ParticleVector;
 
     using SoA = std::conditional_t<
         ParticleType::is_soa_particle,
@@ -1351,6 +1351,58 @@ struct ParticleTile
         return ptd;
     }
 
+    void collectVectors (Vector<ParticleVector*>& pv, Vector<RealVector*>& rv,
+                         Vector<IntVector*>& iv, Vector<typename SoA::IdCPU*>& idcpuv)
+    {
+        if constexpr (!ParticleType::is_soa_particle) {
+            m_aos_tile.collectVectors(pv);
+        } else {
+            amrex::ignore_unused(pv);
+        }
+        m_soa_tile.collectVectors(rv, iv, idcpuv);
+    }
+
+    static void reserve (std::map<ParticleTile<T_ParticleType,NArrayReal,NArrayInt,Allocator>*, int> const& addsizes)
+    {
+        Vector<std::pair<ParticleVector*,int>> pvs;
+        Vector<std::pair<RealVector*,int>> rvs;
+        Vector<std::pair<IntVector*,int>> ivs;
+        Vector<std::pair<typename SoA::IdCPU*,int>> idcpuvs;
+        for (auto [p,s] : addsizes) {
+            Vector<ParticleVector*> pv;
+            Vector<RealVector*> rv;
+            Vector<IntVector*> iv;
+            Vector<typename SoA::IdCPU*> idcpuv;
+            p->collectVectors(pv, rv, iv, idcpuv);
+            for (auto* v : pv) {
+                pvs.emplace_back(v, s);
+            }
+            for (auto* v : rv) {
+                rvs.emplace_back(v, s);
+            }
+            for (auto* v : iv) {
+                ivs.emplace_back(v, s);
+            }
+            for (auto* v : idcpuv) {
+                idcpuvs.emplace_back(v, s);
+            }
+        }
+
+        // xxxxx TODO: Can we come up a better strategy?
+        for (auto [p,s] : pvs) {
+            p->reserve(p->size() + s);
+        }
+        for (auto [p,s] : rvs) {
+            p->reserve(p->size() + s);
+        }
+        for (auto [p,s] : ivs) {
+            p->reserve(p->size() + s);
+        }
+        for (auto [p,s] : idcpuvs) {
+            p->reserve(p->size() + s);
+        }
+    }
+
 private:
 
     AoS m_aos_tile;

Src/Particle/AMReX_StructOfArrays.H

@@ -325,6 +325,28 @@ struct StructOfArrays {
         return arr;
     }
 
+    void collectVectors (Vector<RealVector*>& rv, Vector<IntVector*>& iv,
+                         Vector<IdCPU*>& idcpuv)
+    {
+        if constexpr (use64BitIdCpu == true) {
+            idcpuv.push_back(&m_idcpu);
+        } else {
+            amrex::ignore_unused(idcpuv);
+        }
+        if constexpr (NReal > 0) {
+            for (int i = 0; i < NReal; ++i) { rv.push_back(&(m_rdata[i])); }
+        }
+        if constexpr (NInt > 0) {
+            for (int i = 0; i < NInt;  ++i) { iv.push_back(&(m_idata[i])); }
+        }
+        for (int i = 0; i < int(m_runtime_rdata.size()); ++i) {
+            rv.push_back(&(m_runtime_rdata[i]));
+        }
+        for (int i = 0; i < int(m_runtime_idata.size()); ++i) {
+            iv.push_back(&(m_runtime_idata[i]));
+        }
+    }
+
     int m_num_neighbor_particles{0};
 
 private: