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; // It's always 0 because this function is for RedistributeGPU only and the tiling is off.
+            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,138 @@ struct ParticleTile
         return ptd;
     }
 
+    void collectVectors (Vector<ParticleVector*>& pv,
+                         Vector<typename SoA::IdCPU*>& idcpuv,
+                         Vector<RealVector*>& rv, Vector<IntVector*>& iv)
+    {
+        if constexpr (!ParticleType::is_soa_particle) {
+            m_aos_tile.collectVectors(pv);
+        } else {
+            amrex::ignore_unused(pv);
+        }
+        m_soa_tile.collectVectors(idcpuv, rv, iv);
+    }
+
+    static void reserve (std::map<ParticleTile<T_ParticleType,NArrayReal,NArrayInt,Allocator>*, int> const& addsizes)
+    {
+        if constexpr (!IsArenaAllocator<Allocator<int>>::value) {
+            for (auto [p,s] : addsizes) {
+                p->reserve(p->size()+s);
+            }
+        } else {
+            using PV = std::conditional_t<ParticleType::is_soa_particle,
+                                          IntVector, ParticleVector>;
+            Vector<std::pair<PV*,int>> pvs;
+            Vector<std::pair<typename SoA::IdCPU*,int>> ids;
+            Vector<std::pair<RealVector*,int>> rvs;
+            Vector<std::pair<IntVector*,int>> ivs;
+            for (auto [p,s] : addsizes) {
+                if (s > 0) {
+                    Vector<ParticleVector*> pv;
+                    Vector<typename SoA::IdCPU*> idcpuv;
+                    Vector<RealVector*> rv;
+                    Vector<IntVector*> iv;
+                    p->collectVectors(pv, idcpuv, rv, iv);
+                    if constexpr (!ParticleType::is_soa_particle) {
+                        for (auto* v : pv) {
+                            pvs.emplace_back(v, s);
+                        }
+                    }
+                    for (auto* v : idcpuv) {
+                        ids.emplace_back(v, s);
+                    }
+                    for (auto* v : rv) {
+                        rvs.emplace_back(v, s);
+                    }
+                    for (auto* v : iv) {
+                        ivs.emplace_back(v, s);
+                    }
+                }
+            }
+
+            // xxxxx TODO: Can we come up a better strategy?
+
+            std::sort(pvs.begin(), pvs.end(), [] (auto const& a, auto const& b) {
+                return (a.first->size() + a.second) >
+                       (b.first->size() + b.second);
+            });
+            std::sort(ids.begin(), ids.end(), [] (auto const& a, auto const& b) {
+                return (a.first->size() + a.second) >
+                       (b.first->size() + b.second);
+            });
+            std::sort(rvs.begin(), rvs.end(), [] (auto const& a, auto const& b) {
+                return (a.first->size() + a.second) >
+                       (b.first->size() + b.second);
+            });
+            std::sort(ivs.begin(), ivs.end(), [] (auto const& a, auto const& b) {
+                return (a.first->size() + a.second) >
+                       (b.first->size() + b.second);
+            });
+
+            // Handle big vectcors first
+            {
+                int i_pvs = 0, i_ids = 0, i_rvs = 0, i_ivs = 0;
+                auto n_pvs = int(pvs.size());
+                auto n_ids = int(ids.size());
+                auto n_rvs = int(rvs.size());
+                auto n_ivs = int(ivs.size());
+                while ((i_pvs < n_pvs) || (i_ids < n_ids) || (i_rvs < n_rvs) ||
+                       (i_ivs < n_ivs)) {
+                    std::size_t nbytes = 0;
+                    int ii = -1;
+                    if (i_pvs < n_pvs) {
+                        std::size_t my_bytes = (pvs[i_pvs].first->size()
+                                              + pvs[i_pvs].second) * sizeof(typename PV::value_type);
+                        if (my_bytes > nbytes) {
+                            nbytes = my_bytes;
+                            ii = 0;
+                        }
+                    }
+                    if (i_ids < n_ids) {
+                        std::size_t my_bytes = (ids[i_ids].first->size()
+                                              + ids[i_ids].second) * sizeof(typename SoA::IdCPU::value_type);
+                        if (my_bytes > nbytes) {
+                            nbytes = my_bytes;
+                            ii = 1;
+                        }
+                    }
+                    if (i_rvs < n_rvs) {
+                        std::size_t my_bytes = (rvs[i_rvs].first->size()
+                                              + rvs[i_rvs].second) * sizeof(typename RealVector::value_type);
+                        if (my_bytes > nbytes) {
+                            nbytes = my_bytes;
+                            ii = 2;
+                        }
+                    }
+                    if (i_ivs < n_ivs) {
+                        std::size_t my_bytes = (ivs[i_ivs].first->size()
+                                              + ivs[i_ivs].second) * sizeof(typename IntVector::value_type);
+                        if (my_bytes > nbytes) {
+                            nbytes = my_bytes;
+                            ii = 3;
+                        }
+                    }
+                    if (ii == 0) {
+                        auto [p,s] = pvs[i_pvs++];
+                        p->reserve(p->size() + s);
+                    } else if (ii == 1) {
+                        auto [p,s] = ids[i_ids++];
+                        p->reserve(p->size() + s);
+                    } else if (ii == 2) {
+                        auto [p,s] = rvs[i_rvs++];
+                        p->reserve(p->size() + s);
+                    } else {
+                        auto [p,s] = ivs[i_ivs++];
+                        p->reserve(p->size() + s);
+                    }
+                }
+            }
+
+            // If we fail, we should try to free up some space and try again
+
+        }
+    }
+
 private:
 
     AoS m_aos_tile;

Src/Particle/AMReX_StructOfArrays.H

@@ -325,6 +325,29 @@ struct StructOfArrays {
         return arr;
     }
 
+    void collectVectors (Vector<IdCPU*>& idcpuv, Vector<RealVector*>& rv,
+                         Vector<IntVector*>& iv)
+
+    {
+        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: