From ea47ab9bf6ed4b7d36b4e5c1d705e9f38f8ca4e7 Mon Sep 17 00:00:00 2001
From: Nader Al Awar <naderalawar@gmail.com>
Date: Fri, 17 Jul 2026 11:39:09 -0500
Subject: [PATCH] Serialize hostjit builds and improve test

---
 c/parallel.v2/src/hostjit/compiler.cpp        | 12 +++++++++
 .../compute/test_numba_path_thread_safety.py  | 26 +++++++++++++++----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/c/parallel.v2/src/hostjit/compiler.cpp b/c/parallel.v2/src/hostjit/compiler.cpp
index 769ce784d4c..bc4b6b33645 100644
--- a/c/parallel.v2/src/hostjit/compiler.cpp
+++ b/c/parallel.v2/src/hostjit/compiler.cpp
@@ -1695,20 +1695,32 @@ CUDACompiler::~CUDACompiler()
   delete impl_;
 }
 
+// clang and LLVM are embedded in-process and are not thread-safe for concurrent
+// compilation (shared global registries, cl::opt state, allocator use).
+// cuda.compute releases the GIL around native builds, so distinct ops built from
+// multiple threads -- on a GIL or free-threaded interpreter -- otherwise run
+// clang concurrently and corrupt that shared state. Serialize every native
+// compile/link through one process-wide mutex. This guards only the (cached,
+// one-time) build path; kernel launches are unaffected.
+static std::mutex g_native_compile_mutex;
+
 BitcodeResult CUDACompiler::compileToDeviceBitcode(const std::string& source_code, const CompilerConfig& config)
 {
+  const std::lock_guard<std::mutex> lock(g_native_compile_mutex);
   return impl_->compileToDeviceBitcode(source_code, config);
 }
 
 CompilationResult CUDACompiler::compileToObject(
   const std::string& source_code, const std::string& output_path, const CompilerConfig& config)
 {
+  const std::lock_guard<std::mutex> lock(g_native_compile_mutex);
   return impl_->compileToObject(source_code, output_path, config);
 }
 
 LinkResult CUDACompiler::linkToSharedLibrary(
   const std::vector<std::string>& object_files, const std::string& output_path, const CompilerConfig& config)
 {
+  const std::lock_guard<std::mutex> lock(g_native_compile_mutex);
   return impl_->linkToSharedLibrary(object_files, output_path, config);
 }
 } // namespace hostjit
diff --git a/python/cuda_cccl/tests/compute/test_numba_path_thread_safety.py b/python/cuda_cccl/tests/compute/test_numba_path_thread_safety.py
index 23f99bd0a2a..ef9f3400fb4 100644
--- a/python/cuda_cccl/tests/compute/test_numba_path_thread_safety.py
+++ b/python/cuda_cccl/tests/compute/test_numba_path_thread_safety.py
@@ -48,9 +48,23 @@
     reason="Concurrency tests intentionally run concurrent workers."
 )
 
-THREADS = 4
+# Every input in this file is tiny (32-128 elements), so raising the thread count
+# adds no meaningful GPU-memory pressure -- it just exercises the concurrent
+# build/cache machinery harder.
+THREADS = 8
 ITERATIONS = 3
 
+# The distinct-op build storm is the canonical stress for concurrent NATIVE
+# builds. On the v2 (HostJIT) backend clang is embedded in-process and is not
+# thread-safe, so a regression in the compile-serialization lock corrupts the
+# heap and crashes here -- but only at high concurrency (at 4 threads it fires
+# only under external load). Use many threads so this stays a reliable gate.
+# One barrier-synced round of that many simultaneous builds already corrupts the
+# heap pre-fix (it crashed on the first round every time), and extra rounds are
+# expensive on v2 (builds serialize), so a single iteration is enough here.
+BUILD_STORM_THREADS = 24
+BUILD_STORM_ITERATIONS = 1
+
 
 def _run_threaded(workers):
     # The default timeout turns a worker that dies before reaching the barrier
@@ -110,9 +124,9 @@ def test_concurrent_distinct_python_ops_build_storm():
     any cross-thread op/build contamination.
     """
     num_items = 64
-    for iteration in range(ITERATIONS):
+    for iteration in range(BUILD_STORM_ITERATIONS):
         cuda.compute.clear_all_caches()
-        returned_reducers = [None] * THREADS
+        returned_reducers = [None] * BUILD_STORM_THREADS
 
         def make_thread(worker_id):
             k = 10_000 + worker_id * 7 + iteration
@@ -142,10 +156,12 @@ def thread(barrier):
 
             return thread
 
-        _run_threaded([make_thread(worker_id) for worker_id in range(THREADS)])
+        _run_threaded(
+            [make_thread(worker_id) for worker_id in range(BUILD_STORM_THREADS)]
+        )
 
         build_ids = {id(_single_build_result(r)) for r in returned_reducers}
-        assert len(build_ids) == THREADS
+        assert len(build_ids) == BUILD_STORM_THREADS
 
 
 def test_concurrent_shared_python_op_coalesces():