Make mx.compile cache erasing thread safe

[zcbenz]

Close #3940.

Compiled python function would erase the cache automatically on destruction, but the destruction can happen on any thread so it could happen that a cache entry gets deleted on a different thread from creation.

This PR enforces the erasing to happen on the original cache rather than the cache in the current thread, and makes CompilerCache thread safe to the race condition above. The test is from #4096.

The code is not strictly thread safe to avoid unnecessary overheads based on following assumptions:

1. A thread won't add cache to another thread.
2. A python function can not be destructed while being compiled or used.

[yentur]

A question rather than a correction, since you may have scoped this out already.

find() hands back a CacheEntry& into *ptr, but the caller keeps only that reference, not a copy of the shared_ptr. So cache_.erase(fun_id) drops the last reference and frees the vector while the caller is still filling the entry in. Two heap-use-after-frees under ASAN on 2c69f20, 10 of 10 runs each, CPU-only build.

An erase reaching the tracing thread while it is inside compile_trace:

freed by thread T0:
  #2 detail::compile_erase(weak_ptr<CompilerCache> const&, unsigned long) compile.cpp:1221
previously allocated by thread T2:
  #2 CompilerCache::find(...) compile.cpp:388

The same thing re-entrant on one thread, when a traced body drops another wrapper over the same callable:

freed by thread T0:
  #2 detail::compile_erase(...) compile.cpp:1221
  #5 detail::compile_trace(...) compile.cpp:434
previously allocated by thread T0:
  #2 CompilerCache::find(...)   compile.cpp:388

Repro for the second, no threads needed:

constexpr std::uintptr_t outer_id = 0xf11d;
std::function<std::vector<array>(const std::vector<array>&)> fun =
    [&](const std::vector<array>& inputs) {
      detail::compile_erase(detail::compiler_cache(), outer_id);
      return std::vector<array>{inputs[0] + array(1.0f)};
    };
auto compiled = detail::compile(fun, outer_id);
eval(compiled({array(3.0f)}));

The re-entrant one predates this PR and reproduces on main. The cross-thread one only becomes reachable once erases reach the tracing thread's cache. Holding the shared_ptr in the caller for the duration of the call, rather than only dereferencing it inside find(), would cover both. Is that in scope here, or better as a follow-up?

[zcbenz]

Technically it won't happen in python bindings because it is guaranteed that a function won't be destructed until nothing is using it, i.e. erasing won't happen while compiling the same function. But it might make sense not giving users a foot gun when it is not hard to do.

[sashko-zakharchuk]

The CI failures are a teardown crash introduced by this branch. Three lines reproduce it, no
threads:

import mlx.core as mx

f = mx.compile(lambda x: x + 1)
mx.eval(f(mx.array([1.0])))

Linux, CPU-only build: exit 139 on 2c69f2073, exit 0 on main 306bdcd18, 3 runs each. The
full suite does the same, 835 tests OK and then exit 139, as in the job log.

The symbolized stack for thread 1 is further down that same log, addresses trimmed:

#0 new_threadstate.llvm ()
#1 PyGILState_Ensure.cold ()
#2 nanobind::gil_scoped_acquire::gil_scoped_acquire (nb_misc.h:15)
#3 ThreadCleanup::~ThreadCleanup (python/src/transforms.cpp:415)
#4 __call_tls_dtors ()

~ThreadCleanup used to acquire the GIL only when the cache was not empty, and the atexit
hook empties it during finalization, so on the main thread that branch never ran. Without the
empty check the acquire is unconditional, and libc runs the main thread's TLS destructors
after Py_Finalize, so PyGILState_Ensure reaches new_threadstate with a null interpreter
state, which is the 0x60 in the CI output. The clear itself is already a no-op there, so
only the acquire matters. Cache contents are not involved: with mx.disable_compile() the
cache is never populated and it still crashes 3/3, and so does del f; gc.collect(). Main
survives on the empty check rather than on timing: keep an entry alive past the atexit hook
there and it exits 139 the same way.

An early return fixes it here, repro 0/3 and 835 tests OK:

~ThreadCleanup() {
  if (!Py_IsInitialized()) {
    return;
  }
  nb::gil_scoped_acquire gil;
  mx::detail::compile_clear_cache(cache);
}

That covers the after-finalization case only, it does not change what a thread exiting during
finalization does. Happy to test whatever you prefer.