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); +}