fix(rocm): lock the JIT module cache in get_jit_module #2183

Part of #1801

Problem / Background

The ROCm overlay's process-global JIT module cache has no lock. Two threads that launch JIT kernels at the same time can corrupt it. This was found in the review of PR #2181 (#2149). It has not been observed or reproduced yet.

Current Behavior

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/jit_module.cpp:497-517 (main at b7116d1): get_jit_module_cache() returns a function-static std::unordered_map<std::string, JitModule>, and get_jit_module does an unlocked map.find(key) followed by map.try_emplace(key, device(mlx_device), name, builder, cache). try_emplace runs the JitModule constructor in place, which compiles through hiprtc (custom kernels pass use_disk_cache = false, custom_kernel.cpp:348-355, so every fast::hip_kernel compiles on first use in each process). This leaves a window of tens to hundreds of milliseconds. Two threads that miss the same key both compile and both insert. An insert that grows the table rehashes it while another thread is inside find. get_jit_module_cache() is also exported in jit_module.h:158, but nothing calls it.

JitModule::get_kernel (jit_module.cpp:476-495) also writes the per-kernel configured flag (kernels_ in jit_module.h:114) without a lock. No ROCm caller passes configure_kernel, so this is a benign data race today.

Upstream MLX at the pin (81ba1c6a, mlx/backend/cuda/jit_module.cpp:478) and the CUDA overlay (src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp:563-586) do the cache lookup under a leaked static auto* mtx = new std::shared_mutex. A std::shared_lock covers the find. On a miss, a std::unique_lock covers a second find and then the try_emplace. The CUDA overlay also serializes the lazy per-kernel fill with a std::mutex kernels_mtx_ (jit_module.h:129, jit_module.cpp:539, #1566). The same ROCm file already locks its other process-global, hsaco_cache_dir() (jit_module.cpp:148-152).

Reachability

The race is reachable in production today, though it has not been observed. Every launch of a Compiled primitive (compiled.cpp:875), a fast::hip_kernel port (custom_kernel.cpp:348) and compute_dynamic_offset (slicing.cpp:64) calls get_jit_module, including launches that hit the cache. The server runs model forwards on several OS threads, each with its own thread-local GPU stream and no global eval lock: the BatchScheduler, EmbeddingWorker (src/server/embedding_worker.rs:132), RerankWorker (src/server/rerank_worker.rs:113) and AudioWorker (src/server/audio_worker.rs:165), as listed in src/lib/mlxcel-core/src/streams.rs. DEFAULT_DEVICE_LOCK there guards only default-device swaps. For example, -m <chat> --embedding-model <emb> serves both models at once, so the first embedding request can insert new modules while the chat scheduler is inside find for its decode kernels. New keys also appear late in a process's life, whenever a new template-argument combination (for example a new head dim or dtype) is first used. In tests, the gates (verify-test-rocm, test-fast-rocm) pass --test-threads=1, but a hand run of cargo test --features rocm uses libtest's default thread count. On CUDA, a related JIT race, the per-kernel fill in #1566, crashed the mlxcel-core tests at that thread count. Priority is medium: the race is reachable in a supported server configuration, and its failure is a crash or a corrupted kernel handle, not a wrong number. It is not high because nobody has observed it, and single-model serving is far less exposed.

Proposed Solution

Copy upstream's locking into the ROCm overlay. Do not design new locking.

• jit_module.cpp: move get_jit_module_cache() into the anonymous namespace and make the storage leaked, static auto* cache = new std::unordered_map<std::string, JitModule>, as upstream does. Add static auto* mtx = new std::shared_mutex. Do the device-index key find under std::shared_lock. On a miss, take std::unique_lock, repeat the find, and try_emplace only if the key is still missing. Keep the key format "<device index>:<name>" and the signature unchanged. Add #include <shared_mutex>.
• jit_module.h: delete the get_jit_module_cache() declaration, so nothing can reach the map without the lock. Add std::mutex kernels_mtx_ to JitModule.
• JitModule::get_kernel: keep the lock-free kernels_.find, since the map is filled only in the constructor. Take std::lock_guard<std::mutex> on kernels_mtx_ around the configured-flag check and set, as the CUDA overlay does at jit_module.cpp:539.
• Rejected: a per-key std::once_flag or a future, which would let compiles of different keys run in parallel. It diverges from upstream, and a hiprtc compile happens once per key per process.
• The leaked cache means hipModuleUnload no longer runs at exit. This matches upstream's stated reason: user code may still run JIT code after the main thread tears down. It also avoids the unload-at-exit failures on a faulted device described in LOCAL_FIXES item 7.

Scope

In scope: patches-rocm/mlx/backend/rocm/jit_module.cpp and jit_module.h; a test-only bridge fixture; a new ROCm test file; and a LOCAL_FIXES entry.

Out of scope: the CUDA and Metal backends (already locked); cache keying (#2149 / PR #2181); the disk hsaco cache, whose same-key writes become single-writer within a process once this lands.

Implementation Notes

• Reuse: Model the fixture on rocm_fault_probe_array (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5178-5225): #ifdef MLXCEL_BRIDGE_ROCM_BACKEND, fast::hip_kernel, and a throwing stub on other backends.
• Fixture: std::unique_ptr<MlxArray> rocm_jit_race_probe_array(const MlxArray& input, int32_t variant), declared in mlx_cxx_bridge.h and the lib.rs ffi block. It wraps mlxcel_core::rocm_faults::jit_race_probe_array(input, variant) -> Result<UniquePtr<MlxArray>, cxx::Exception>. The kernel name is "mlxcel_jit_race_probe_v" + std::to_string(variant), with input inp and output out. The body is out[thread_index()] = inp[thread_index()] * 2.0f + 1.0f;. Use one block of n threads, with n = input.size() and n <= 1024. The output is f32 with shape {n}.
• Edge cases: two threads miss the same key, so only one JitModule may be constructed and both must get the same reference; an insert that rehashes while other threads read; and different devices, where the key prefix stays.
• Error handling: unchanged. A compile or load failure throws out of the constructor. Under the unique_lock, the unlock happens when the exception unwinds, and nothing is inserted, so the next caller retries.
• LOCAL_FIXES: add an entry to src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md under "Fixes to the fork's kernels" with the next free number (34 if PR fix(rocm): key custom-kernel JIT modules by their generated source #2181 lands first with 33). It names mlx/backend/rocm/jit_module.cpp and mlx/backend/rocm/jit_module.h, says what was unlocked and how it is now locked, and ends with "Applies to the fork; kept in mlxcelverse under the 2026-10-06 fork policy, not proposed there." Add nothing under docs/mlxcelverse/upstream/.

Acceptance Criteria

• get_jit_module in the ROCm overlay matches the pin's CUDA locking (shared lock on find, then unique lock with a second find before try_emplace), and get_jit_module_cache is gone from jit_module.h.

• JitModule::get_kernel sets the configured flag under kernels_mtx_.

• tests/rocm_jit_module_concurrency.rs (#![cfg(feature = "rocm")]) contains same_kernel_first_launch_from_many_threads. For 16 rounds it starts 8 threads. Each thread installs its own stream with streams::new_thread_local_generation_stream and install_thread_local_default_stream, waits on a std::sync::Barrier, calls the probe with that round's variant on arange(256) f32, evaluates the result, and asserts every element equals 2*i + 1 exactly.

• The same file contains distinct_kernels_first_launch_while_others_hit_cache. One thread repeatedly launches an already-warm variant while 8 threads each first-launch their own new variant behind a barrier, for 16 rounds. All outputs must be exact.

• With the lock in get_jit_module removed, run the new test file 20 times. Report in the PR how many runs crashed or failed. If none did, say so: the test is then a crash guard, not a deterministic reproduction.

• The LOCAL_FIXES entry exists with no upstream marking, and make verify-rocm-overlay passes.

• Added during implementation (PR fix(rocm): lock the JIT module cache and the HIP event pool #2196): HipEventPool in patches-rocm/mlx/backend/rocm/event.hip was also an unlocked process-global and failed the new tests in 20 of 20 runs with the JIT cache locked (hipEventRecord ... invalid resource handle, SIGABRT, SIGSEGV); it is now locked with a std::mutex, recorded as LOCAL_FIXES item 36.

Verification

cargo test --profile test-fast --features rocm --test rocm_jit_module_concurrency
cargo clippy --features rocm --test rocm_jit_module_concurrency -- -D warnings
cargo clippy -p mlxcel-core --features rocm --lib -- -D warnings
make verify-rocm-overlay
make verify-rocm

Run these on a ROCm host (gfx1151). A pass means both tests pass, with no SIGSEGV or SIGABRT, and verify-rocm reports 0 failed.