From 1ae7897f6ec46c389f25089b1a5e43fb950fa088 Mon Sep 17 00:00:00 2001
From: Misaki Kasumi <misakikasumi@outlook.com>
Date: Tue, 22 Mar 2022 02:34:05 +0800
Subject: [PATCH] LocalAllocator: add alloc_hinted

In current implementation, `alloc` expects small allocations are more
likely, which is likely true for most cases, and puts small allocations
in the fast path. However, in some scenario, the allocator users may know
an approximate alloc size in compile time. For example, they may be
allocating a certein type of objects that may have a typical size, but
the actual size in runtime may vary; or they may know the least size of
some objects.

By introducing a method `add_hinted`, we can let users to hint the
allocator. It accepts a compile-time size hint, based on which a hinted
fast path is chosen, instead of always favoring small allocations.

Currently, we use the size hint to adjust the fast path only in
LocalAllocator. Further we can use hint to do more optimizations, like,
adding `add_range_hinted` for backends, pass the hint from
LocalAllocator to them, and adjust their fast path to passing the
allocation to their parent or allocating by their own.
---
 src/mem/localalloc.h | 55 ++++++++++++++++++++++++++++++++++----------
 1 file changed, 43 insertions(+), 12 deletions(-)

diff --git a/src/mem/localalloc.h b/src/mem/localalloc.h
index d9bf18493..f267f4410 100644
--- a/src/mem/localalloc.h
+++ b/src/mem/localalloc.h
@@ -167,7 +167,7 @@ namespace snmalloc
      * passed to the core allocator.
      */
     template<ZeroMem zero_mem>
-    SNMALLOC_SLOW_PATH capptr::Alloc<void> alloc_not_small(size_t size)
+    SNMALLOC_FAST_PATH capptr::Alloc<void> alloc_not_small_fast(size_t size)
     {
       if (size == 0)
       {
@@ -206,6 +206,12 @@ namespace snmalloc
       });
     }
 
+    template<ZeroMem zero_mem>
+    SNMALLOC_SLOW_PATH capptr::Alloc<void> alloc_not_small(size_t size)
+    {
+      return alloc_not_small_fast<zero_mem>(size);
+    }
+
     template<ZeroMem zero_mem>
     SNMALLOC_FAST_PATH capptr::Alloc<void> small_alloc(size_t size)
     {
@@ -241,6 +247,12 @@ namespace snmalloc
         domesticate, size, slowpath);
     }
 
+    template<ZeroMem zero_mem>
+    SNMALLOC_SLOW_PATH capptr::Alloc<void> small_alloc_slow(size_t size)
+    {
+      return small_alloc<zero_mem>(size);
+    }
+
     /**
      * Send all remote deallocation to other threads.
      */
@@ -416,8 +428,8 @@ namespace snmalloc
     /**
      * Allocate memory of a dynamically known size.
      */
-    template<ZeroMem zero_mem = NoZero>
-    SNMALLOC_FAST_PATH ALLOCATOR void* alloc(size_t size)
+    template<size_t size_hint, ZeroMem zero_mem = NoZero>
+    SNMALLOC_FAST_PATH ALLOCATOR void* alloc_hinted(size_t size)
     {
 #ifdef SNMALLOC_PASS_THROUGH
       // snmalloc guarantees a lot of alignment, so we can depend on this
@@ -429,27 +441,46 @@ namespace snmalloc
         memset(result, 0, size);
       return result;
 #else
-      // Perform the - 1 on size, so that zero wraps around and ends up on
-      // slow path.
-      if (SNMALLOC_LIKELY(
-            (size - 1) <= (sizeclass_to_size(NUM_SMALL_SIZECLASSES - 1) - 1)))
+      if constexpr (
+        (size_hint - 1) <= (sizeclass_to_size(NUM_SMALL_SIZECLASSES - 1) - 1))
       {
-        // Small allocations are more likely. Improve
-        // branch prediction by placing this case first.
-        return capptr_reveal(small_alloc<zero_mem>(size));
+        // Perform the - 1 on size, so that zero wraps around and ends up on
+        // slow path.
+        if (SNMALLOC_LIKELY(
+              (size - 1) <= (sizeclass_to_size(NUM_SMALL_SIZECLASSES - 1) - 1)))
+        {
+          return capptr_reveal(small_alloc<zero_mem>(size));
+        }
+
+        return capptr_reveal(alloc_not_small<zero_mem>(size));
       }
+      else
+      {
+        if (SNMALLOC_UNLIKELY(
+              (size - 1) <= (sizeclass_to_size(NUM_SMALL_SIZECLASSES - 1) - 1)))
+        {
+          return capptr_reveal(small_alloc_slow<zero_mem>(size));
+        }
 
-      return capptr_reveal(alloc_not_small<zero_mem>(size));
+        return capptr_reveal(alloc_not_small_fast<zero_mem>(size));
+      }
 #endif
     }
 
+    template<ZeroMem zero_mem = NoZero>
+    SNMALLOC_FAST_PATH ALLOCATOR void* alloc(size_t size)
+    {
+      // Small allocations are more likely
+      return alloc_hinted<1, zero_mem>(size);
+    }
+
     /**
      * Allocate memory of a statically known size.
      */
     template<size_t size, ZeroMem zero_mem = NoZero>
     SNMALLOC_FAST_PATH ALLOCATOR void* alloc()
     {
-      return alloc<zero_mem>(size);
+      return alloc_hinted<size, zero_mem>(size);
     }
 
     /*