Merge pull request #49 from msimberg/heap-config

Make heap segment limits and sizes configurable

Author: boeschf

include/hwmalloc/heap.hpp

@@ -1,7 +1,7 @@
 /*
  * ghex-org
  *
- * Copyright (c) 2014-2023, ETH Zurich
+ * Copyright (c) 2014-2025, ETH Zurich
  * All rights reserved.
  *
  * Please, refer to the LICENSE file in the root directory.
@@ -14,6 +14,7 @@
 #include <hwmalloc/fancy_ptr/void_ptr.hpp>
 #include <hwmalloc/fancy_ptr/const_void_ptr.hpp>
 #include <hwmalloc/fancy_ptr/unique_ptr.hpp>
+#include <hwmalloc/heap_config.hpp>
 #include <hwmalloc/allocator.hpp>
 #include <vector>
 #include <unordered_map>
@@ -56,6 +57,9 @@ class heap
     template<typename T>
     using unique_ptr = unique_ptr<T, block_type>;
 
+    // Note: sizes below are defaults and can be changed through heap_config and
+    // environment variables.
+    //
     // There are 5 size classes that the heap uses. For each size class it relies on a
     // fixed_size_heap. The size classes are:
     // - tiny:  heaps with linearly increasing block sizes, each heap backed by 16KiB segments
@@ -98,78 +102,56 @@ class heap
     //    :                                                                :  m_huge_heaps: map
 
   private:
-    static constexpr std::size_t log2_c(std::size_t n) noexcept
+    static std::size_t tiny_bucket_index(std::size_t n, std::size_t tiny_increment,
+        std::size_t tiny_increment_shift) noexcept
     {
-        return ((n < 2) ? 1 : 1 + log2_c(n >> 1));
+        return ((n + tiny_increment - 1) >> tiny_increment_shift) - 1;
     }
 
-    static const std::size_t s_tiny_limit = (1u << 7);   //   128
-    static const std::size_t s_small_limit = (1u << 10); //  1024
-    static const std::size_t s_large_limit = (1u << 16); // 65536
-
-    static const std::size_t s_bucket_shift = log2_c(s_tiny_limit) - 1;
-
-    static const std::size_t s_tiny_segment = 0x04000;  // 16KiB
-    static const std::size_t s_small_segment = 0x08000; // 32KiB
-    static const std::size_t s_large_segment = 0x10000; // 64KiB
-
-    static const std::size_t s_tiny_increment_shift = 3;
-    static const std::size_t s_tiny_increment = (1u << s_tiny_increment_shift); // = 8
-
-    static const std::size_t s_num_tiny_heaps = s_tiny_limit / s_tiny_increment;
-    static const std::size_t s_num_small_heaps = log2_c(s_small_limit) - log2_c(s_tiny_limit);
-    static const std::size_t s_num_large_heaps = log2_c(s_large_limit) - log2_c(s_small_limit);
-
-    static std::size_t tiny_bucket_index(std::size_t n) noexcept
+    static std::size_t bucket_index(std::size_t n, std::size_t bucket_shift) noexcept
     {
-        return ((n + s_tiny_increment - 1) >> s_tiny_increment_shift) - 1;
-    }
-
-    static std::size_t bucket_index(std::size_t n) noexcept
-    {
-        return log2_c((n - 1) >> s_bucket_shift) - 1;
-    }
-
-    static constexpr std::size_t round_to_pow_of_2(std::size_t n) noexcept
-    {
-        return 1u << log2_c(n - 1);
+        return detail::log2_c((n - 1) >> bucket_shift) - 1;
     }
 
   private:
+    heap_config m_config;
     Context*    m_context;
     std::size_t m_max_size;
-    bool        m_never_free;
-    std::size_t m_num_reserve_segments;
     heap_vector m_tiny_heaps;
     heap_vector m_heaps;
     heap_map    m_huge_heaps;
     std::mutex  m_mutex;
 
   public:
-    heap(Context* context, bool never_free = false, std::size_t num_reserve_segments = 1)
-    : m_context{context}
-    , m_max_size(std::max(round_to_pow_of_2(s_large_limit * 2), s_large_limit))
-    , m_never_free{never_free}
-    , m_num_reserve_segments{num_reserve_segments}
-    , m_tiny_heaps(s_tiny_limit / s_tiny_increment)
-    , m_heaps(bucket_index(m_max_size) + 1)
+    heap(Context* context, heap_config const& config = get_default_heap_config())
+    : m_config{config}
+    , m_context{context}
+    , m_max_size(
+          std::max(detail::round_to_pow_of_2(m_config.m_large_limit * 2), m_config.m_large_limit))
+    , m_tiny_heaps(m_config.m_tiny_limit / m_config.m_tiny_increment)
+    , m_heaps(bucket_index(m_max_size, m_config.m_bucket_shift) + 1)
     {
         for (std::size_t i = 0; i < m_tiny_heaps.size(); ++i)
             m_tiny_heaps[i] = std::make_unique<fixed_size_heap_type>(m_context,
-                s_tiny_increment * (i + 1), s_tiny_segment, m_never_free, m_num_reserve_segments);
+                m_config.m_tiny_increment * (i + 1), m_config.m_tiny_segment_size,
+                m_config.m_never_free, m_config.m_num_reserve_segments);
 
-        for (std::size_t i = 0; i < s_num_small_heaps; ++i)
+        for (std::size_t i = 0; i < m_config.m_num_small_heaps; ++i)
             m_heaps[i] = std::make_unique<fixed_size_heap_type>(m_context,
-                (s_tiny_limit << (i + 1)), s_small_segment, m_never_free, m_num_reserve_segments);
+                (m_config.m_tiny_limit << (i + 1)), m_config.m_small_segment_size,
+                m_config.m_never_free, m_config.m_num_reserve_segments);
 
-        for (std::size_t i = 0; i < s_num_large_heaps; ++i)
-            m_heaps[i + s_num_small_heaps] = std::make_unique<fixed_size_heap_type>(m_context,
-                (s_small_limit << (i + 1)), s_large_segment, m_never_free, m_num_reserve_segments);
+        for (std::size_t i = 0; i < m_config.m_num_large_heaps; ++i)
+            m_heaps[i + m_config.m_num_small_heaps] = std::make_unique<fixed_size_heap_type>(
+                m_context, (m_config.m_small_limit << (i + 1)), m_config.m_large_segment_size,
+                m_config.m_never_free, m_config.m_num_reserve_segments);
 
-        for (std::size_t i = 0; i < m_heaps.size() - (s_num_small_heaps + s_num_large_heaps); ++i)
-            m_heaps[i + s_num_small_heaps + s_num_large_heaps] =
-                std::make_unique<fixed_size_heap_type>(m_context, (s_large_limit << (i + 1)),
-                    (s_large_limit << (i + 1)), m_never_free, m_num_reserve_segments);
+        for (std::size_t i = 0;
+             i < m_heaps.size() - (m_config.m_num_small_heaps + m_config.m_num_large_heaps); ++i)
+            m_heaps[i + m_config.m_num_small_heaps + m_config.m_num_large_heaps] =
+                std::make_unique<fixed_size_heap_type>(m_context,
+                    (m_config.m_large_limit << (i + 1)), (m_config.m_large_limit << (i + 1)),
+                    m_config.m_never_free, m_config.m_num_reserve_segments);
     }
 
     heap(heap const&) = delete;
@@ -194,20 +176,22 @@ class heap
 
     pointer allocate(std::size_t size, std::size_t numa_node)
     {
-        if (size <= s_tiny_limit)
-            return {m_tiny_heaps[tiny_bucket_index(size)]->allocate(numa_node)};
+        if (size <= m_config.m_tiny_limit)
+            return {m_tiny_heaps[tiny_bucket_index(size, m_config.m_tiny_increment,
+                                     m_config.m_tiny_increment_shift)]
+                        ->allocate(numa_node)};
         else if (size <= m_max_size)
-            return {m_heaps[bucket_index(size)]->allocate(numa_node)};
+            return {m_heaps[bucket_index(size, m_config.m_bucket_shift)]->allocate(numa_node)};
         else
         {
             fixed_size_heap_type* h;
             {
                 std::lock_guard<std::mutex> lock(m_mutex);
-                const auto                  s = round_to_pow_of_2(size);
+                const auto                  s = detail::round_to_pow_of_2(size);
                 auto&                       u_ptr = m_huge_heaps[s];
                 if (!u_ptr)
-                    u_ptr = std::make_unique<fixed_size_heap_type>(m_context, s, s, m_never_free,
-                        m_num_reserve_segments);
+                    u_ptr = std::make_unique<fixed_size_heap_type>(m_context, s, s,
+                        m_config.m_never_free, m_config.m_num_reserve_segments);
                 h = u_ptr.get();
             }
             return {h->allocate(numa_node)};
@@ -223,20 +207,23 @@ class heap
 #if HWMALLOC_ENABLE_DEVICE
     pointer allocate(std::size_t size, std::size_t numa_node, int device_id)
     {
-        if (size <= s_tiny_limit)
-            return {m_tiny_heaps[tiny_bucket_index(size)]->allocate(numa_node, device_id)};
+        if (size <= m_config.m_tiny_limit)
+            return {m_tiny_heaps[tiny_bucket_index(size, m_config.m_tiny_increment,
+                                     m_config.m_tiny_increment_shift)]
+                        ->allocate(numa_node, device_id)};
         else if (size <= m_max_size)
-            return {m_heaps[bucket_index(size)]->allocate(numa_node, device_id)};
+            return {m_heaps[bucket_index(size, m_config.m_bucket_shift)]->allocate(numa_node,
+                device_id)};
         else
         {
             fixed_size_heap_type* h;
             {
                 std::lock_guard<std::mutex> lock(m_mutex);
-                const auto                  s = round_to_pow_of_2(size);
+                const auto                  s = detail::round_to_pow_of_2(size);
                 auto&                       u_ptr = m_huge_heaps[s];
                 if (!u_ptr)
-                    u_ptr = std::make_unique<fixed_size_heap_type>(m_context, s, s, m_never_free,
-                        m_num_reserve_segments);
+                    u_ptr = std::make_unique<fixed_size_heap_type>(m_context, s, s,
+                        m_config.m_never_free, m_config.m_num_reserve_segments);
                 h = u_ptr.get();
             }
             return {h->allocate(numa_node, device_id)};
@@ -293,29 +280,4 @@ class heap
     }
 };
 
-template<typename Context>
-const std::size_t heap<Context>::s_tiny_limit;
-template<typename Context>
-const std::size_t heap<Context>::s_small_limit;
-template<typename Context>
-const std::size_t heap<Context>::s_large_limit;
-template<typename Context>
-const std::size_t heap<Context>::s_bucket_shift;
-template<typename Context>
-const std::size_t heap<Context>::s_tiny_segment;
-template<typename Context>
-const std::size_t heap<Context>::s_small_segment;
-template<typename Context>
-const std::size_t heap<Context>::s_large_segment;
-template<typename Context>
-const std::size_t heap<Context>::s_tiny_increment_shift;
-template<typename Context>
-const std::size_t heap<Context>::s_tiny_increment;
-template<typename Context>
-const std::size_t heap<Context>::s_num_tiny_heaps;
-template<typename Context>
-const std::size_t heap<Context>::s_num_small_heaps;
-template<typename Context>
-const std::size_t heap<Context>::s_num_large_heaps;
-
 } // namespace hwmalloc

include/hwmalloc/heap_config.hpp

@@ -0,0 +1,63 @@
+/*
+ * ghex-org
+ *
+ * Copyright (c) 2014-2025, ETH Zurich
+ * All rights reserved.
+ *
+ * Please, refer to the LICENSE file in the root directory.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+#pragma once
+
+#include <cstddef>
+
+namespace hwmalloc
+{
+namespace detail
+{
+inline constexpr std::size_t
+log2_c(std::size_t n) noexcept
+{
+    return ((n < 2) ? 1 : 1 + log2_c(n >> 1));
+}
+
+inline constexpr std::size_t
+round_to_pow_of_2(std::size_t n) noexcept
+{
+    return 1u << log2_c(n - 1);
+}
+} // namespace detail
+
+struct heap_config
+{
+    static constexpr bool        never_free_default = false;
+    static constexpr std::size_t num_reserve_segments_default = 16u;
+    static constexpr std::size_t tiny_limit_default = 128u;             // 128B
+    static constexpr std::size_t small_limit_default = 4096u;           // 4KiB
+    static constexpr std::size_t large_limit_default = 2097152u;        // 2MiB
+    static constexpr std::size_t tiny_segment_size_default = 65536u;    // 64KiB
+    static constexpr std::size_t small_segment_size_default = 65536u;   // 64KiB
+    static constexpr std::size_t large_segment_size_default = 2097152u; // 2MiB
+
+    bool                         m_never_free;
+    std::size_t                  m_num_reserve_segments;
+    std::size_t                  m_tiny_limit;
+    std::size_t                  m_small_limit;
+    std::size_t                  m_large_limit;
+    std::size_t                  m_bucket_shift = detail::log2_c(m_tiny_limit) - 1;
+    std::size_t                  m_tiny_segment_size;
+    std::size_t                  m_small_segment_size;
+    std::size_t                  m_large_segment_size;
+    static constexpr std::size_t m_tiny_increment_shift = 3;
+    static constexpr std::size_t m_tiny_increment = (1u << m_tiny_increment_shift);
+    std::size_t                  m_num_tiny_heaps = m_tiny_limit / m_tiny_increment;
+    std::size_t m_num_small_heaps = detail::log2_c(m_small_limit) - detail::log2_c(m_tiny_limit);
+    std::size_t m_num_large_heaps = detail::log2_c(m_large_limit) - detail::log2_c(m_small_limit);
+
+    heap_config(bool never_free, std::size_t num_reserve_segments, std::size_t tiny_limit,
+        std::size_t small_limit, std::size_t large_limit, std::size_t tiny_segment_size,
+        std::size_t small_segment_size, std::size_t large_segment_size);
+};
+
+heap_config const& get_default_heap_config();
+} // namespace hwmalloc

src/CMakeLists.txt

@@ -1,3 +1,5 @@
+target_sources(hwmalloc PRIVATE heap_config.cpp)
+
 if (NUMA_LIBRARY)
     target_sources(hwmalloc PRIVATE numa.cpp)
 else()