From da9a3c502a5f259625ee09df26f0d2751887fc22 Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:18:46 +0900 Subject: [PATCH 1/3] Grad mode and the torch-function mode stack are per thread, as upstream's are (#76) Both were process-wide: torch.no_grad() in one thread stopped grad_fn recording in another, and torch.device('meta') on one thread moved another thread's factories to meta. Wrong with the GIL on, not only on 3.15t. tensor.rs GRAD_ENABLED is a thread_local Cell; the bootstrap's grad flag and _MODE_STACK sit on threading.local. tests/api/test_threadstate.py: 3 tests, green on upstream torch, red on the shim before, green after. --- .github/scripts/gate_suites.py | 1 + tests/api/test_threadstate.py | 103 ++++++++++++++++++++++ torchnative/rust/torch_c/src/bootstrap.py | 54 +++++++++--- torchnative/rust/torch_c/src/tensor.rs | 21 +++-- 4 files changed, 158 insertions(+), 21 deletions(-) create mode 100644 tests/api/test_threadstate.py diff --git a/.github/scripts/gate_suites.py b/.github/scripts/gate_suites.py index 105afdd9..b3a8412d 100644 --- a/.github/scripts/gate_suites.py +++ b/.github/scripts/gate_suites.py @@ -83,6 +83,7 @@ SUITES = { "_support/test_shim.py": BOTH, "api/test_import.py": BOTH, + "api/test_threadstate.py": BOTH, "api/test_tnnamespace.py": BOTH, "api/test_tntransformers.py": BOTH, "bindings/test_aliasinc.py": BOTH, diff --git a/tests/api/test_threadstate.py b/tests/api/test_threadstate.py new file mode 100644 index 00000000..214ca0c2 --- /dev/null +++ b/tests/api/test_threadstate.py @@ -0,0 +1,103 @@ +"""State that upstream keeps per thread must not leak between threads (#76). + +Upstream keeps grad mode (`c10/core/GradMode.h`) and the `torch.device(...)` +mode stack per thread. The shim held both process-wide, so `torch.no_grad()` +in one thread silently stopped another thread's `grad_fn` recording. That is +wrong with the GIL on, not only on a free-threaded interpreter: the GIL +serialises bytecode, it does not make a global per thread. + +Each test parks one thread inside a context manager (a no-grad region, a device +mode) and asserts from another thread that the state it sees is its own. The +same tests pass against upstream torch, which is how they were checked before +the fix. + +What this cannot see: a race that leaves the right answer. These tests check +isolation, not memory safety; the RNG stream free race of #76 needs the +free-threaded hammer in tests/api/test_freethreaded.py. +""" + +import os +import sys +import threading + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "_support")) +sys.path.insert(0, os.path.join(str(next(p for p in __import__("pathlib").Path(__file__).resolve().parents if (p / "AGENTS.md").is_file())), "torchnative", "python")) + +import torch # noqa: E402 +import _skip # noqa: E402 + + +def _park(enter, seen): + """Run `enter()` in a thread, record what it sees inside, hold it open.""" + inside, release = threading.Event(), threading.Event() + + def body(): + with enter(): + seen["inside"] = torch.is_grad_enabled() + inside.set() + release.wait(10) + seen["after"] = torch.is_grad_enabled() + + t = threading.Thread(target=body) + t.start() + assert inside.wait(10), "parked thread never entered" + return t, release + + +def test_no_grad_in_another_thread_does_not_stop_recording_here(): + seen = {} + t, release = _park(torch.no_grad, seen) + try: + x = torch.ones(2, requires_grad=True) + y = x * 2 + here = torch.is_grad_enabled() + finally: + release.set() + t.join(10) + assert seen["inside"] is False, seen + assert here is True, "grad mode leaked from another thread's no_grad" + assert y.grad_fn is not None, "grad_fn not recorded: another thread's no_grad applied here" + + +def test_overlapping_no_grad_regions_restore_each_threads_own_mode(): + # Thread A enters no_grad, then the main thread enters and leaves its own. + # With one global flag, the main thread's exit restores A's False (or A's + # exit restores the main thread's), and someone is left with the wrong mode. + seen = {} + t, release = _park(torch.no_grad, seen) + try: + with torch.no_grad(): + assert torch.is_grad_enabled() is False + after_main = torch.is_grad_enabled() + finally: + release.set() + t.join(10) + assert after_main is True, "main thread left with another thread's mode" + assert seen["after"] is True, seen + + +def test_a_device_mode_in_another_thread_does_not_apply_here(): + inside, release = threading.Event(), threading.Event() + + def body(): + with torch.device("meta"): + inside.set() + release.wait(10) + + t = threading.Thread(target=body) + t.start() + try: + assert inside.wait(10) + dev = torch.empty(2).device + finally: + release.set() + t.join(10) + assert dev.type == "cpu", f"another thread's torch.device('meta') applied here: {dev}" + + +if __name__ == "__main__": + assert hasattr(torch._C, "_aten_implemented") or os.environ.get("TORCHNATIVE_ORACLE") == "1", \ + "imported upstream torch, not the shim" + raise SystemExit(1 if _skip.run_tests( + [(n, f) for n, f in list(globals().items()) if n.startswith("test_") and callable(f)], + "api/test_threadstate") else 0) diff --git a/torchnative/rust/torch_c/src/bootstrap.py b/torchnative/rust/torch_c/src/bootstrap.py index b47cd319..ad5c3cc0 100644 --- a/torchnative/rust/torch_c/src/bootstrap.py +++ b/torchnative/rust/torch_c/src/bootstrap.py @@ -4969,13 +4969,40 @@ def _get_tensor_metadata(tensor): # `DeviceContext`; `torch/__init__.py` has `set_default_device` and # `get_default_device`. All five `_C` names they bottom out in are below. # -# **The stack is process-wide where upstream's is thread-local.** -# `PythonTorchFunctionTLS` is per-thread; this is one list. Recorded rather than -# fixed: a mode entered on one thread would be seen by another, which upstream -# would not do. Nothing in this shim's measured paths is multi-threaded, and the -# fix (a `threading.local`) is a two-line change in this block if that stops -# being true. -_MODE_STACK: list = [] +# **The stack is per thread**, as upstream's `PythonTorchFunctionTLS` is. It was +# one process-wide list, so `with torch.device("meta")` on one thread made +# another thread's `torch.empty` land on meta (#76). The proxy keeps the +# list-shaped spelling every caller below uses, over a `threading.local`. +class _ThreadModeStack: + __slots__ = ("_tls",) + + def __init__(self): + self._tls = threading.local() + + def _list(self) -> list: + try: + return self._tls.stack + except AttributeError: + self._tls.stack = [] + return self._tls.stack + + def append(self, mode): + self._list().append(mode) + + def pop(self): + return self._list().pop() + + def __len__(self): + return len(self._list()) + + def __bool__(self): + return bool(self._list()) + + def __getitem__(self, index): + return self._list()[index] + + +_MODE_STACK = _ThreadModeStack() def _through_torch_function_modes(func, args, kwargs): @@ -7181,11 +7208,10 @@ def _install_grad_mode(module, varfns) -> None: `torch.is_grad_enabled()` and `torch._C.is_grad_enabled()` disagree. """ state = { - "grad": True, # `torch/autograd/grad_mode.py:340` and `:393`. Both are real # backend-configuration switches (thread pool, layout enforcement) with # nothing behind them here, and both are context managers that restore - # what they read -- so, like `grad`, they have to round-trip. + # what they read -- so, like grad mode, they have to round-trip. "multithreading": True, "layout_enforcement": False, # `torch.is_inference_mode_enabled()`. Upstream spells this ONLY at @@ -7194,7 +7220,7 @@ def _install_grad_mode(module, varfns) -> None: # as the table-less refusal made `fake_tensor.py:1801` stop the whole of # `torch.export` on a predicate that has an obvious answer. # - # It lives HERE, in the same dict as `grad`, rather than as a constant, + # It lives HERE, in this dict rather than as a constant, # because docs/graph/EXPORT.md §2.2 is about exactly the other choice: a # constant `False` would make `with torch.inference_mode():` a block # that enters, reports itself absent, and changes nothing. The setter @@ -7203,8 +7229,12 @@ def _install_grad_mode(module, varfns) -> None: "inference": False, } + # Grad mode is per thread, as upstream's `GradMode` is (#76); the other + # entries in `state` stay process-wide, as their upstream switches are. + grad_tls = threading.local() + def is_grad_enabled(): - return state["grad"] + return getattr(grad_tls, "grad", True) def _set_grad_enabled(mode): # Two writes, one source of truth. `state` stays the value @@ -7214,7 +7244,7 @@ def _set_grad_enabled(mode): # differentiate. See `tensor.rs`'s `GRAD_ENABLED`, and # `test_no_grad_gates_grad_fn` for the two being checked against each # other rather than assumed equal. - state["grad"] = bool(mode) + grad_tls.grad = bool(mode) module._shim_set_grad_enabled_flag(bool(mode)) def is_inference_mode_enabled(): diff --git a/torchnative/rust/torch_c/src/tensor.rs b/torchnative/rust/torch_c/src/tensor.rs index 16708a6a..b76b209d 100644 --- a/torchnative/rust/torch_c/src/tensor.rs +++ b/torchnative/rust/torch_c/src/tensor.rs @@ -4048,10 +4048,13 @@ pub fn has_storage(value: &Bound<'_, PyAny>) -> PyResult { /// of an `AtomicBool` is the same shape as `capture::is_active`, which /// docs/graph/CAPTURE.md §7 already measured at the same door. /// -/// `Ordering::Relaxed` for the same reason capture uses it: there is nothing -/// else for this flag to be ordered *against*. A thread that flips it and then -/// dispatches does both under the GIL. -static GRAD_ENABLED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(true); +/// **Per thread**, as upstream's `GradMode` is (`c10/core/GradMode.h`). It was +/// one process-wide `AtomicBool`, so `torch.no_grad()` in one thread stopped +/// another thread's `grad_fn` recording, with the GIL on (#76). A `Cell` read is +/// no dearer at the door than the relaxed atomic load it replaces. +thread_local! { + static GRAD_ENABLED: std::cell::Cell = const { std::cell::Cell::new(true) }; +} /// `torch._C._set_throw_on_mutable_data_ptr(tensor)`. /// @@ -4111,7 +4114,7 @@ pub fn throws_on_mutable_data_ptr(tensor: PyRef<'_, PyTensorBase>) -> bool { #[pyfunction] #[pyo3(name = "_shim_set_grad_enabled_flag")] pub fn set_grad_enabled_flag(value: bool) { - GRAD_ENABLED.store(value, std::sync::atomic::Ordering::Relaxed); + GRAD_ENABLED.with(|g| g.set(value)); } /// Grad mode off for the duration of a scope, restored on drop. @@ -4125,20 +4128,20 @@ pub struct NoGradGuard(bool); impl NoGradGuard { pub fn enter() -> Self { - Self(GRAD_ENABLED.swap(false, std::sync::atomic::Ordering::Relaxed)) + Self(GRAD_ENABLED.with(|g| g.replace(false))) } } impl Drop for NoGradGuard { fn drop(&mut self) { - GRAD_ENABLED.store(self.0, std::sync::atomic::Ordering::Relaxed); + GRAD_ENABLED.with(|g| g.set(self.0)); } } #[pyfunction] #[pyo3(name = "_shim_grad_enabled_flag")] pub fn grad_enabled_flag() -> bool { - GRAD_ENABLED.load(std::sync::atomic::Ordering::Relaxed) + GRAD_ENABLED.with(|g| g.get()) } /// Ops whose output is **not** a differentiable function of their tensor @@ -4289,7 +4292,7 @@ pub fn mark_from_op( kwargs: Option<&Bound<'_, pyo3::types::PyDict>>, out: &Py, ) -> bool { - if !GRAD_ENABLED.load(std::sync::atomic::Ordering::Relaxed) { + if !GRAD_ENABLED.with(|g| g.get()) { return false; } if NOT_DIFFERENTIABLE.contains(&op) { From f506e52bf3d8114cf80c40bb8e1edaac23c7a921 Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:42:39 +0900 Subject: [PATCH 2/3] RNG streams are Arc-held, so freeing a generator cannot pull a stream out from under a live draw (#76) Streams were Box::leak'ed and reclaimed by Box::from_raw in stream_free, sound only while a kernel's guard and the generator's __del__ could not overlap under the GIL. stream() now returns a GenGuard that holds its own Arc clone; stream_free drops the table's reference. The new cargo test frees a locked stream and reuses its memory: SIGSEGV on the old code, green now. --- torchnative/rust/torch_c/src/rng.rs | 97 ++++++++++++++++++++++------- 1 file changed, 76 insertions(+), 21 deletions(-) diff --git a/torchnative/rust/torch_c/src/rng.rs b/torchnative/rust/torch_c/src/rng.rs index 388c1e25..76ce489c 100644 --- a/torchnative/rust/torch_c/src/rng.rs +++ b/torchnative/rust/torch_c/src/rng.rs @@ -236,19 +236,22 @@ pub fn default_generator() -> MutexGuard<'static, CpuGenerator> { /// Every `torch.Generator()` owns one `CpuGenerator`, kept here under a small /// integer id that the Python object carries as `_shim_gen_id`. /// -/// The stream is leaked (`Box::leak`) so a kernel can hold a -/// `MutexGuard<'static, CpuGenerator>` -- the same type `default_generator()` -/// returns -- and every kernel body stays as it was upstream-shaped: it takes -/// `&mut CpuGenerator` and does not care whose. `stream_free` reclaims it when -/// the Python object dies. That is sound because a kernel holds its guard only -/// while it runs, under the GIL, and `__del__` cannot run in the middle of one. +/// Each stream is an `Arc`, and `stream()` hands out a guard that holds its own +/// clone of it. It used to be `Box::leak`ed, with `stream_free` reclaiming it by +/// `Box::from_raw`, which was sound only because a kernel's guard and the +/// Python `__del__` that frees the stream could not both run at once "under +/// the GIL". On a free-threaded interpreter they can, and the guard would point +/// at freed memory (#76). Now `stream_free` only drops the table's reference: +/// a kernel that is still drawing keeps the generator alive until it finishes. /// /// `None` is the default generator, so a call that names no generator and a /// call that names `torch.default_generator` take the same path. -static STREAMS: OnceLock>>> = OnceLock::new(); +type Stream = std::sync::Arc>; + +static STREAMS: OnceLock>> = OnceLock::new(); static NEXT_STREAM: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); -fn streams() -> MutexGuard<'static, HashMap>> { +fn streams() -> MutexGuard<'static, HashMap> { match STREAMS.get_or_init(|| Mutex::new(HashMap::new())).lock() { Ok(g) => g, Err(p) => p.into_inner(), @@ -262,18 +265,38 @@ pub const DEFAULT_GENERATOR_SEED: u64 = 67280421310721; pub fn stream_new(seed: u64) -> u64 { let id = NEXT_STREAM.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let gen: &'static Mutex = Box::leak(Box::new(Mutex::new(CpuGenerator::new(seed)))); - streams().insert(id, gen); + streams().insert(id, std::sync::Arc::new(Mutex::new(CpuGenerator::new(seed)))); id } pub fn stream_free(id: u64) { - if let Some(gen) = streams().remove(&id) { - // SAFETY: `gen` came from `Box::leak` in `stream_new`, it was just - // removed from the only table that handed out references, and kernels - // lock a stream only while holding the GIL, which the caller of - // `stream_free` (a Python `__del__`) also holds -- so no guard is live. - drop(unsafe { Box::from_raw(gen as *const Mutex as *mut Mutex) }); + // Drops the table's reference only. A guard handed out by `stream()` + // holds its own, so the generator outlives this call if a kernel is + // still drawing from it. + streams().remove(&id); +} + +/// A locked generator: the default one, or a stream kept alive by the guard. +/// +/// Field order is the soundness argument. `guard` borrows from the `Arc` in +/// `_keep` (its `'static` is a lie told to the borrow checker), and fields drop +/// in declaration order, so the lock is released before the last reference +/// this guard owns can go. +pub struct GenGuard { + guard: MutexGuard<'static, CpuGenerator>, + _keep: Option, +} + +impl std::ops::Deref for GenGuard { + type Target = CpuGenerator; + fn deref(&self) -> &CpuGenerator { + &self.guard + } +} + +impl std::ops::DerefMut for GenGuard { + fn deref_mut(&mut self) -> &mut CpuGenerator { + &mut self.guard } } @@ -281,17 +304,23 @@ pub fn stream_free(id: u64) { /// default when `id` is `None`. An unknown id is a freed generator -- a bug on /// the Python side -- and panics rather than quietly drawing from another /// stream. -pub fn stream(id: Option) -> MutexGuard<'static, CpuGenerator> { +pub fn stream(id: Option) -> GenGuard { let Some(id) = id else { - return default_generator(); + return GenGuard { guard: default_generator(), _keep: None }; }; - let gen: &'static Mutex = *streams() + let keep: Stream = streams() .get(&id) + .cloned() .unwrap_or_else(|| panic!("torch._C shim: generator stream {id} was freed")); - match gen.lock() { + // SAFETY: the pointer is into the heap allocation `keep` owns, and `keep` + // moves into the same `GenGuard`, declared after `guard`, so it is dropped + // after the guard; the `Arc` allocation does not move when `keep` does. + let m: &'static Mutex = unsafe { &*std::sync::Arc::as_ptr(&keep) }; + let guard = match m.lock() { Ok(g) => g, Err(p) => p.into_inner(), - } + }; + GenGuard { guard, _keep: Some(keep) } } // --------------------------------------------------------------------------- @@ -835,3 +864,29 @@ pub fn register(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<() m.add_function(wrap_pyfunction!(shim_reseed, m)?)?; Ok(()) } + +#[cfg(test)] +mod stream_lifetime_tests { + use super::*; + + /// #76: a kernel's guard must keep its stream alive when the Python side + /// frees the generator mid-draw. Freeing, then allocating same-sized + /// streams with other seeds, is how the old `Box::leak`/`Box::from_raw` + /// scheme let the guard read another generator's state. + #[test] + fn a_live_guard_keeps_drawing_its_own_stream_after_stream_free() { + let seed = 12345; + let id = stream_new(seed); + let mut guard = stream(Some(id)); + stream_free(id); + let others: Vec = (0..256).map(|i| stream_new(1_000_000 + i)).collect(); + let mut reference = CpuGenerator::new(seed); + for _ in 0..64 { + assert_eq!(guard.random64(), reference.random64(), "guard drew from a reused allocation"); + } + drop(guard); + for o in others { + stream_free(o); + } + } +} From 992c69c9c5c3081f53957f174e9610b32ba7f8fc Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:25:39 +0900 Subject: [PATCH 3/3] Capture's activity flag counts open regions instead of one bool (#76); test_threadstate sets TORCH_USE_RTLD_GLOBAL CAPTURING was cleared by whichever thread ended its capture first, so a thread still mid-capture stopped recording and its later ops vanished from its trace. It is now an AtomicUsize of open regions (one per thread at most, since capture_begin refuses nesting); the door still pays one relaxed load. test_threadstate.py never ran in the gate before this commit: it died at import torch for want of TORCH_USE_RTLD_GLOBAL, which every other suite that imports torch sets. Now 4 ok in the gate (2010 -> 2014). --- tests/api/test_threadstate.py | 35 +++++++++++++++++++++++++ torchnative/rust/torch_c/src/capture.rs | 16 ++++++++--- 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/tests/api/test_threadstate.py b/tests/api/test_threadstate.py index 214ca0c2..cca278a7 100644 --- a/tests/api/test_threadstate.py +++ b/tests/api/test_threadstate.py @@ -20,6 +20,9 @@ import sys import threading +# The shim imports upstream extension libraries that resolve only with global +# symbols, as every other suite that imports torch sets. +os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1") sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "_support")) sys.path.insert(0, os.path.join(str(next(p for p in __import__("pathlib").Path(__file__).resolve().parents if (p / "AGENTS.md").is_file())), "torchnative", "python")) @@ -95,6 +98,38 @@ def body(): assert dev.type == "cpu", f"another thread's torch.device('meta') applied here: {dev}" +def test_ending_a_capture_on_one_thread_does_not_stop_another_threads(): + # #76: the "is anyone recording" flag was one bool, cleared by whichever + # thread ended its capture first, while the recorders are per thread. The + # other thread's later ops then skipped recording and vanished from its + # trace. Here the main thread captures and ends while thread A is mid-capture. + C = torch._C + if not hasattr(C, "_capture_begin"): + raise _skip.Skip("no capture API (upstream torch)") + began, release, out = threading.Event(), threading.Event(), {} + + def body(): + x = torch.ones(3) + C._capture_begin([x]) + began.set() + release.wait(10) + y = x.neg() + out["ops"] = [n["op"] for n in C._capture_end(y).nodes] + + t = threading.Thread(target=body) + t.start() + try: + assert began.wait(10) + a = torch.ones(2) + C._capture_begin([a]) + C._capture_end(a.exp()) + finally: + release.set() + t.join(10) + assert any("neg" in op for op in out.get("ops", [])), \ + f"thread A's op was dropped after another thread ended its capture: {out}" + + if __name__ == "__main__": assert hasattr(torch._C, "_aten_implemented") or os.environ.get("TORCHNATIVE_ORACLE") == "1", \ "imported upstream torch, not the shim" diff --git a/torchnative/rust/torch_c/src/capture.rs b/torchnative/rust/torch_c/src/capture.rs index 3bc7280c..a65ec7db 100644 --- a/torchnative/rust/torch_c/src/capture.rs +++ b/torchnative/rust/torch_c/src/capture.rs @@ -61,7 +61,13 @@ use crate::tensor::PyTensorBase; /// Global rather than thread-local because reading a `thread_local!` costs a /// TLS lookup, and the recorder itself is thread-local anyway -- a second /// thread that takes the branch finds no recorder and falls straight back out. -static CAPTURING: AtomicBool = AtomicBool::new(false); +/// +/// A **count** of open regions, not a bool. As a bool it was cleared by +/// whichever thread ended its capture first, and every later op of a thread +/// still mid-capture skipped recording and vanished from its trace (#76). +/// `capture_begin` refuses a second region on the same thread, so each thread +/// contributes at most one, and only a thread that took its recorder decrements. +static CAPTURING: AtomicUsize = AtomicUsize::new(0); thread_local! { static RECORDER: RefCell> = const { RefCell::new(None) }; @@ -76,7 +82,7 @@ thread_local! { #[inline(always)] pub fn is_active() -> bool { - CAPTURING.load(Ordering::Relaxed) + CAPTURING.load(Ordering::Relaxed) != 0 } /// Whether the eager tape records at all. @@ -1738,7 +1744,7 @@ pub fn capture_begin(py: Python<'_>, inputs: &Bound<'_, PyAny>) -> PyResult<()> let _ = py; RECORDER.with(|cell| *cell.borrow_mut() = Some(rec)); - CAPTURING.store(true, Ordering::Relaxed); + CAPTURING.fetch_add(1, Ordering::Relaxed); Ok(()) } @@ -1753,7 +1759,9 @@ pub fn capture_abandon() -> PyResult<()> { fn take_recorder() -> PyResult { let taken = RECORDER.with(|cell| cell.borrow_mut().take()); - CAPTURING.store(false, Ordering::Relaxed); + if taken.is_some() { + CAPTURING.fetch_sub(1, Ordering::Relaxed); + } taken.ok_or_else(not_recording) }