gh-113141: Upgrade mimalloc to upstream v3 - #159718

The following commit authors need to sign the Contributor License Agreement:

• daanl@outlook.com

Documentation build overview

📚 cpython-previews | 🛠️ Build #35312361 | 📁 Comparing cff014b against main (415ebff)

  🔍 Preview build  

@daanx can you please sign the CLA?

@daanx can you please sign the CLA?

Ah yes, but not really for the current PR as that includes the full mimalloc v3 vendored sources and that is a separate project under the MIT license.

I can of course sign off for the changes I made on the CPython side. The current PR contains both, maybe we can split it? I think commit cedd47b is the CPython part only, while the other commits are just importing mimalloc v3 as a vendored library.

Using MI_WIN_INIT_USE_FLS is not really encouraged for mimalloc. If we can use the default that would be best.

There's a crash on the Windows builds w/o it currently (and it seems to show up basically at startup but the freeze programs make it through it). I was planning on digging into it more later but this is what Claude came up as a fix - I just don't have immediate access to a Windows machine to see where the actual crash is coming from. But that's why it's there for now.

Right now I'm just trying to get things to green CI and then I'll clean up the final result.

@corona10 Hey here it is!

For TSAN fail see: 2599016

IIRC, it was already happened at

About guardrail for memory pressure

any changes expected to increase memory above 15% (with a hard limit of 20%) should be discussed with the SC in advance.

Perf run is here waiting for results: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39346123598

@daanx can you please sign the CLA?

Ah yes, but not really for the current PR as that includes the full mimalloc v3 vendored sources and that is a separate project under the MIT license.

I can of course sign off for the changes I made on the CPython side. The current PR contains both, maybe we can split it? I think commit cedd47b is the CPython part only, while the other commits are just importing mimalloc v3 as a vendored library.

FWIW the docs due already call out mimalloc license: https://docs.python.org/3/license.html#mimalloc although we should probably update the date range w/ the v3 update.

perf run is here waiting for results: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39346123598

2x to 1.14x but still higher than mimalloc v2

New benchmark run: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39373144391

New benchmark run: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39373144391

That was still a 6% regression on memory (but things are headed in the right direction, THP tweak about it from 2x to 1.15, and then the _Py_qsbr_poll fix got us here), here's another run https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39476892845 which has commit 3c8be4f that sets mi_option_purge_delay and mi_option_arena_purge_mult to their v2 values.

Down to a 4% regression now: https://github.com/facebookexperimental/free-threading-benchmarking/tree/main/results/bm-20261020-3.16.0a1+dev-3c8be4f-NOGIL

Another run with new options for mi_option_arena_purge_min_size, and mi_option_page_full_retain and a tweaked mi_option_arena_purge_mult: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39497633255

Down to 3%

This is now neutral: https://github.com/facebookexperimental/free-threading-benchmarking/tree/main/results/bm-20261020-3.16.0a1+dev-a36efb6-NOGIL

It includes one more upstream fix which is a36efb6 that re-enables the mi_arenas_try_purge call which tries to purge expired decommits. This was commented out but not removed in the 3x branch.

@daanx Thanks for merging all of the other upstream PRs I sent. I can send another one to re-enable this but I can also throw it under another ifdef if you're concerned about having it on in the common case. But I think this basically brings us to parity where we could merge this. I'd still want to bring back the latest commits you have upstream but my previous run had that as neutral. I also am assuming your dev3 branch isn't yet released?

Another benchmark run w/ the latest and greatest mimalloc/dev3 branch merged in: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/39651451425

We're now down to 4 vendored in changes which I've opened upstream PRs for (the last being the most questionable I think):
microsoft/mimalloc#1445
microsoft/mimalloc#1446
microsoft/mimalloc#1447
microsoft/mimalloc#1448

This should also resolve:
#118072
#134070
#113141
#159111
#134070
#133571

One more run after merging against and updating the define for 3.6 #160360

https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/41690003577

It looks like the slice shift tweak hits assertions, doing another perf run w/o it and with the base having it: https://github.com/facebookexperimental/free-threading-benchmarking/actions/runs/41693317724

While looking into #161420, I tried this PR. In #161420, worker threads allocate 16 MiB bytes objects that the main thread frees, a few at a time. That needs about 150 MiB, but mimalloc 2.5.2 holds back the blocks other threads freed, up to about 10 GB.

The v3 code in this PR would fix that, but one of the option defaults the PR sets for CPython turns the fix off.

Why v3 fixes it. With its default options, v3 abandons a page that becomes full: the page no longer belongs to the allocating thread (_mi_page_to_full() calls _mi_page_abandon()). When another thread frees a block on an abandoned page, that free collects the page right away (mi_free_try_collect_mt()), so the memory can be reused at once.

Why it comes back here. This PR sets page_full_retain to -1 in set_mimalloc_defaults(), and that turns abandoning off (allow_page_abandon = (page_full_retain >= 0)). Full pages then stay with the allocating thread, and frees from other threads are only collected during that thread's housekeeping, every 1,000 generic allocations.

Peak memory, with about 150 MiB needed:

Could abandoning stay enabled, at least for pages of large blocks? The C program and all numbers are in this comment on microsoft/mimalloc#1117.

Hi @maikschulze , @DinoV : I agree that using page_full_retain = -1 is not a good setting in general; however, I just pushed fixes to the mimalloc dev3 branch that collect the large allocations more aggressively (and a small extra fix for the mi_thread_id_of to clear it on thread exit). That should fix the issue above; let me know.

However, we should try again also with the default page_full_retain (2). I think there is a mistake in the current PR:

op->ob_tid = mi_page_is_abandoned(page) ? 0 : mi_page_thread_id(page);
if (op->ob_tid == 0) merge_refcount(op, 0);

this should become:

op->ob_tid = mi_thread_id_of(op);   // this gets a CPython thread id (not a mimalloc abandoned page tid) (todo: make this more efficient by avoiding overwriting the tid in the object header (and avoiding this call)?
if (op->ob_tid == 0) merge_refcount(op, 0);

edit: actually, it would be even better to call the internal mi_page_qsbr_thread_id(page) as we have the page anyways (instead of mi_thread_id_of(op)) -- but not sure if that is visible there?

I think with this change we don't need the page_full_retain and can use abandoned pages (which should improve matters). Fingers crossed :-)

(not sure if you can easily pull from the latest dev3 ? but nice to try if we are getting to the point of being a drop-in update. I will make a versioned release later this week).

Thank you, @daanx! I tried the C version of the example (8 threads, 16 MiB blocks freed by the main thread, about 144 MiB in flight) on the latest dev3: with page_full_retain = -1 ("-1" below), the peak now stays at what the blocks in flight need, and a run is as fast as with the default options again. Before the fix, it spent most of its CPU time in the kernel.

Good see that @maikschulze; @DinoV , I tested the latest mimalloc dev3 with the fix I proposed:

--- a/Python/gc_free_threading.c
+++ b/Python/gc_free_threading.c
@@ -311,8 +311,8 @@ gc_restore_tid(PyObject *op)
         op->ob_tid = 0;
     }
     else {
-        // Reclaimed pages may have a new owner; abandoned pages have none.
-        op->ob_tid = mi_page_is_abandoned(page) ? 0 : mi_page_thread_id(page);
+        // Full-page abandonment preserves the QSBR thread ownership.
+        op->ob_tid = mi_page_qsbr_thread_id(page);
         if (op->ob_tid == 0) {
             merge_refcount(op, 0);
         }

and it all seems to work now with default mimalloc settings, i.e. page_full_retain=2 and page_reclaim_on_free=0.

import argparse
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
import ctypes
import faulthandler
import gc
import threading
import sys
import weakref

parser = argparse.ArgumentParser()
parser.add_argument("--reclaim", action="store_true")
args = parser.parse_args()

if sys.implementation.name != "cpython" or sys._is_gil_enabled():
    raise SystemExit("Requires free-threaded CPython with the GIL disabled")

TIMEOUT = 30
COUNT = 20_000

faulthandler.dump_traceback_later(TIMEOUT, exit=True)
gc.disable()


def object_thread_id(obj):
    # ob_tid is the first word of a free-threaded CPython object.
    return ctypes.c_size_t.from_address(id(obj)).value


class Node:
    __slots__ = ("value", "link", "__weakref__")

    def __init__(self, value):
        self.value = value
        self.link = None


state = {}
ready = threading.Event()
reclaimed = threading.Event()
finish = threading.Event()


def allocate():
    objects = [Node(i) for i in range(COUNT)]
    state["owner"] = object_thread_id(objects[0])
    state["freed"] = {id(objects[i]) for i in range(0, COUNT, 2)}
    state["garbage"] = [
        weakref.ref(objects[i]) for i in range(0, COUNT, 2)
    ]

    # Half become unreachable cycles, to be freed by GC on another thread.
    for i in range(0, COUNT, 2):
        objects[i].link = objects[i]
        objects[i] = None

    state["objects"] = objects
    ready.set()
    assert finish.wait(TIMEOUT)
    return sum(obj.value for obj in objects if obj is not None)


def reclaim_pages():
    objects = [Node(i) for i in range(COUNT)]
    state["reclaimer"] = object_thread_id(objects[0])
    state["reused"] = sum(
        id(obj) in state["freed"] for obj in objects
    )
    state["new"] = objects
    reclaimed.set()
    assert finish.wait(TIMEOUT)
    return sum(
        obj.value for obj in state["objects"] if obj is not None
    )


expected_sum = sum(range(1, COUNT, 2))

with ThreadPoolExecutor(max_workers=2) as pool:
    owner = pool.submit(allocate)
    other = None
    try:
        assert ready.wait(TIMEOUT)
        objects = state["objects"]
        refs = [
            weakref.ref(obj) for obj in objects if obj is not None
        ]
        assert {
            object_thread_id(obj)
            for obj in objects if obj is not None
        } == {state["owner"]}

        # Free cycles outside their allocating thread, avoiding owner reclaim.
        gc.collect()
        assert all(ref() is None for ref in state["garbage"])

        expected_owners = {state["owner"]}
        if args.reclaim:
            other = pool.submit(reclaim_pages)
            assert reclaimed.wait(TIMEOUT)
            assert state["reused"] > 0, "No freed slots were reused"
            expected_owners.add(state["reclaimer"])

        for _ in range(2):
            gc.collect()
            owners = Counter(
                object_thread_id(obj)
                for obj in objects if obj is not None
            )
            assert owners.keys() <= expected_owners, owners
            if args.reclaim:
                assert owners[state["reclaimer"]] > 0, owners
            assert sum(
                obj.value for obj in objects if obj is not None
            ) == expected_sum
    finally:
        finish.set()

    assert owner.result() == expected_sum
    if other is not None:
        assert other.result() == expected_sum

del objects
state.clear()
gc.collect()
assert all(ref() is None for ref in refs)

faulthandler.cancel_dump_traceback_later()
print("PASS:", "reclaimed pages" if args.reclaim else "abandoned pages")

for retain in 0 2; do
    MIMALLOC_PAGE_FULL_RETAIN=$retain MIMALLOC_PAGE_RECLAIM_ON_FREE=0 \
        ./python -X gil=0 qsbr_gc_test.py

    MIMALLOC_PAGE_FULL_RETAIN=$retain MIMALLOC_PAGE_RECLAIM_ON_FREE=0 \
        ./python -X gil=0 qsbr_gc_test.py --reclaim
done

Thanks, I reran the #161420 pattern against this PR, in both directions, in case it helps with the page_full_retain choice. In free-threaded builds of main (e8c0584877), large blocks that one thread makes and another frees are much cheaper with this PR's mimalloc v3 (48883a281d merged in, with dev3 6a7e041 vendored instead of v3.6.1) than with mimalloc 2.5.2. On memory, with -1 both directions stay close to main's level; only when worker threads exit does about one block per thread wait until the next garbage collection. With 2, the blocks seem to wait in the thread that frees them, and when that is a pool worker, a large part of what was processed stays in memory. I'm not sure what's best here: I only looked at large blocks passed between threads, which is a small part of the picture.

On speed: 2.5.2 makes an madvise call on every free of a block above 4 MiB from another thread; v3 doesn't. Workers make a bytearray(n) and the main thread frees it, 8 in flight, with the PR's -1 (fastest of 5 rounds of 1 GiB, median of 3 processes, PGO + LTO builds, macOS arm64), GiB/s:

On memory, the pattern: 16 MiB bytes objects, 8 in flight, 10 GiB in total (10 rounds of 64). Either the workers make them and the main thread frees them, as in #161420, or the main thread reads them from a file and the workers checksum and free them; with one ThreadPoolExecutor kept, or a new one per round, so that the workers exit. "With collections" means gc.collect() after every round; on its own, the pattern triggers 4 or 5 collections. Release builds without PGO, macOS arm64 and Linux arm64 (Docker on the same machine); page_full_retain set through MIMALLOC_PAGE_FULL_RETAIN, the other options as the PR sets them (page_reclaim_on_free is 0). Peak RSS in MiB, medians of 5:

• With -1 and one pool kept, both directions stay within 60 MiB of main. When the workers exit, about one block per exited worker waits (1.4, 2.6 and 5.1 GiB after 10, 20 and 40 rounds on macOS, against 0.7 to 0.8 GiB on main; 2 is lower in this case) until a collection, where _PyMem_ProcessDelayedNoDealloc() collects the abandoned pages. As far as I can tell from counters in an instrumented build, the worker's last page is empty but still delayed when the thread exits; _mi_page_abandon() unlinks it from the thread's QSBR list, and _mi_arenas_page_abandon() leaves a singleton on no list that an allocation or a free would look at (71 to 77 pages, 1.2 to 1.3 GiB, for 80 exited workers).
• With 2, as far as I can tell, each full singleton page is abandoned right away, and the thread that frees the block reclaims the empty, delayed page into its own QSBR list. When the main thread frees, up to 0.4 GiB waits in its list. When the workers free, about 525 of the 640 pages go to their lists, which they collect only 9 to 18 times in the run, so up to 3.8 GiB waits there, and it grows with the amount processed (5.2, 9.7 and 13.7 GiB for 10, 20 and 40 GiB on macOS, medians of 3). A collection on the main thread doesn't release it, since it only collects its own thread's lists and the abandoned pages. With -1, the main thread keeps these pages and collects its list at every large allocation.
• How much waits with 2 depends on timing (macOS, single runs): if the main thread copies the blocks in memory instead of reading them from a file, the second row stays at 0.1 GiB; if the workers sleep for 3 ms instead of checksumming, it reaches 1.2 GiB with the copy and 7.9 GiB with the file.
• Your gc_restore_tid change gives about the same numbers for 2 (in the table's row order: 731, 4,659, 1,036, 945, 11,205 and 785 MiB).

# Large blocks made on one thread and freed on another, 8 in flight: 10 rounds of 64 blocks of 16 MiB.
# usage: python pool_repro.py produce|consume new|keep [gc]
#   produce: workers make the blocks, the main thread frees them (as in gh-161420)
#   consume: the main thread reads the blocks from a file, workers checksum and free them
#   new: a new ThreadPoolExecutor per round (the workers exit); keep: one pool for all rounds
#   gc: call gc.collect() after every round
import gc, resource, sys, zlib
from collections import deque
from concurrent.futures import ThreadPoolExecutor

direction, mode, collect = sys.argv[1], sys.argv[2], sys.argv[3:] == ["gc"]
BLOCK, COUNT, THREADS, ROUNDS = 16 << 20, 64, 8, 10
pattern = bytes(range(256)) * (BLOCK // 256)
source = open("/dev/zero", "rb", buffering=0)

def make():
    return memoryview(pattern).tobytes()  # a new bytes object, as decompress() returns

kept = ThreadPoolExecutor(THREADS) if mode == "keep" else None
for _ in range(ROUNDS):
    pool = kept or ThreadPoolExecutor(THREADS)
    pending = deque()
    for _ in range(COUNT):
        if direction == "produce":
            pending.append(pool.submit(make))  # the block is freed on the main thread
        else:
            pending.append(pool.submit(zlib.crc32, source.read(BLOCK)))  # freed on the worker
        if len(pending) >= THREADS:
            pending.popleft().result()
    while pending:
        pending.popleft().result()
    if kept is None:
        pool.shutdown()  # the workers exit
    if collect:
        gc.collect()

peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
peak = peak >> 20 if sys.platform == "darwin" else peak >> 10
print(f"{direction} {mode}{' gc' if collect else ''}: peak RSS {peak} MiB")

for retain in -1 2; do
  for args in "produce keep" "consume keep" "produce new" "produce new gc" "consume keep gc"; do
    MIMALLOC_PAGE_FULL_RETAIN=$retain ./python pool_repro.py $args
  done
done

   if (!mi_heap_qsbr_is_enabled(heap)) return false;
+  return false;  // PROBE ONLY: never delay empty pages for QSBR (unsafe in general)
   const bool delayed = mi_page_qsbr_should_delay(&page->qsbr, &heap->qsbr, page->capacity * mi_page_block_size(page));

Separately: in v3, a block above 31 MiB seems to be purged on every free. As far as I can see, _mi_os_good_alloc_size() rounds such a request up to at least 32 MiB, which reaches the default arena_purge_immediate_size (microsoft/mimalloc#1422, raised to 32 MiB in bbe184ae). When each thread allocates and frees its own block, v3 runs at 54 GiB/s at 31 MiB against 229 with 2.5.2 on macOS (PGO + LTO, medians of 5), and at 105 against 207 on Linux (release builds, medians of 3); with MIMALLOC_ARENA_PURGE_IMMEDIATE_SIZE=1GiB it is back at 188 and 180, and at 30 MiB all of them run at about 195 to 230. A bytearray that grows to 30 MiB reaches it too, as its last step asks for 33.75 MiB.

I used an AI assistant to run these measurements and draft this comment; I reviewed the results myself.

Thanks @maikschulze -- I pushed an update to the latest dev3 branch; I think the main issue is/was that foreign threads should not reclaim abandoned empty qsbr pages (as that can lead to them being not collected for a while). Hope that fixes the differences seen in your experiment. (of course, note that the benchmark is synthetic and tests a worst-case scenario with pure allocators and free'ers).