From f193bdaa3c4cf7394afdcf899e38c9ea3c8edb50 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Tue, 1 Sep 2026 08:16:56 +0000 Subject: [PATCH] perf: make the small-object allocation fast path a leaf This PR shortens the fused small-object allocation entry point from 18 to 14 instructions and gives constructor allocation a single merged header store, worth 0.75% of user instructions on the workloads measured below. Three changes, none of which need a patch to mimalloc. The page refill moves into a cold helper that the entry points tail-jump to, so the fast path no longer builds a stack frame for a call it rarely makes; the helper also writes `m_cs_sz`, which is what leaves `sz` dead across the branch and keeps `lean_alloc_small_object_core` a leaf as well. Restating the callers' precondition on `sz` (word-aligned, non-zero, at most `MI_SMALL_SIZE_MAX`) folds mimalloc's bin-index rounding into a plain byte offset and drops the mask before the header store; the bound is a safety condition in its own right, since a larger `sz` would index `pages_free_direct` out of bounds. Finally, a second entry point `lean_alloc_small_object_raw` leaves the whole header to the caller. The fast path of `lean_alloc_small_object_core`: ```diff -push rbx ; frame, only for the refill call below mov rax, [rip+...] ; &lean_g_tls inc QWORD PTR fs:[rax] ; heartbeat++ mov rcx, QWORD PTR fs:[rax+0x8] ; theap mov esi, edi -lea rax, [rsi+0x7] ; round sz up to a bin index -and rax, 0xfffffffffffffff8 -mov rdx, [rcx+rax*1+0x120] ; theap->pages_free_direct[..] +mov rdx, [rcx+rsi*1+0x120] ; sz is the byte offset already mov rax, [rdx+0x8] ; page->free test rax, rax -je .Lcold ; inline cold path, ending in pop rbx; ret +je lean_alloc_small_generic_core ; tail jump, nothing to tear down mov rcx, [rax] ; pop the block mov [rdx+0x8], rcx inc WORD PTR [rdx+0x10] ; page->used++ mov QWORD PTR [rax], 0x0 ; mimalloc's free-list scrub -mov WORD PTR [rax+0x4], di ; m_cs_sz, masked to 16 bits -pop rbx +mov DWORD PTR [rax+0x4], edi ; m_cs_sz, no mask needed ret ``` `lean_alloc_small_object_raw` is what lets `lean_alloc_ctor` write all four header fields with adjacent stores that merge, rather than having the callee store `m_cs_sz` and the caller fill in around it. Allocating a `List.cons`: ```diff mov edi, 0x18 -call lean_alloc_small_object_core ; stores m_cs_sz itself -mov DWORD PTR [rax], 0x1 ; m_rc -mov WORD PTR [rax+0x6], 0x102 ; m_other, m_tag; cannot merge across m_cs_sz +call lean_alloc_small_object_raw ; leaves the header alone +movabs rcx, 0x102001800000001 ; m_rc | m_cs_sz | m_other | m_tag +mov QWORD PTR [rax], rcx ; one store mov QWORD PTR [rax+0x8], r14 mov QWORD PTR [rax+0x10], rbx ``` Measured by swapping `libleanshared.so` between two full builds, single-threaded with `-DElab.async=false` on pinned cores: `tests/elab_bench/big_omega.lean` goes from 18.946G to 18.803G user instructions (-0.754%) and `import Lean` from 1.9365G to 1.9219G (-0.753%), each reproducible to within 0.01%; cycles on the former fall about 1.2%. Heartbeat counts are unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- src/include/lean/lean.h | 35 +++++++++++++++++++++++++- src/runtime/mimalloc.cpp | 53 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 84 insertions(+), 4 deletions(-) diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h index 8a0e614c85c4..917c39d4f130 100644 --- a/src/include/lean/lean.h +++ b/src/include/lean/lean.h @@ -494,6 +494,11 @@ Requires `sz` to be a positive multiple of `LEAN_OBJECT_SIZE_DELTA` of at most `MI_SMALL_SIZE_MAX`. Initializes `m_cs_sz`. */ LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz); +/* +As above, but leaving the entire header to the caller, so that a caller writing all four header +fields can do so with a single store. +*/ +LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_raw(unsigned sz); #endif #ifndef __cplusplus @@ -796,9 +801,37 @@ static inline uint8_t * lean_ctor_scalar_cptr(lean_object * o) { static inline lean_object * lean_alloc_ctor(unsigned tag, unsigned num_objs, unsigned scalar_sz) { assert(tag <= LeanMaxCtorTag && num_objs < LEAN_MAX_CTOR_FIELDS && scalar_sz < LEAN_MAX_CTOR_SCALARS_SIZE); - lean_object * o = lean_alloc_ctor_memory(lean_usize_add_checked(lean_usize_add_checked(sizeof(lean_ctor_object), lean_usize_mul_checked(sizeof(void*), num_objs)), scalar_sz)); + size_t sz = lean_usize_add_checked(lean_usize_add_checked(sizeof(lean_ctor_object), lean_usize_mul_checked(sizeof(void*), num_objs)), scalar_sz); +#ifdef LEAN_MIMALLOC + // NOTE: `sz` is known at compile time for most callers, folding the branches below + size_t sz1 = lean_align(sz, LEAN_OBJECT_SIZE_DELTA); + lean_object * o; + if (LEAN_LIKELY(sz1 <= MI_SMALL_SIZE_MAX)) { + o = lean_alloc_small_object_raw((unsigned)sz1); + } else { + lean_inc_heartbeat(); + void * mem = mi_malloc(sz1); + if (mem == 0) lean_internal_panic_out_of_memory(); + o = (lean_object*)mem; + } + if (sz1 > sz) { + /* Zero the last word so that the (sz1 - sz) uninitialized trailing bytes do not make the + structural comparisons in `maxsharing.cpp` and `compact.cpp` miss sharing. */ + ((size_t*)((char*)o + sz1))[-1] = 0; + } + /* Write the full header with adjacent stores; for the constant arguments of compiled code they + merge into one. `lean_set_st_header` cannot merge, as it must preserve the `m_cs_sz` that + `lean_alloc_small_object_core` has already written. */ + lean_internal_set_rc(o, 1); + o->m_cs_sz = (unsigned)sz1; + o->m_other = num_objs; + o->m_tag = tag; + return o; +#else + lean_object * o = lean_alloc_ctor_memory((unsigned)sz); lean_set_st_header(o, tag, num_objs); return o; +#endif } static inline b_lean_obj_res lean_ctor_get(b_lean_obj_arg o, unsigned i) { diff --git a/src/runtime/mimalloc.cpp b/src/runtime/mimalloc.cpp index 9a4d944d7111..7047f8649f0c 100644 --- a/src/runtime/mimalloc.cpp +++ b/src/runtime/mimalloc.cpp @@ -36,17 +36,64 @@ extern "C" void lean_mi_theap_cache_init(void) { lean_g_tls.mi_theap_default = mi_theap_get_default(); } +/* The callers guarantee `sz > 0 && sz % LEAN_OBJECT_SIZE_DELTA == 0 && sz <= MI_SMALL_SIZE_MAX`. + The bound is also a safety condition, as a larger `sz` would index `pages_free_direct` out of + bounds. Restating both lets the compiler fold the bin index into a plain byte offset and emit + the `m_cs_sz` store below without a mask. */ +static inline mi_page_t * lean_small_page(mi_theap_t * theap, unsigned sz) { + lean_assert(sz > 0 && (sz % MI_INTPTR_SIZE) == 0 && sz <= MI_SMALL_SIZE_MAX); +#if defined(__GNUC__) || defined(__clang__) + if ((sz % MI_INTPTR_SIZE) != 0 || sz == 0 || sz > MI_SMALL_SIZE_MAX) __builtin_unreachable(); +#endif + return _mi_theap_get_free_small_page(theap, sz); +} + +/* Wraps mimalloc's generic allocation routine. This and `lean_alloc_small_generic_core` are out + of line so that the matching entry points below can tail-jump to them and stay leaf functions; + a call that returned would cost them a prologue and epilogue on every allocation, to serve a + path they almost never take. Only this routine can fail, so the OOM check belongs here: + `mi_theap_malloc_small` would force it to the call site, where it would prevent the tail + jump. */ +static mi_decl_noinline lean_object * lean_alloc_small_generic_raw(mi_theap_t * theap, size_t sz) { + void * mem = _mi_malloc_generic(theap, sz, 0, NULL); + if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory(); + return (lean_object *)mem; +} + +/* Adds the `m_cs_sz` store, as `lean_alloc_small_object_core` does to + `lean_alloc_small_object_raw`. Repeating it here leaves that entry point nothing to do after the + call; otherwise `sz` would have to survive it, the call could not be a tail jump, and it would + need a stack frame. */ +static mi_decl_noinline lean_object * lean_alloc_small_generic_core(mi_theap_t * theap, size_t sz) { + lean_object * o = lean_alloc_small_generic_raw(theap, sz); + o->m_cs_sz = sz; + return o; +} + extern "C" LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz) { lean_runtime_tls * tls = &lean_g_tls; tls->heartbeat++; lean_assert(sz > 0 && sz % LEAN_OBJECT_SIZE_DELTA == 0 && sz <= MI_SMALL_SIZE_MAX); /* Feeding the cached theap into mimalloc saves the load of mimalloc's own thread-local: the heartbeat update and the theap read share one TLS address computation. */ - void * mem = mi_theap_malloc_small(tls->mi_theap_default, sz); - if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory(); - lean_object * o = (lean_object *)mem; + mi_theap_t * const theap = tls->mi_theap_default; + mi_page_t * const page = lean_small_page(theap, sz); + if (LEAN_UNLIKELY(page->free == NULL)) return lean_alloc_small_generic_core(theap, sz); + /* No OOM check: only the generic path can fail, and the test above diverts to it; see + `lean_alloc_small_generic_raw` for why the check belongs there. */ + lean_object * o = (lean_object *)mi_page_malloc_zero(theap, page, sz, false, NULL); /* `m_cs_sz` must be the exact (aligned) requested size, not mimalloc's potentially larger block size: `lean_small_object_size` and `leangz` rely on it. */ o->m_cs_sz = sz; return o; } + +extern "C" LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_raw(unsigned sz) { + lean_runtime_tls * tls = &lean_g_tls; + tls->heartbeat++; + mi_theap_t * const theap = tls->mi_theap_default; + mi_page_t * const page = lean_small_page(theap, sz); + if (LEAN_UNLIKELY(page->free == NULL)) return lean_alloc_small_generic_raw(theap, sz); + /* Cannot be `NULL`; see `lean_alloc_small_object_core`. */ + return (lean_object *)mi_page_malloc_zero(theap, page, sz, false, NULL); +}