From 70b5d5486fffa0353319d7550b3e6262855906c5 Mon Sep 17 00:00:00 2001 From: Sam Gross Date: Mon, 18 Dec 2023 16:28:38 -0500 Subject: [PATCH] WIP --- Include/internal/mimalloc/mimalloc/internal.h | 2 + Include/internal/pycore_mimalloc.h | 23 ++++++++ Include/internal/pycore_tstate.h | 7 ++- Lib/test/test_import/__init__.py | 4 +- Objects/mimalloc/heap.c | 22 ++++++-- Objects/mimalloc/init.c | 24 ++++---- Objects/obmalloc.c | 36 ++++++++++++ Python/pystate.c | 56 +++++++++++++++++++ 8 files changed, 152 insertions(+), 22 deletions(-) diff --git a/Include/internal/mimalloc/mimalloc/internal.h b/Include/internal/mimalloc/mimalloc/internal.h index f076bc6a40f9773..cb6e211de5bb639 100644 --- a/Include/internal/mimalloc/mimalloc/internal.h +++ b/Include/internal/mimalloc/mimalloc/internal.h @@ -85,6 +85,7 @@ mi_threadid_t _mi_thread_id(void) mi_attr_noexcept; mi_heap_t* _mi_heap_main_get(void); // statically allocated main backing heap void _mi_thread_done(mi_heap_t* heap); void _mi_thread_data_collect(void); +void _mi_tld_init(mi_tld_t* tld, mi_heap_t* bheap); // os.c void _mi_os_init(void); // called from process init @@ -170,6 +171,7 @@ size_t _mi_bin_size(uint8_t bin); // for stats uint8_t _mi_bin(size_t size); // for stats // "heap.c" +void _mi_heap_init_ex(mi_heap_t* heap, mi_tld_t* tld, mi_arena_id_t arena_id); void _mi_heap_destroy_pages(mi_heap_t* heap); void _mi_heap_collect_abandon(mi_heap_t* heap); void _mi_heap_set_default_direct(mi_heap_t* heap); diff --git a/Include/internal/pycore_mimalloc.h b/Include/internal/pycore_mimalloc.h index c29dc82a42762aa..789b48da9d28308 100644 --- a/Include/internal/pycore_mimalloc.h +++ b/Include/internal/pycore_mimalloc.h @@ -9,11 +9,34 @@ # error "pycore_mimalloc.h must be included before mimalloc.h" #endif +#define _Py_MIMALLOC_HEAP_MEM 0 // PyMem_Malloc() and friends +#define _Py_MIMALLOC_HEAP_OBJECT 1 // non-GC objects +#define _Py_MIMALLOC_HEAP_GC 2 // GC objects without pre-header +#define _Py_MIMALLOC_HEAP_GC_PRE 3 // GC objects with pre-header +#define _Py_MIMALLOC_HEAP_COUNT 4 + #include "pycore_pymem.h" #define MI_DEBUG_UNINIT PYMEM_CLEANBYTE #define MI_DEBUG_FREED PYMEM_DEADBYTE #define MI_DEBUG_PADDING PYMEM_FORBIDDENBYTE +#ifdef Py_DEBUG +# define MI_DEBUG 1 +#else +# define MI_DEBUG 0 +#endif #include "mimalloc.h" +#include "mimalloc/types.h" +#include "mimalloc/internal.h" + +struct _mimalloc_thread_state { +#ifdef Py_GIL_DISABLED + mi_heap_t *current_object_heap; + mi_heap_t heaps[_Py_MIMALLOC_HEAP_COUNT]; + mi_tld_t tld; +#else + char _unused; // empty structs are not allowed +#endif +}; #endif // Py_INTERNAL_MIMALLOC_H diff --git a/Include/internal/pycore_tstate.h b/Include/internal/pycore_tstate.h index 17f3e8656417735..856ddd5e7e5ff07 100644 --- a/Include/internal/pycore_tstate.h +++ b/Include/internal/pycore_tstate.h @@ -8,6 +8,8 @@ extern "C" { # error "this header requires Py_BUILD_CORE define" #endif +#include "pycore_mimalloc.h" // struct _mimalloc_thread_state + // Every PyThreadState is actually allocated as a _PyThreadStateImpl. The // PyThreadState fields are exposed as part of the C API, although most fields @@ -16,7 +18,10 @@ typedef struct _PyThreadStateImpl { // semi-public fields are in PyThreadState. PyThreadState base; - // TODO: add private fields here +#ifdef Py_GIL_DISABLED + struct _mimalloc_thread_state mimalloc; +#endif + } _PyThreadStateImpl; diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py index 48c0a43f29e27f8..0145d36f836c586 100644 --- a/Lib/test/test_import/__init__.py +++ b/Lib/test/test_import/__init__.py @@ -26,7 +26,8 @@ from test.support import os_helper from test.support import ( STDLIB_DIR, swap_attr, swap_item, cpython_only, is_emscripten, - is_wasi, run_in_subinterp, run_in_subinterp_with_config, Py_TRACE_REFS) + is_wasi, run_in_subinterp, run_in_subinterp_with_config, Py_TRACE_REFS, + Py_GIL_DISABLED) from test.support.import_helper import ( forget, make_legacy_pyc, unlink, unload, ready_to_import, DirsOnSysPath, CleanImport) @@ -2018,6 +2019,7 @@ def parse(cls, text): return self +# @unittest.skipIf(Py_GIL_DISABLED, "deallocates objects from a different interpreter") @requires_singlephase_init class SinglephaseInitTests(unittest.TestCase): diff --git a/Objects/mimalloc/heap.c b/Objects/mimalloc/heap.c index 4eb622ed4bad765..30dd18923e4cf6c 100644 --- a/Objects/mimalloc/heap.c +++ b/Objects/mimalloc/heap.c @@ -206,18 +206,28 @@ mi_heap_t* mi_heap_get_backing(void) { return bheap; } -mi_decl_nodiscard mi_heap_t* mi_heap_new_in_arena(mi_arena_id_t arena_id) { - mi_heap_t* bheap = mi_heap_get_backing(); - mi_heap_t* heap = mi_heap_malloc_tp(bheap, mi_heap_t); // todo: OS allocate in secure mode? - if (heap == NULL) return NULL; +void _mi_heap_init_ex(mi_heap_t* heap, mi_tld_t* tld, mi_arena_id_t arena_id) +{ _mi_memcpy_aligned(heap, &_mi_heap_empty, sizeof(mi_heap_t)); - heap->tld = bheap->tld; + heap->tld = tld; heap->thread_id = _mi_thread_id(); heap->arena_id = arena_id; - _mi_random_split(&bheap->random, &heap->random); + if (heap == tld->heap_backing) { + _mi_random_init(&heap->random); + } + else { + _mi_random_split(&tld->heap_backing->random, &heap->random); + } heap->cookie = _mi_heap_random_next(heap) | 1; heap->keys[0] = _mi_heap_random_next(heap); heap->keys[1] = _mi_heap_random_next(heap); +} + +mi_decl_nodiscard mi_heap_t* mi_heap_new_in_arena(mi_arena_id_t arena_id) { + mi_heap_t* bheap = mi_heap_get_backing(); + mi_heap_t* heap = mi_heap_malloc_tp(bheap, mi_heap_t); // todo: OS allocate in secure mode? + if (heap == NULL) return NULL; + _mi_heap_init_ex(heap, bheap->tld, arena_id); heap->no_reclaim = true; // don't reclaim abandoned pages or otherwise destroy is unsafe // push on the thread local heaps list heap->next = heap->tld->heaps; diff --git a/Objects/mimalloc/init.c b/Objects/mimalloc/init.c index 7dfa76577371171..376e14b49b7c7bb 100644 --- a/Objects/mimalloc/init.c +++ b/Objects/mimalloc/init.c @@ -297,24 +297,20 @@ static bool _mi_heap_init(void) { mi_thread_data_t* td = mi_thread_data_zalloc(); if (td == NULL) return false; - mi_tld_t* tld = &td->tld; - mi_heap_t* heap = &td->heap; + _mi_tld_init(&td->tld, &td->heap); + _mi_heap_init_ex(&td->heap, &td->tld, _mi_arena_id_none()); + _mi_heap_set_default_direct(&td->heap); + } + return false; +} + +void _mi_tld_init(mi_tld_t* tld, mi_heap_t* bheap) { _mi_memcpy_aligned(tld, &tld_empty, sizeof(*tld)); - _mi_memcpy_aligned(heap, &_mi_heap_empty, sizeof(*heap)); - heap->thread_id = _mi_thread_id(); - _mi_random_init(&heap->random); - heap->cookie = _mi_heap_random_next(heap) | 1; - heap->keys[0] = _mi_heap_random_next(heap); - heap->keys[1] = _mi_heap_random_next(heap); - heap->tld = tld; - tld->heap_backing = heap; - tld->heaps = heap; tld->segments.stats = &tld->stats; tld->segments.os = &tld->os; tld->os.stats = &tld->stats; - _mi_heap_set_default_direct(heap); - } - return false; + tld->heap_backing = bheap; + tld->heaps = bheap; } // Free the thread local default heap (called from `mi_thread_done`) diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c index 99c95d90658b08e..883adcb1c19b6e0 100644 --- a/Objects/obmalloc.c +++ b/Objects/obmalloc.c @@ -88,19 +88,37 @@ _PyMem_RawFree(void *Py_UNUSED(ctx), void *ptr) void * _PyMem_MiMalloc(void *ctx, size_t size) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM]; + return mi_heap_malloc(heap, size); +#else return mi_malloc(size); +#endif } void * _PyMem_MiCalloc(void *ctx, size_t nelem, size_t elsize) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM]; + return mi_heap_calloc(heap, nelem, elsize); +#else return mi_calloc(nelem, elsize); +#endif } void * _PyMem_MiRealloc(void *ctx, void *ptr, size_t size) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM]; + return mi_heap_realloc(heap, ptr, size); +#else return mi_realloc(ptr, size); +#endif } void @@ -112,20 +130,38 @@ _PyMem_MiFree(void *ctx, void *ptr) void * _PyObject_MiMalloc(void *ctx, size_t nbytes) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = tstate->mimalloc.current_object_heap; + return mi_heap_malloc(heap, nbytes); +#else return mi_malloc(nbytes); +#endif } void * _PyObject_MiCalloc(void *ctx, size_t nelem, size_t elsize) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = tstate->mimalloc.current_object_heap; + return mi_heap_calloc(heap, nelem, elsize); +#else return mi_calloc(nelem, elsize); +#endif } void * _PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes) { +#ifdef Py_GIL_DISABLED + _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET(); + mi_heap_t *heap = tstate->mimalloc.current_object_heap; + return mi_heap_realloc(heap, ptr, nbytes); +#else return mi_realloc(ptr, nbytes); +#endif } void diff --git a/Python/pystate.c b/Python/pystate.c index e18eb0186d00103..53a85260c4060c3 100644 --- a/Python/pystate.c +++ b/Python/pystate.c @@ -236,6 +236,9 @@ tstate_is_bound(PyThreadState *tstate) static void bind_gilstate_tstate(PyThreadState *); static void unbind_gilstate_tstate(PyThreadState *); +static void tstate_mimalloc_bind(PyThreadState *); +static void tstate_mimalloc_clear(PyThreadState *); + static void bind_tstate(PyThreadState *tstate) { @@ -256,6 +259,9 @@ bind_tstate(PyThreadState *tstate) tstate->native_thread_id = PyThread_get_thread_native_id(); #endif + // mimalloc state needs to be initialized from the active thread. + tstate_mimalloc_bind(tstate); + tstate->_status.bound = 1; } @@ -1533,6 +1539,8 @@ PyThreadState_Clear(PyThreadState *tstate) tstate->on_delete(tstate->on_delete_data); } + tstate_mimalloc_clear(tstate); + tstate->_status.cleared = 1; // XXX Call _PyThreadStateSwap(runtime, NULL) here if "current". @@ -2495,3 +2503,51 @@ _PyThreadState_MustExit(PyThreadState *tstate) } return 1; } + +/********************/ +/* mimalloc support */ +/********************/ + +static void +tstate_mimalloc_bind(PyThreadState *tstate) +{ +#ifdef Py_GIL_DISABLED + struct _mimalloc_thread_state *mts = &((_PyThreadStateImpl*)tstate)->mimalloc; + + // Initialize the mimalloc thread state. This must be called from the + // same thread that will use the thread state. The "mem" heap doubles as + // the "backing" heap. + mi_tld_t *tld = &mts->tld; + _mi_tld_init(tld, &mts->heaps[_Py_MIMALLOC_HEAP_MEM]); + + // Initialize each heap + for (Py_ssize_t i = 0; i < _Py_MIMALLOC_HEAP_COUNT; i++) { + _mi_heap_init_ex(&mts->heaps[i], tld, _mi_arena_id_none()); + } + + // By default, object allocations use _Py_MIMALLOC_HEAP_OBJECT. + // _PyObject_GC_New() and similar functions temporarily override this to + // use one of the GC heaps. + mts->current_object_heap = &mts->heaps[_Py_MIMALLOC_HEAP_OBJECT]; +#endif +} + +static void +tstate_mimalloc_clear(PyThreadState *tstate) +{ +#ifdef Py_GIL_DISABLED + if (!tstate->_status.bound) { + // The mimalloc heaps are only initialized when the thread is bound. + return; + } + + _PyThreadStateImpl *tstate_impl = (_PyThreadStateImpl *)tstate; + for (Py_ssize_t i = 0; i < _Py_MIMALLOC_HEAP_COUNT; i++) { + // Abandon all segments in use by this thread. This pushes them to + // a shared pool to later be reclaimed by other threads. It's important + // to do this before the thread state is destroyed so that objects + // remain visible to the GC. + _mi_heap_collect_abandon(&tstate_impl->mimalloc.heaps[i]); + } +#endif +}