Merge pull request #6393 from STEllAR-GROUP/caching_allocator

Optionally disable caching allocator
STEllAR-GROUP · Dec 6, 2023 · 1cf717d · 1cf717d
2 parents 8d573fc + 8530848
commit 1cf717d
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 23 deletions.
diff --git a/libs/core/allocator_support/CMakeLists.txt b/libs/core/allocator_support/CMakeLists.txt
@@ -4,15 +4,35 @@
 # Distributed under the Boost Software License, Version 1.0. (See accompanying
 # file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
 
+# Allow to disable caching allocator
+hpx_option(
+  HPX_ALLOCATOR_SUPPORT_WITH_CACHING BOOL
+  "Enable caching allocator. (default: ON)" ON ADVANCED
+  CATEGORY "Modules"
+  MODULE ALLOCATOR_SUPPORT
+)
+
+if(HPX_ALLOCATOR_SUPPORT_WITH_CACHING)
+  hpx_add_config_define_namespace(
+    DEFINE HPX_ALLOCATOR_SUPPORT_HAVE_CACHING NAMESPACE ALLOCATOR_SUPPORT
+  )
+endif()
+
 set(allocator_support_headers
     hpx/allocator_support/aligned_allocator.hpp
     hpx/allocator_support/allocator_deleter.hpp
     hpx/allocator_support/detail/new.hpp
     hpx/allocator_support/internal_allocator.hpp
-    hpx/allocator_support/thread_local_caching_allocator.hpp
     hpx/allocator_support/traits/is_allocator.hpp
 )
 
+if(HPX_ALLOCATOR_SUPPORT_WITH_CACHING)
+  set(allocator_support_headers
+      ${allocator_support_headers}
+      hpx/allocator_support/thread_local_caching_allocator.hpp
+  )
+endif()
+
 # cmake-format: off
 set(allocator_support_compat_headers
     hpx/allocator_support.hpp => hpx/modules/allocator_support.hpp

diff --git a/libs/core/allocator_support/include/hpx/allocator_support/thread_local_caching_allocator.hpp b/libs/core/allocator_support/include/hpx/allocator_support/thread_local_caching_allocator.hpp
@@ -7,7 +7,9 @@
 #pragma once
 
 #include <hpx/config.hpp>
+#include <hpx/allocator_support/config/defines.hpp>
 
+#include <cstddef>
 #include <memory>
 #include <new>
 #include <stack>
@@ -16,8 +18,10 @@
 
 namespace hpx::util {
 
+#if defined(HPX_ALLOCATOR_SUPPORT_HAVE_CACHING) &&                             \
+    !((defined(HPX_HAVE_CUDA) && defined(__CUDACC__)) ||                       \
+        defined(HPX_HAVE_HIP))
     ///////////////////////////////////////////////////////////////////////////
-#if !((defined(HPX_HAVE_CUDA) && defined(__CUDACC__)) || defined(HPX_HAVE_HIP))
     template <typename T = char, typename Allocator = std::allocator<T>>
     struct thread_local_caching_allocator
     {
@@ -61,6 +65,44 @@ namespace hpx::util {
             allocated_cache& operator=(allocated_cache&&) = delete;
 
             ~allocated_cache()
+            {
+                clear_cache();
+            }
+
+            pointer allocate(size_type n)
+            {
+                pointer p;
+                if (data.empty())
+                {
+                    p = traits::allocate(alloc, n);
+                    if (p == nullptr)
+                    {
+                        throw std::bad_alloc();
+                    }
+                }
+                else
+                {
+                    p = data.top().first;
+                    data.pop();
+                }
+
+                ++allocated;
+                return p;
+            }
+
+            void deallocate(pointer p, size_type n) noexcept
+            {
+                data.push(std::make_pair(p, n));
+                if (++deallocated > 2 * (allocated + 16))
+                {
+                    clear_cache();
+                    allocated = 0;
+                    deallocated = 0;
+                }
+            }
+
+        private:
+            void clear_cache() noexcept
             {
                 while (!data.empty())
                 {
@@ -72,12 +114,14 @@ namespace hpx::util {
 
             HPX_NO_UNIQUE_ADDRESS Allocator alloc;
             std::stack<std::pair<T*, size_type>> data;
+            std::size_t allocated = 0;
+            std::size_t deallocated = 0;
         };
 
-        std::stack<std::pair<T*, size_type>>& cache()
+        allocated_cache& cache()
         {
             thread_local allocated_cache allocated_data(alloc);
-            return allocated_data.data;
+            return allocated_data;
         }
 
     public:
@@ -114,29 +158,12 @@ namespace hpx::util {
             {
                 throw std::bad_array_new_length();
             }
-
-            pointer p;
-
-            if (auto& c = cache(); c.empty())
-            {
-                p = traits::allocate(alloc, n);
-                if (p == nullptr)
-                {
-                    throw std::bad_alloc();
-                }
-            }
-            else
-            {
-                p = c.top().first;
-                c.pop();
-            }
-
-            return p;
+            return cache().allocate(n);
         }
 
         void deallocate(pointer p, size_type n) noexcept
         {
-            cache().push(std::make_pair(p, n));
+            cache().deallocate(p, n);
         }
 
         [[nodiscard]] constexpr size_type max_size() noexcept