[BUG] TVM-FFI exported functions hang under multithreads #3706

# JAX repro: a CuTe DSL kernel exported with --enable-tvm-ffi and registered as an XLA FFI
# target via jax_tvm_ffi.register_ffi_target. In a single-process multi-GPU jit, XLA's
# per-device threads make the first call of the kernel concurrently, which leaks the init
# lock in libcute_dsl_runtime.so and the step never finishes. Needs >= 3 local GPUs.
# Usage: python3 tvmffi_jax_repro.py            (all local GPUs -> hangs; kill it)
#        CUDA_VISIBLE_DEVICES=0 python3 ...     (control -> OK)
import os

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")

import jax
import jax.numpy as jnp
import jax_tvm_ffi
from jax.sharding import Mesh, PartitionSpec as P

import cutlass.cute as cute
from cutlass import Float32

N = 256  # elements per device; one CTA

@cute.kernel
def add_one_kernel(x: cute.Tensor, y: cute.Tensor):
    i, _, _ = cute.arch.thread_idx()
    y[i] = x[i] + 1.0

@cute.jit
def add_one(x: cute.Tensor, y: cute.Tensor, stream):
    add_one_kernel(x, y).launch(grid=(1, 1, 1), block=(N, 1, 1), stream=stream)

fake = lambda: cute.runtime.make_fake_compact_tensor(
    Float32, (N,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=16
)
compiled = cute.compile(add_one, fake(), fake(), cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
                        options="--enable-tvm-ffi")
jax_tvm_ffi.register_ffi_target("cute_add_one", compiled, ["args", "rets"], platform="gpu")

def per_device(x):
    return jax.ffi.ffi_call("cute_add_one", jax.ShapeDtypeStruct(x.shape, x.dtype))(x)

mesh = Mesh(jax.local_devices(), ("x",))
step = jax.jit(jax.shard_map(per_device, mesh=mesh, in_specs=P("x"), out_specs=P("x")))
x = jnp.arange(mesh.size * N, dtype=jnp.float32)
out = step(x).block_until_ready()  # hangs here on >= 3 GPUs
print(f"OK: {mesh.size}-device step finished; result correct={bool(jnp.array_equal(out, x + 1))}")

# Minimal repro: libcute_dsl_runtime.so leaks its init lock when threads make the first
# call of a --enable-tvm-ffi function concurrently. With 3 threads, one loads the module,
# one returns still holding the lock, and the third spins forever: the script never exits.
import threading

import cutlass.cute as cute

@cute.kernel
def noop():
    pass

@cute.jit
def launch():
    noop().launch(grid=(1, 1, 1), block=(1, 1, 1))

f = cute.compile(launch, options="--enable-tvm-ffi")
barrier = threading.Barrier(3)

def first_call_of_f():
    barrier.wait()
    f()

threads = [threading.Thread(target=first_call_of_f) for _ in range(3)]
for t in threads:
    t.start()
for t in threads:
    t.join()
print("OK: no hang")