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: