[davidtai]

Summary

compile(std::function fun, bool shapeless) keys a closure (a function with no addressable target) by the address of a heap copy pfun. The deleter of pfun erased the trace only from the compile cache of the thread that called compile(). Compile caches are thread-local. A trace that another thread recorded stayed in that thread's cache after pfun was freed:

• The trace keeps its arrays alive (inputs, outputs, tape, captured constants) until that thread exits.
• The key is a freed heap address. A later closure can get the same address. When that closure is called on the same thread with the same input shapes, dtypes and default stream, the cache returns the old trace and the old graph runs.

Fix: the deleter erases through the fork's cache registry (detail::compile_erase(fun_id), PR #10), which erases the key from every live thread cache. The erase runs before delete, so no new allocation can get the address while an old trace is still in a cache.

Note: the cross-thread registry (CompileCacheRegistry) is a fork addition (PR #10), not upstream code. Upstream has only the per-cache compile_erase(cache, fun_id).

Changes

• mlx/compile.cpp: the pfun deleter calls detail::compile_erase(fun_id) instead of detail::compile_erase(cache, fun_id) with the creating thread's cache.
• tests/compile_tests.cpp: test freeing a compiled closure erases every thread cache. It compiles a closure that captures an array, calls it once on a second thread (the trace goes into that thread's cache only), frees the closure while the second thread is still alive, and checks captured.is_uniquely_referenced(). A trace that stays holds the captured array, so the check fails. The array is built from host data and the worker does not evaluate the output, so no device work holds a reference. The test uses only existing API (array::is_uniquely_referenced, fork PR for buffer uniqueness). No test-only production API.
• fork.yaml: section "Freed compiled closures leave every thread cache".

Not changed (open finding): the Python binding PyCompiledFun::~PyCompiledFun erases only the cache of the thread that called the function last. Same class of hole; out of scope for this PR.

Test first

• Test only, 30437ff3: fails on this test only, compile_tests.cpp:954 CHECK( captured.is_uniquely_referenced() ), in every C++ job: Build and Test https://github.com/Layr-Labs/mlx/actions/runs/40108428325 (Fedora x86_64 + aarch64, ASAN, UBSAN: 280/281 pass) and Fork Tests https://github.com/Layr-Labs/mlx/actions/runs/40108428358 (build-cpp, GPU and CPU: 315/316 pass).
• Test + fix, ff4bb35b: Build and Test https://github.com/Layr-Labs/mlx/actions/runs/40108806485, Fork Tests https://github.com/Layr-Labs/mlx/actions/runs/40108806495, Fork Diff Check https://github.com/Layr-Labs/mlx/actions/runs/40108806500.

Merge order

This PR is step 1 of a pin chain. Merge in this order, each with Create a merge commit (never squash; a squash makes the pinned SHAs unreachable from main, as with mlx #38 / mlx-c #17):

1. Layr-Labs/mlx (this PR)
2. Layr-Labs/mlx-c re-pin PR (GIT_TAG = this PR's head)
3. Layr-Labs/mlx-swift re-pin PR (gitlinks = this PR's head and the mlx-c PR head)
4. Layr-Labs/mlx-swift-lm re-pin PR (mlx-swift revision = the mlx-swift PR head)

Draft upstream issue (not opened)

compile(std::function): freeing a closure leaves its traces in other threads' compile caches

Since compile caches became thread-local, mlx::core::compile(std::function<...> fun, bool shapeless) handles a function with no addressable target like this: it copies fun into a heap pfun, uses pfun.get() as fun_id, and gives pfun a deleter that calls detail::compile_erase(cache, fun_id) with cache = detail::compile_cache() taken when compile() runs.

The compiled function can be called on any thread. Each call records the trace in the calling thread's cache. When the last copy of the compiled function is freed, the deleter erases fun_id only from the cache of the thread that called compile(). Traces on other threads stay until those threads exit.

Effects:

1. The traces keep their arrays (inputs, outputs, tape, captured constants) alive. In a long-lived thread pool this is a leak that grows with each closure compiled on one thread and called on another.
2. The key is a freed heap address. A new closure can get the same address. If it is called on a thread that holds the old trace, with the same input shapes, dtypes and default stream, CompileCache::find returns the old entry and the old graph runs in place of the new function.

Repro (C++, doctest): compile a closure that captures an array c; call it once on a second thread; keep that thread alive; free the compiled function on the first thread. c still has an extra reference (the second thread's trace). With the fix, c.array_desc_.use_count() == 1.

Possible fix: keep a process-wide registry of live thread caches (each CompileCache adds itself in its constructor and removes itself in its destructor), and make the pfun deleter erase fun_id from every registered cache before delete. Layr-Labs/mlx carries this registry (also used for the process-wide compile_erase(fun_id) / compile_clear_cache() that older bindings call) and this fix.

The Python binding has a related path: PyCompiledFun::~PyCompiledFun calls compile_erase(cache, fun_id) with the cache of the thread that called the function last (call_impl sets cache). fun_id is the address of the Python function object, which Python reuses freely.

🤖 Generated with Claude Code