From cc86abf80a932bfd85082583c10f0768c9925ef8 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 `lean_alloc_small_object_core` from 18 to 14 instructions by keeping mimalloc's generic allocation path out of line, so that the fast path needs no stack frame. The entry point now tests `page->free` itself and tail-jumps to a cold helper on a miss, instead of calling `mi_theap_malloc_small` and checking its result. Only the generic routine can fail, so both the OOM check and the `m_cs_sz` store move to that helper; leaving either at the call site would make the cold call one that has to return, and the fast path would pay a prologue and epilogue on every allocation to serve a path it almost never takes. Restating the callers' precondition on `sz` (word-aligned, non-zero, at most `MI_SMALL_SIZE_MAX`) additionally folds mimalloc's bin-index rounding into a plain byte offset and lets the header store drop its mask. That bound is a safety condition in its own right, as a larger `sz` would index `pages_free_direct` out of bounds, so it is also asserted. ```diff -push rbx ; frame, only for the generic 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 generic 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 ``` 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.781G to 18.650G user instructions (-0.694%) and `import Lean` from 1.9356G to 1.9215G (-0.729%), each reproducible to within 0.02%. LLVM's hot-cold splitting does not do this: with `-mllvm -hot-cold-split=true` the pass fires 27 times elsewhere in the translation unit and leaves this function byte-identical, because it outlines a cold region into a function the original still calls and returns from, which does not remove the frame. Co-Authored-By: Claude Opus 5 --- src/include/lean/lean.h | 2 ++ src/runtime/mimalloc.cpp | 34 +++++++++++++++++++++++++++++++--- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h index 8a0e614c85c4..3dc2341b716a 100644 --- a/src/include/lean/lean.h +++ b/src/include/lean/lean.h @@ -492,6 +492,8 @@ inlined into it. Requires `sz` to be a positive multiple of `LEAN_OBJECT_SIZE_DELTA` of at most `MI_SMALL_SIZE_MAX`. Initializes `m_cs_sz`. + +Panics on exhaustion rather than returning `NULL`, so callers need no OOM check. */ LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz); #endif diff --git a/src/runtime/mimalloc.cpp b/src/runtime/mimalloc.cpp index 9a4d944d7111..1f3e69de7e64 100644 --- a/src/runtime/mimalloc.cpp +++ b/src/runtime/mimalloc.cpp @@ -36,15 +36,43 @@ 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, out of line so that the entry point below can + tail-jump here and stay a leaf function; a call that returned would cost it a prologue and + epilogue on every allocation, to serve a path it almost never takes. Two things have to stay on + this side for that to work: the OOM check, as only this routine can fail, and the `m_cs_sz` + store, without which `sz` would have to survive the call. */ +static mi_decl_noinline lean_object * lean_alloc_small_generic_core(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(); + lean_object * o = (lean_object *)mem; + 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_core` 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;