[BUG] c++ compile cache should be thread safe

[davidkoski]

Describe the bug

The c++ compile cache should be thread safe -- whatever mutable state is behind mlx::core::detail::compile should be thread safe so it can safely be called from multiple threads (common in Swift, maybe C++).

See also:

• CompiledFunction is not thread safe mlx-swift#226
• Crashes occur when running multiple LLMEvaluators. mlx-swift#225

To Reproduce

See:

• Crashes occur when running multiple LLMEvaluators. mlx-swift#225

Note that the fix for this was to protect the inner function that mutates the arguments (as that is where the crash was observed), during discussion with @awni we think that the cache of graphs also needs locking (or whatever technique is appropriate) as it is also mutable state.

Expected behavior

Inner state of mlx::core should be thread safe or documented as not (like MLXArray).

Desktop (please complete the following information):

• OS Version: macOS15
• Version 0.23.1

[awni]

Inner state of mlx::core should be thread safe

There are still a few places this is not true still. Off the top of my head:

• rng state
• compiler cache
• transform state (e.g. InTracing)

[ashthedash2k]

I'd be interested in working on this issue, I have knowledge of threading and caches so I might be able to help out? I have not really ever contributed to open source, the code base is pretty large, I was wondering if I could get some pointers / a starting direction! Thanks so much.