- Evaluate policy.is_cold() early in fastpath and pass down via template boolean IsCold.

tcmalloc/cpu_cache.h

@@ -331,6 +331,7 @@ class CpuCache {
   // Separate allocation fast/slow paths.
   // The fast path succeeds iff the thread has already cached the slab pointer
   // (done by AllocateSlow) and there is an available object in the slab.
+  template <bool IsCold = false>
   [[nodiscard]] void* absl_nullable AllocateFast(size_t size_class);
   [[nodiscard]] void* absl_nullable AllocateSlow(size_t size_class);
   // A slightly faster version of AllocateSlow that may be called only
@@ -773,10 +774,11 @@ void* CpuCache<Forwarder>::Allocate(size_t size_class) {
 }
 
 template <class Forwarder>
+template <bool IsCold>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* CpuCache<Forwarder>::AllocateFast(
     size_t size_class) {
   TC_ASSERT_GT(size_class, 0);
-  return freelist_.Pop(size_class);
+  return freelist_.template Pop<IsCold>(size_class);
 }
 
 template <class Forwarder>

tcmalloc/internal/percpu_tcmalloc.h

@@ -201,6 +201,7 @@ class TcmallocSlab {
 
   // Remove an item (LIFO) from the current CPU's slab. If the slab is empty,
   // invokes <underflow_handler> and returns its result.
+  template <bool IsCold = false>
   [[nodiscard]] void* Pop(size_t size_class);
 
   // Add up to <len> items to the current cpu slab from the array located at
@@ -801,10 +802,11 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void PrefetchNextObject(
 
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ && defined(__x86_64__)
 template <size_t NumClasses>
+template <bool IsCold>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
     size_t size_class) {
   TC_ASSERT_NE(size_class, 0);
-  void* next;
+  void* next = nullptr;
   void* result;
   uintptr_t scratch, current;
 
@@ -838,7 +840,9 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
   // Important! code below this must not affect any flags (i.e.: ccc)
   // If so, the above code needs to explicitly set a ccc return value.
 #endif
+      ".if %c[is_cold] == 0\n"
       "movq -16(%[scratch], %[current], 8), %[next]\n"
+      ".endif\n"
       "lea -1(%[current]), %[current]\n"
       "movw %w[current], (%[scratch], %[size_class], 4)\n"
       // Commit
@@ -854,7 +858,7 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
 #else
         [begin_mark_bit] "n"(absl::countr_zero(kBeginMark)),
 #endif
-        [size_class] "r"(size_class)
+        [size_class] "r"(size_class), [is_cold] "n"(IsCold ? 1 : 0)
       : "cc", "memory"
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ_ASM_GOTO_OUTPUT
       : underflow_path
@@ -865,14 +869,18 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
     goto underflow_path;
   }
 #endif
-  TC_ASSERT(next);
+  if constexpr (!IsCold) {
+    TC_ASSERT(next);
+  }
   TC_ASSERT(result);
   TSANAcquire(result);
 
-  // The next pop will be from current-1, but because we prefetch the previous
-  // element we've already just read that, so prefetch current-2.
-  PrefetchSlabMemory(scratch + (current - 2) * sizeof(void*));
-  PrefetchNextObject(next);
+  if constexpr (!IsCold) {
+    // The next pop will be from current-1, but because we prefetch the previous
+    // element we've already just read that, so prefetch current-2.
+    PrefetchSlabMemory(scratch + (current - 2) * sizeof(void*));
+    PrefetchNextObject(next);
+  }
   return AssumeNotNull(result);
 underflow_path:
   return nullptr;
@@ -881,12 +889,13 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
 
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ && defined(__aarch64__)
 template <size_t NumClasses>
+template <bool IsCold>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
     size_t size_class) {
   TC_ASSERT_NE(size_class, 0);
   void* result;
+  void* prefetch = nullptr;
   void* region_start;
-  void* prefetch;
   uintptr_t scratch;
   uintptr_t previous;
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ_ASM_GOTO_OUTPUT
@@ -914,9 +923,13 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
       // previous = scratch - 1
       "sub %w[previous], %w[scratch], #1\n"
       "add %[scratch],%[region_start], %[scratch], LSL #3\n"
+      ".if %c[is_cold] == 0\n"
       // Note we are using a pre-indexed ldp. After this instruction, scratch
       // points to the next location in slab memory, which we will prefetch.
       "ldp %[prefetch], %[result], [%[scratch], #-16]!\n"
+      ".else\n"
+      "ldr %[result], [%[scratch], #-8]\n"
+      ".endif\n"
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ_ASM_GOTO_OUTPUT
       "tbnz %[result], #%c[begin_mark_bit], %l[underflow_path]\n"
 #else
@@ -946,7 +959,7 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
         [cached_slabs_mask] "r"(TCMALLOC_CACHED_SLABS_MASK),
         [begin_mark_mask] "n"(kBeginMark), [scratch2] "=&r"(scratch2),
 #endif
-        [size_class_lsl2] "r"(size_class << 2)
+        [size_class_lsl2] "r"(size_class << 2), [is_cold] "n"(IsCold ? 1 : 0)
       : TCMALLOC_RSEQ_CLOBBER, "memory"
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ_ASM_GOTO_OUTPUT
         ,
@@ -960,8 +973,10 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
   }
 #endif
   TSANAcquire(result);
-  PrefetchSlabMemory(scratch);
-  PrefetchNextObject(prefetch);
+  if constexpr (!IsCold) {
+    PrefetchSlabMemory(scratch);
+    PrefetchNextObject(prefetch);
+  }
   return AssumeNotNull(result);
 underflow_path:
   return nullptr;
@@ -970,6 +985,7 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
 
 #if !TCMALLOC_INTERNAL_PERCPU_USE_RSEQ
 template <size_t NumClasses>
+template <bool IsCold>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab<NumClasses>::Pop(
     size_t size_class) {
   return nullptr;

tcmalloc/sizemap.h

@@ -207,8 +207,8 @@ class SizeMap {
   // TODO(b/171978365): Replace the output parameter with returning
   // absl::optional<uint32_t>.
   template <typename Policy>
-  [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE SizeMapResult
-  GetSizeClass(Policy policy, size_t size) const {
+  [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE inline SizeMapResult GetSizeClass(
+      Policy policy, size_t size, bool is_cold) const {
     const size_t align = static_cast<size_t>(policy.align());
     TC_ASSERT(absl::has_single_bit(align));
 
@@ -224,7 +224,7 @@ class SizeMap {
     // Note, if security heap partitioning is enabled, only data (partition 0)
     // is added to the cold heap. See the comment for kClassArraySizePartitions
     // for more details.
-    if (kHasExpandedClasses && policy.is_cold()) {
+    if (kHasExpandedClasses && is_cold) {
       TC_ASSERT(policy.allocation_type() == AllocationType::New);
       TC_ASSERT_LT(idx + (policy.security_partition() + kColdRegisterStride) *
                              kClassArraySize,
@@ -271,11 +271,17 @@ class SizeMap {
     return {true, size_class};
   }
 
+  template <typename Policy>
+  [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE inline SizeMapResult GetSizeClass(
+      Policy policy, size_t size) const {
+    return GetSizeClass(policy, size, policy.is_cold());
+  }
+
   // Returns size class for given size, or 0 if this instance has not been
   // initialized yet. REQUIRES: size <= kMaxSize.
   template <typename Policy>
-  [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE size_t
-  SizeClass(Policy policy, size_t size) const {
+  [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE inline size_t SizeClass(
+      Policy policy, size_t size) const {
     ASSUME(size <= kMaxSize);
     return GetSizeClass(policy, size).size_class;
   }

tcmalloc/tcmalloc.cc

@@ -1189,14 +1189,24 @@ alloc_small_sampled_hooks_or_perthread(size_t size, size_t size_class,
   if (ABSL_PREDICT_FALSE(size_class == 0)) {
     // This happens on the first call then the size class table is not inited.
     TC_ASSERT(tc_globals.IsInited());
-    auto ret = tc_globals.sizemap().GetSizeClass(policy, size);
+    const bool is_cold = policy.is_cold();
+    auto ret = tc_globals.sizemap().GetSizeClass(policy, size, is_cold);
     size_class = ret.size_class;
     TC_CHECK(ret.is_small);
   }
   void* res;
   // If we are here because of sampling, try AllocateFast first.
-  if (ABSL_PREDICT_TRUE(weight == 0) ||
-      (res = tc_globals.cpu_cache().AllocateFast(size_class)) == nullptr) {
+  if (ABSL_PREDICT_TRUE(weight == 0)) {
+    const bool is_cold = policy.is_cold();
+    if (ABSL_PREDICT_FALSE(is_cold)) {
+      res = tc_globals.cpu_cache().AllocateFast<true>(size_class);
+    } else {
+      res = tc_globals.cpu_cache().AllocateFast<false>(size_class);
+    }
+  } else {
+    res = nullptr;
+  }
+  if (ABSL_PREDICT_FALSE(res == nullptr)) {
     if (UsePerCpuCache(tc_globals)) {
       res = tc_globals.cpu_cache().AllocateSlow(size_class);
     } else {
@@ -1271,8 +1281,9 @@ static inline Pointer ABSL_ATTRIBUTE_ALWAYS_INLINE fast_alloc(size_t size,
   // path. If malloc is not yet initialized, we may end up with size_class == 0
   // (regardless of size), but in this case should also delegate to the slow
   // path by the fast path check further down.
+  const bool is_cold = policy.is_cold();
   const auto [is_small, size_class] =
-      tc_globals.sizemap().GetSizeClass(policy, size);
+      tc_globals.sizemap().GetSizeClass(policy, size, is_cold);
   if (ABSL_PREDICT_FALSE(!is_small)) {
     SLOW_PATH_BARRIER();
     TCMALLOC_MUSTTAIL return slow_alloc_large(size, policy);
@@ -1294,7 +1305,12 @@ static inline Pointer ABSL_ATTRIBUTE_ALWAYS_INLINE fast_alloc(size_t size,
   // - cpu / thread cache data has been initialized.
   // - the allocation is not subject to sampling / gwp-asan.
   // - no new/delete hook is installed and required to be called.
-  void* ret = tc_globals.cpu_cache().AllocateFast(size_class);
+  void* ret;
+  if (ABSL_PREDICT_FALSE(is_cold)) {
+    ret = tc_globals.cpu_cache().AllocateFast<true>(size_class);
+  } else {
+    ret = tc_globals.cpu_cache().AllocateFast<false>(size_class);
+  }
   if (ABSL_PREDICT_FALSE(ret == nullptr)) {
     SLOW_PATH_BARRIER();
     return slow_alloc_small(size, size_class, policy);