From 32315136b41441302a9530fcedf97f98a4f77cf5 Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:19:52 +0900 Subject: [PATCH] Dispatch registrations: build the table locally, publish it whole under a lock (#76 B6) _scan_dispatch_registrations filled the shared dict in place, and its "already built" check was true as soon as the keys were inserted, before any name was. Eight threads asking at once at a 1 us switch interval got back seven empty lists and one 743-name list. The table is now built into a fresh dict and bound to the global under a lock, so a reader sees no table or the finished one. - test_threadstate: 8 threads behind a Barrier, three fresh subprocesses; red before (seen [0,0,0,0,0,0,0,743]), green after. - 3.15t with PYTHON_GIL=0 (free-threaded _C, .scratch stage): 743 x 8, three runs of three. --- tests/api/test_threadstate.py | 51 +++++++++++++++++++++++ torchnative/rust/torch_c/src/bootstrap.py | 37 ++++++++++++---- 2 files changed, 79 insertions(+), 9 deletions(-) diff --git a/tests/api/test_threadstate.py b/tests/api/test_threadstate.py index 3bf6585..a9a4d50 100644 --- a/tests/api/test_threadstate.py +++ b/tests/api/test_threadstate.py @@ -130,6 +130,57 @@ def body(): f"thread A's op was dropped after another thread ended its capture: {out}" +_REGISTRATIONS_CHILD = r""" +import json, sys, threading +import torch +assert hasattr(torch._C, "_aten_implemented"), "the child imported upstream, not the shim" +sys.setswitchinterval(1e-6) +key = "CompositeImplicitAutograd" +n = 8 +barrier = threading.Barrier(n) +seen = [None] * n + +def body(i): + barrier.wait() + seen[i] = torch._C._dispatch_get_registrations_for_dispatch_key(key) + +threads = [threading.Thread(target=body, args=(i,)) for i in range(n)] +for t in threads: + t.start() +for t in threads: + t.join(60) +final = torch._C._dispatch_get_registrations_for_dispatch_key(key) +print(json.dumps({"seen": [len(s) if s is not None else None for s in seen], + "same": [s == final for s in seen], + "final": len(final), "unique": len(set(final))})) +""" + + +def test_dispatch_registrations_built_by_eight_threads_at_once_are_whole(): + # #76 B6: the table was filled lazily into the shared dict, and the + # `if _DISPATCH_REGISTRATIONS: return` check is true as soon as the keys + # are inserted, before any name is. A second thread got back a partial + # table, and two threads filling at once appended every name twice. Each + # run is a fresh subprocess because the table is built once per process. + import json + import subprocess + if not hasattr(torch._C, "_aten_implemented"): + raise _skip.Skip("upstream torch: no file-built registration table") + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join(p for p in sys.path[:2]) + for attempt in range(3): + proc = subprocess.run([sys.executable, "-c", _REGISTRATIONS_CHILD], + capture_output=True, text=True, timeout=600, env=env) + lines = [l for l in proc.stdout.splitlines() if l.startswith("{")] + assert proc.returncode == 0 and lines, proc.stderr[-2000:] + out = json.loads(lines[-1]) + assert out["final"] > 0, out + assert out["unique"] == out["final"], ( + f"run {attempt}: the table holds duplicate names, two threads filled it at once: {out}") + assert all(out["same"]), ( + f"run {attempt}: a thread got back a different table than the finished one: {out}") + + def test_backward_on_another_thread_than_the_forward(): # #78: upstream's graph hangs off the tensors, so a forward recorded in one # thread can be differentiated from another. The shim's eager tape was per diff --git a/torchnative/rust/torch_c/src/bootstrap.py b/torchnative/rust/torch_c/src/bootstrap.py index b575993..6d928d1 100644 --- a/torchnative/rust/torch_c/src/bootstrap.py +++ b/torchnative/rust/torch_c/src/bootstrap.py @@ -1711,7 +1711,9 @@ def _aten_tags(module, qualname: str, overload: str) -> list: #: `` -> `["aten::[.]", ...]`, built on first use. -_DISPATCH_REGISTRATIONS: dict = {} +#: `None` until a whole table is published; see `_scan_dispatch_registrations`. +_DISPATCH_REGISTRATIONS: dict | None = None +_DISPATCH_REGISTRATIONS_LOCK = threading.Lock() #: The keys `native_functions.yaml` is authoritative about. These four are #: *alias* keys: they say how an op is put together, not which backend runs it, @@ -2011,18 +2013,35 @@ def _scan_dispatch_registrations() -> dict: builtin -- the same 176-name gap `_is_absent_inplace_variant` documents, seen from another side. """ - if _DISPATCH_REGISTRATIONS: + global _DISPATCH_REGISTRATIONS + table = _DISPATCH_REGISTRATIONS + if table is not None: + return table + with _DISPATCH_REGISTRATIONS_LOCK: + if _DISPATCH_REGISTRATIONS is None: + _DISPATCH_REGISTRATIONS = _build_dispatch_registrations() return _DISPATCH_REGISTRATIONS - for key in _FILE_DECLARED_DISPATCH_KEYS: - _DISPATCH_REGISTRATIONS[key] = [] + + +def _build_dispatch_registrations() -> dict: + """The scan `_scan_dispatch_registrations` publishes, into a fresh dict. + + #76 B6: this used to fill the shared dict in place, and the "already + built" check was true as soon as the keys were inserted, before any name + was. Eight threads asking at once got back seven empty lists, and two + threads filling at once appended every name twice. Built here, then + published whole under `_DISPATCH_REGISTRATIONS_LOCK`, a reader sees either + no table or the finished one. + """ + table = {key: [] for key in _FILE_DECLARED_DISPATCH_KEYS} source = _native_functions_source() if not os.path.isabs(source or ""): - return _DISPATCH_REGISTRATIONS + return table try: with open(source, encoding="utf-8") as handle: lines = handle.read().splitlines() except OSError: - return _DISPATCH_REGISTRATIONS + return table state = {"name": None, "keys": set(), "dispatch": False, "structured": False, "delegate": False, "in_block": False} @@ -2035,8 +2054,8 @@ def flush(): if not state["dispatch"] and not state["structured"] and not state["delegate"]: keys.add("CompositeImplicitAutograd") for key in keys: - if key in _DISPATCH_REGISTRATIONS: - _DISPATCH_REGISTRATIONS[key].append(f"aten::{name}") + if key in table: + table[key].append(f"aten::{name}") for line in lines: if line.startswith("- func:"): @@ -2072,7 +2091,7 @@ def flush(): for key in stripped.split(":", 1)[0].split(","): state["keys"].add(key.strip()) flush() - return _DISPATCH_REGISTRATIONS + return table def _dispatch_registrations(key: str) -> list: