Fix concurrent Metal kernel cache lookup

[kitty-eu-org]

Summary

Device::get_kernel performs its cache-hit lookup while holding a shared lock. The lookup used operator[] on the outer library_kernels_ map, which inserts a new entry on a miss and therefore mutates an unordered_map while other readers may be doing the same. Concurrent first-time custom Metal kernel compilation can consequently crash.

This change uses find for both cache levels in the shared-lock fast path. Cache misses continue through the existing exclusive-lock path, which populates the maps.

Reproduction

The issue can be reproduced with multiple Python threads compiling distinct mx.fast.metal_kernel sources at the same time:

import threading
import mlx.core as mx

n_threads = 16
barrier = threading.Barrier(n_threads)
results = [None] * n_threads
errors = []

def worker(index):
    try:
        stream = mx.new_stream(mx.gpu)
        value = mx.zeros((1,), dtype=mx.float32, stream=stream)
        kernel = mx.fast.metal_kernel(
            name="concurrent_python_kernel",
            input_names=["inp"],
            output_names=["out"],
            source=(
                "uint i = thread_position_in_grid.x; "
                f"out[i] = inp[i] + {index}.0f;"
            ),
        )
        barrier.wait()
        result = kernel(
            inputs=[value],
            grid=(1, 1, 1),
            threadgroup=(1, 1, 1),
            output_shapes=[value.shape],
            output_dtypes=[value.dtype],
            stream=stream,
        )[0]
        mx.eval(result)
        results[index] = result.item()
        mx.clear_streams()
    except BaseException as error:
        errors.append((index, repr(error)))

threads = [threading.Thread(target=worker, args=(i,)) for i in range(n_threads)]
for thread in threads:
    thread.start()
for thread in threads:
    thread.join()
if errors:
    raise RuntimeError(errors)
assert results == list(map(float, range(n_threads)))

Tests

• cmake -S . -B build -DMLX_BUILD_METAL=ON -DMLX_BUILD_TESTS=ON -DMLX_BUILD_PYTHON_BINDINGS=OFF
• cmake --build build --target tests -j2
• ctest --test-dir build --output-on-failure (263/263 passed)
• New regression test repeated 1000 times
• Python reproduction repeated for 100 rounds with 16 threads per round

[zcbenz]

Nice catch!