Add GrowthStrategy argument to PODVector resize and reserve (#4763)

Follow-up to #4734

Author: AlexanderSinn

Src/Base/AMReX_PODVector.H

@@ -4,6 +4,7 @@
 
 #include <AMReX.H>
 #include <AMReX_Arena.H>
+#include <AMReX_Enum.H>
 #include <AMReX_GpuLaunch.H>
 #include <AMReX_GpuAllocators.H>
 #include <AMReX_GpuDevice.H>
@@ -244,6 +245,8 @@ namespace amrex
         }
     }
 
+    AMREX_ENUM(GrowthStrategy, Poisson, Exact, Geometric);
+
     namespace VectorGrowthStrategy
     {
         extern AMREX_EXPORT Real growth_factor;
@@ -258,17 +261,35 @@ namespace amrex
         void Initialize ();
     }
 
-    inline std::size_t grow_podvector_capacity (std::size_t s)
+    inline std::size_t grow_podvector_capacity (GrowthStrategy strategy, std::size_t new_size,
+                                                std::size_t old_capacity, std::size_t sizeof_T)
     {
-        if (s <= 900) {
-            // 3*sqrt(900) = 900/10. Note that we don't need to be precise
-            // here. Even if later we change the else block to
-            // 4*std::sqrt(s), it's not really an issue to still use 900
-            // here.
-            return s + s/10;
-        } else {
-            return s + std::size_t(3*std::sqrt(s));
+        switch (strategy) {
+            case GrowthStrategy::Poisson:
+                if (new_size <= 900) {
+                    // 3*sqrt(900) = 900/10. Note that we don't need to be precise
+                    // here. Even if later we change the else block to
+                    // 4*std::sqrt(s), it's not really an issue to still use 900
+                    // here.
+                    return new_size + new_size/10;
+                } else {
+                    return new_size + std::size_t(3*std::sqrt(new_size));
+                }
+            case GrowthStrategy::Exact:
+                return new_size;
+            case GrowthStrategy::Geometric:
+                if (old_capacity == 0) {
+                    return std::max(64/sizeof_T, new_size);
+                } else {
+                    Real const gf = VectorGrowthStrategy::GetGrowthFactor();
+                    if (amrex::almostEqual(gf, Real(1.5))) {
+                        return std::max((old_capacity*3+1)/2, new_size);
+                    } else {
+                        return std::max(std::size_t(gf*Real(old_capacity+1)), new_size);
+                    }
+                }
         }
+        return 0; // unreachable
     }
 
     template <class T, class Allocator = std::allocator<T> >
@@ -309,7 +330,8 @@ namespace amrex
         {}
 
         explicit PODVector (size_type a_size)
-            : m_size(a_size), m_capacity(grow_podvector_capacity(a_size))
+            : m_size(a_size),
+              m_capacity(grow_podvector_capacity(GrowthStrategy::Poisson, a_size, 0, sizeof(T)))
         {
             if (m_capacity != 0) {
                 m_data = allocate(m_capacity);
@@ -322,7 +344,7 @@ namespace amrex
         PODVector (size_type a_size, const value_type& a_value,
                    const allocator_type& a_allocator = Allocator())
             : Allocator(a_allocator), m_size(a_size),
-              m_capacity(grow_podvector_capacity(a_size))
+              m_capacity(grow_podvector_capacity(GrowthStrategy::Poisson, a_size, 0, sizeof(T)))
         {
             if (m_capacity != 0) {
                 m_data = allocate(m_capacity);
@@ -337,7 +359,8 @@ namespace amrex
                    const allocator_type& a_allocator = Allocator())
             : Allocator(a_allocator),
               m_size    (a_initializer_list.size()),
-              m_capacity(grow_podvector_capacity(a_initializer_list.size()))
+              m_capacity(grow_podvector_capacity(
+                GrowthStrategy::Poisson, a_initializer_list.size(), 0, sizeof(T)))
         {
             if (m_capacity != 0) {
                 m_data = allocate(m_capacity);
@@ -661,20 +684,82 @@ namespace amrex
 
         [[nodiscard]] const_reverse_iterator crend () const noexcept { return const_reverse_iterator(begin()); }
 
-        void resize (size_type a_new_size)
+        /** Resizes the PODVector to a_new_size elements. New elements are either uninitialized or,
+         * if amrex.init_snan is set to true, initialized to NaN.
+         * To minimize reallocations, the capacity of the vector may be increased to a value larger
+         * than the requested size. This behavior can be controlled using the strategy parameter.
+         * Available strategies are:
+         *
+         * - GrowthStrategy::Poisson (Default)
+         *      Sets the new capacity to a_new_size + 3*sqrt(a_new_size).
+         *      This is useful for vectors that store particle data, with particles moving between
+         *      tiles due to a thermal process. In this case the expected number of particles
+         *      per tile follows a Poisson distribution, which has a standard deviation of sqrt(n).
+         *      Adding three sigma to the size greatly minimizes reallocations in Redistribute()
+         *      of successive timesteps.
+         *
+         * - GrowthStrategy::Exact
+         *      Sets the new capacity exactly to the requested size.
+         *      Can be used if the PODVector will never be resized to a bigger size.
+         *
+         * - GrowthStrategy::Geometric
+         *      Sets the new capacity to max(growth_factor * old_capacity, a_new_size),
+         *      where growth_factor is 1.5 by default and can be adjusted with
+         *      amrex.vector_growth_factor. This strategy is needed to guarantee O(n) runtime
+         *      if the PODVector is resized in a loop to consecutively increasing sizes up to n.
+         *      In contrast, the Poisson and Exact strategies would have O(n*sqrt(n)) and O(n^2)
+         *      runtimes, respectively. However, it is not recommended to use a PODVector this way.
+         *      Instead, the reserve function should be used once upfront.
+         *
+         * \param[in] a_new_size the new size of the PODVector
+         * \param[in] strategy the growth strategy to set the new capacity
+         */
+        void resize (size_type a_new_size,
+                     GrowthStrategy strategy = GrowthStrategy::Poisson)
         {
             auto old_size = m_size;
-            resize_without_init_snan(a_new_size);
+            resize_without_init_snan(a_new_size, strategy);
             if (old_size < a_new_size) {
                 detail::maybe_init_snan(m_data + old_size,
                                         m_size - old_size, (Allocator const&)(*this));
             }
         }
 
-        void resize (size_type a_new_size, const T& a_val)
+        /** Resizes the PODVector to a_new_size elements. New elements are set to a_val.
+         * To minimize reallocations, the capacity of the vector may be increased to a value larger
+         * than the requested size. This behavior can be controlled using the strategy parameter.
+         * Available strategies are:
+         *
+         * - GrowthStrategy::Poisson (Default)
+         *      Sets the new capacity to a_new_size + 3*sqrt(a_new_size).
+         *      This is useful for vectors that store particle data, with particles moving between
+         *      tiles due to a thermal process. In this case the expected number of particles
+         *      per tile follows a Poisson distribution, which has a standard deviation of sqrt(n).
+         *      Adding three sigma to the size greatly minimizes reallocations in Redistribute()
+         *      of successive timesteps.
+         *
+         * - GrowthStrategy::Exact
+         *      Sets the new capacity exactly to the requested size.
+         *      Can be used if the PODVector will never be resized to a bigger size.
+         *
+         * - GrowthStrategy::Geometric
+         *      Sets the new capacity to max(growth_factor * old_capacity, a_new_size),
+         *      where growth_factor is 1.5 by default and can be adjusted with
+         *      amrex.vector_growth_factor. This strategy is needed to guarantee O(n) runtime
+         *      if the PODVector is resized in a loop to consecutively increasing sizes up to n.
+         *      In contrast, the Poisson and Exact strategies would have O(n*sqrt(n)) and O(n^2)
+         *      runtimes, respectively. However, it is not recommended to use a PODVector this way.
+         *      Instead, the reserve function should be used once upfront.
+         *
+         * \param[in] a_new_size the new size of the PODVector
+         * \param[in] a_val the value of the new elements
+         * \param[in] strategy the growth strategy to set the new capacity
+         */
+        void resize (size_type a_new_size, const T& a_val,
+                     GrowthStrategy strategy = GrowthStrategy::Poisson)
         {
             size_type old_size = m_size;
-            resize_without_init_snan(a_new_size);
+            resize_without_init_snan(a_new_size, strategy);
             if (old_size < a_new_size)
             {
                 detail::uninitializedFillNImpl(m_data + old_size,
@@ -683,10 +768,39 @@ namespace amrex
             }
         }
 
-        void reserve (size_type a_capacity)
+        /** Sets the capacity of the PODVector to at least a_capacity without changing the size.
+         * To minimize reallocations, the capacity of the vector may be increased to a value larger
+         * than the requested one. This behavior can be controlled using the strategy parameter.
+         * Available strategies are:
+         *
+         * - GrowthStrategy::Poisson (Default)
+         *      Sets the new capacity to a_capacity + 3*sqrt(a_capacity).
+         *      This is useful for vectors that store particle data, with particles moving between
+         *      tiles due to a thermal process. In this case the expected number of particles
+         *      per tile follows a Poisson distribution, which has a standard deviation of sqrt(n).
+         *      Adding three sigma to the size greatly minimizes reallocations in Redistribute()
+         *      of successive timesteps.
+         *
+         * - GrowthStrategy::Exact
+         *      Sets the new capacity exactly to the requested a_capacity.
+         *      Can be used if the PODVector will never be resized to a bigger size.
+         *
+         * - GrowthStrategy::Geometric
+         *      Sets the new capacity to max(growth_factor * old_capacity, a_capacity),
+         *      where growth_factor is 1.5 by default and can be adjusted with
+         *      amrex.vector_growth_factor. This strategy is needed to guarantee O(n) runtime
+         *      if the PODVector is resized in a loop to consecutively increasing sizes up to n.
+         *      In contrast, the Poisson and Exact strategies would have O(n*sqrt(n)) and O(n^2)
+         *      runtimes, respectively. However, it is not recommended to use a PODVector this way.
+         *      Instead, the reserve function should be used once upfront.
+         *
+         * \param[in] a_capacity the new minimum capacity of the PODVector
+         * \param[in] strategy the growth strategy to set the new capacity
+         */
+        void reserve (size_type a_capacity, GrowthStrategy strategy = GrowthStrategy::Poisson)
         {
             if (m_capacity < a_capacity) {
-                reserve_doit(grow_podvector_capacity(a_capacity));
+                reserve_doit(grow_podvector_capacity(strategy, a_capacity, m_capacity, sizeof(T)));
             }
         }
 
@@ -738,16 +852,8 @@ namespace amrex
         // this is where we would change the growth strategy for push_back
         [[nodiscard]] size_type GetNewCapacityForPush () const noexcept
         {
-            if (m_capacity == 0) {
-                return std::max(64/sizeof(T), size_type(1));
-            } else {
-                Real const gf = VectorGrowthStrategy::GetGrowthFactor();
-                if (amrex::almostEqual(gf, Real(1.5))) {
-                    return (m_capacity*3+1)/2;
-                } else {
-                    return size_type(gf*Real(m_capacity+1));
-                }
-            }
+            return grow_podvector_capacity(GrowthStrategy::Geometric, m_capacity + 1,
+                                           m_capacity, sizeof(T));
         }
 
         void UpdateDataPtr (FatPtr<T> const& fp)
@@ -817,10 +923,10 @@ namespace amrex
             m_capacity = new_capacity;
         }
 
-        void resize_without_init_snan (size_type a_new_size)
+        void resize_without_init_snan (size_type a_new_size, GrowthStrategy strategy)
         {
             if (m_capacity < a_new_size) {
-                reserve(a_new_size);
+                reserve(a_new_size, strategy);
             }
             m_size = a_new_size;
         }

Src/Particle/AMReX_ArrayOfStructs.H

@@ -91,9 +91,13 @@ public:
         m_data.swap(other.m_data);
     }
 
-    void resize (size_t count) { m_data.resize(count); }
+    void resize (size_t count, GrowthStrategy strategy = GrowthStrategy::Poisson) {
+        m_data.resize(count, strategy);
+    }
 
-    void reserve (size_t capacity) { m_data.reserve(capacity); }
+    void reserve (size_t capacity, GrowthStrategy strategy = GrowthStrategy::Poisson) {
+        m_data.reserve(capacity, strategy);
+    }
 
     Iterator erase ( ConstIterator first, ConstIterator second) { return m_data.erase(first, second); }
 

Src/Particle/AMReX_ParticleTile.H

@@ -965,20 +965,20 @@ struct ParticleTile
         }
     }
 
-    void resize (std::size_t count)
+    void resize (std::size_t count, GrowthStrategy strategy = GrowthStrategy::Poisson)
     {
         if constexpr (!ParticleType::is_soa_particle) {
-            m_aos_tile.resize(count);
+            m_aos_tile.resize(count, strategy);
         }
-        m_soa_tile.resize(count);
+        m_soa_tile.resize(count, strategy);
     }
 
-    void reserve (std::size_t capacity)
+    void reserve (std::size_t capacity, GrowthStrategy strategy = GrowthStrategy::Poisson)
     {
         if constexpr (!ParticleType::is_soa_particle) {
-            m_aos_tile.reserve(capacity);
+            m_aos_tile.reserve(capacity, strategy);
         }
-        m_soa_tile.reserve(capacity);
+        m_soa_tile.reserve(capacity, strategy);
     }
 
     ///

Src/Particle/AMReX_StructOfArrays.H

@@ -267,34 +267,42 @@ struct StructOfArrays {
 
     [[nodiscard]] int getNumNeighbors () const { return m_num_neighbor_particles; }
 
-    void resize (size_t count)
+    void resize (size_t count, GrowthStrategy strategy = GrowthStrategy::Poisson)
     {
         if constexpr (use64BitIdCpu == true) {
-            m_idcpu.resize(count);
+            m_idcpu.resize(count, strategy);
         }
         if constexpr (NReal > 0) {
-            for (int i = 0; i < NReal; ++i) { m_rdata[i].resize(count); }
+            for (int i = 0; i < NReal; ++i) { m_rdata[i].resize(count, strategy); }
         }
         if constexpr (NInt > 0) {
-            for (int i = 0; i < NInt;  ++i) { m_idata[i].resize(count); }
+            for (int i = 0; i < NInt;  ++i) { m_idata[i].resize(count, strategy); }
+        }
+        for (int i = 0; i < int(m_runtime_rdata.size()); ++i) {
+            m_runtime_rdata[i].resize(count, strategy);
+        }
+        for (int i = 0; i < int(m_runtime_idata.size()); ++i) {
+            m_runtime_idata[i].resize(count, strategy);
         }
-        for (int i = 0; i < int(m_runtime_rdata.size()); ++i) { m_runtime_rdata[i].resize(count); }
-        for (int i = 0; i < int(m_runtime_idata.size()); ++i) { m_runtime_idata[i].resize(count); }
     }
 
-    void reserve (size_t capacity)
+    void reserve (size_t capacity, GrowthStrategy strategy = GrowthStrategy::Poisson)
     {
         if constexpr (use64BitIdCpu == true) {
-            m_idcpu.reserve(capacity);
+            m_idcpu.reserve(capacity, strategy);
         }
         if constexpr (NReal > 0) {
-            for (int i = 0; i < NReal; ++i) { m_rdata[i].reserve(capacity); }
+            for (int i = 0; i < NReal; ++i) { m_rdata[i].reserve(capacity, strategy); }
         }
         if constexpr (NInt > 0) {
-            for (int i = 0; i < NInt;  ++i) { m_idata[i].reserve(capacity); }
+            for (int i = 0; i < NInt;  ++i) { m_idata[i].reserve(capacity, strategy); }
+        }
+        for (int i = 0; i < int(m_runtime_rdata.size()); ++i) {
+            m_runtime_rdata[i].reserve(capacity, strategy);
+        }
+        for (int i = 0; i < int(m_runtime_idata.size()); ++i) {
+            m_runtime_idata[i].reserve(capacity, strategy);
         }
-        for (int i = 0; i < int(m_runtime_rdata.size()); ++i) { m_runtime_rdata[i].reserve(capacity); }
-        for (int i = 0; i < int(m_runtime_idata.size()); ++i) { m_runtime_idata[i].reserve(capacity); }
     }
 
     [[nodiscard]] uint64_t* idcpuarray () {