diff --git a/tcmalloc/cpu_cache.h b/tcmalloc/cpu_cache.h index 6234480c6..826c8b765 100644 --- a/tcmalloc/cpu_cache.h +++ b/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 [[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::Allocate(size_t size_class) { } template +template inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* CpuCache::AllocateFast( size_t size_class) { TC_ASSERT_GT(size_class, 0); - return freelist_.Pop(size_class); + return freelist_.template Pop(size_class); } template diff --git a/tcmalloc/internal/percpu_tcmalloc.h b/tcmalloc/internal/percpu_tcmalloc.h index 056cacaee..8ac395312 100644 --- a/tcmalloc/internal/percpu_tcmalloc.h +++ b/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 and returns its result. + template [[nodiscard]] void* Pop(size_t size_class); // Add up to 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 +template inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab::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::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::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::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::Pop( #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ && defined(__aarch64__) template +template inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab::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::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::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::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::Pop( #if !TCMALLOC_INTERNAL_PERCPU_USE_RSEQ template +template inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* TcmallocSlab::Pop( size_t size_class) { return nullptr; diff --git a/tcmalloc/sizemap.h b/tcmalloc/sizemap.h index 69101c532..155d2c68f 100644 --- a/tcmalloc/sizemap.h +++ b/tcmalloc/sizemap.h @@ -207,8 +207,8 @@ class SizeMap { // TODO(b/171978365): Replace the output parameter with returning // absl::optional. template - [[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(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 + [[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 - [[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; } diff --git a/tcmalloc/tcmalloc.cc b/tcmalloc/tcmalloc.cc index 7b2754439..2c13a60d0 100644 --- a/tcmalloc/tcmalloc.cc +++ b/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(size_class); + } else { + res = tc_globals.cpu_cache().AllocateFast(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(size_class); + } else { + ret = tc_globals.cpu_cache().AllocateFast(size_class); + } if (ABSL_PREDICT_FALSE(ret == nullptr)) { SLOW_PATH_BARRIER(); return slow_alloc_small(size, size_class, policy);