From 85eb18afe74ca208582cfa517181668840fd01b1 Mon Sep 17 00:00:00 2001
From: Weiqun Zhang <weiqunzhang@lbl.gov>
Date: Mon, 17 Jul 2023 18:26:54 -0700
Subject: [PATCH] CArena: Implement alloc_in_place

First we try to see if CArena has free space that can be used to extend the
given pointer to the specified size range. If that's unsuccessful, we use
the normal alloc function to allocate memory.

If in-place allocation is successful during vector resizing, we can avoid
the cost of copying data to the new allocation address.
---
 Src/Base/AMReX_CArena.H   | 14 ++++++++
 Src/Base/AMReX_CArena.cpp | 76 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/Src/Base/AMReX_CArena.H b/Src/Base/AMReX_CArena.H
index ee081d88f82..92147ecfcf5 100644
--- a/Src/Base/AMReX_CArena.H
+++ b/Src/Base/AMReX_CArena.H
@@ -47,6 +47,18 @@ public:
     //! Allocate some memory.
     [[nodiscard]] void* alloc (std::size_t nbytes) final;
 
+    /**
+     * Try to allocate in-place by extending the capacity of given pointer.
+     */
+    [[nodiscard]] std::pair<void*,std::size_t>
+    alloc_in_place (void* pt, std::size_t szmin, std::size_t szmax) final;
+
+    /**
+     * Try to shrink in-place
+     */
+    [[nodiscard]] void*
+    shrink_in_place (void* pt, std::size_t sz) final;
+
     /**
     * \brief Free up allocated memory.  Merge neighboring free memory chunks
     * into largest possible chunk.
@@ -87,6 +99,8 @@ public:
 
 protected:
 
+    void* alloc_protected (std::size_t nbytes);
+
     std::size_t freeUnused_protected () final;
 
     //! The nodes in our free list and block list.
diff --git a/Src/Base/AMReX_CArena.cpp b/Src/Base/AMReX_CArena.cpp
index 8c5c0dc7e18..eb7c246a126 100644
--- a/Src/Base/AMReX_CArena.cpp
+++ b/Src/Base/AMReX_CArena.cpp
@@ -42,9 +42,13 @@ void*
 CArena::alloc (std::size_t nbytes)
 {
     std::lock_guard<std::mutex> lock(carena_mutex);
-
     nbytes = Arena::align(nbytes == 0 ? 1 : nbytes);
+    return alloc_protected(nbytes);
+}
 
+void*
+CArena::alloc_protected (std::size_t nbytes)
+{
     MemStat* stat = nullptr;
 #ifdef AMREX_TINY_PROFILING
     if (m_do_profiling) {
@@ -127,6 +131,76 @@ CArena::alloc (std::size_t nbytes)
     return vp;
 }
 
+std::pair<void*,std::size_t>
+CArena::alloc_in_place (void* pt, std::size_t szmin, std::size_t szmax)
+{
+    std::lock_guard<std::mutex> lock(carena_mutex);
+
+    std::size_t nbytes_max = Arena::align(szmax == 0 ? 1 : szmax);
+
+    if (pt != nullptr) { // Try to allocate in-place first
+        auto busy_it = m_busylist.find(Node(pt,nullptr,0));
+        AMREX_ALWAYS_ASSERT(busy_it != m_busylist.end());
+        AMREX_ASSERT(m_freelist.find(*busy_it) == m_freelist.end());
+
+        if (busy_it->size() >= szmax) {
+            return std::make_pair(pt, busy_it->size());
+        }
+
+        void* next_block = (char*)pt + busy_it->size();
+        auto next_it = m_freelist.find(Node(next_block,nullptr,0));
+        if (next_it != m_freelist.end() && busy_it->coalescable(*next_it)) {
+            std::size_t total_size = busy_it->size() + next_it->size();
+            if (total_size >= szmax) {
+                // Must use nbytes_max instead of szmax for alignment.
+                std::size_t new_size = std::min(total_size, nbytes_max);
+                std::size_t left_size = total_size - new_size;
+                if (left_size <= 64) {
+                    m_freelist.erase(next_it);
+                    new_size = total_size;
+                } else {
+                    auto& free_node = const_cast<Node&>(*next_it);
+                    free_node.block((char*)pt + new_size);
+                    free_node.size(left_size);
+                }
+                std::size_t extra_size = new_size - busy_it->size();
+#ifdef AMREX_TINY_PROFILING
+                if (m_do_profiling) {
+                    TinyProfiler::memory_alloc(extra_size, m_profiling_stats);
+                }
+#endif
+                m_actually_used += extra_size;
+                const_cast<Node&>(*busy_it).size(new_size);
+                return std::make_pair(pt, new_size);
+            } else if (total_size >= szmin) {
+                m_freelist.erase(next_it);
+                std::size_t extra_size = total_size - busy_it->size();
+#ifdef AMREX_TINY_PROFILING
+                if (m_do_profiling) {
+                    TinyProfiler::memory_alloc(extra_size, m_profiling_stats);
+                }
+#endif
+                m_actually_used += extra_size;
+                const_cast<Node&>(*busy_it).size(total_size);
+                return std::make_pair(pt, total_size);
+            }
+        }
+
+        if (busy_it->size() >= szmin) {
+            return std::make_pair(pt, busy_it->size());
+        }
+    }
+
+    void* newp = alloc_protected(nbytes_max);
+    return std::make_pair(newp, nbytes_max);
+}
+
+void*
+CArena::shrink_in_place (void* /*pt*/, std::size_t sz)
+{
+    return alloc(sz); // xxxxx TODO
+}
+
 void
 CArena::free (void* vp)
 {