fix(device): register the device factory as the FX device builtin

README.md

@@ -408,7 +408,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain.
 
 ```
 # pyproject.toml or requirements.txt
-torchada>=0.1.93
+torchada>=0.1.94
 ```
 
 ### Step 2: Conditional Import

README_CN.md

@@ -387,7 +387,7 @@ if torchada.is_gpu_device(device):  # 在 CUDA 和 MUSA 上都能工作
 
 ```
 # pyproject.toml 或 requirements.txt
-torchada>=0.1.93
+torchada>=0.1.94
 ```
 
 ### 步骤 2：条件导入

benchmarks/benchmark_history.json

@@ -3,7 +3,7 @@
   "description": "Historical benchmark results for torchada performance tracking",
   "results": [
     {
-      "version": "0.1.93",
+      "version": "0.1.94",
       "date": "2026-01-29",
       "platform": "MUSA",
       "pytorch_version": "2.7.1",

examples/extension_setup.py

@@ -69,7 +69,7 @@ def get_extensions():
     if extensions:
         setup(
             name="my_cuda_extension",
-            version="0.1.93",
+            version="0.1.94",
             ext_modules=extensions,
             cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
             python_requires=">=3.8",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.93"
+version = "0.1.94"
 description = "Adapter package for torch_musa to act exactly like PyTorch CUDA"
 readme = "README.md"
 license = {text = "MIT"}

src/torchada/__init__.py

@@ -24,7 +24,7 @@
     from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME
 """
 
-__version__ = "0.1.93"
+__version__ = "0.1.94"
 
 from . import cuda, utils
 

src/torchada/_patch.py

@@ -895,6 +895,17 @@ def _patch_torch_device():
     except (AttributeError, ImportError, RuntimeError, TypeError):
         logger.debug("Unable to register torch.device as a TorchScript builtin", exc_info=True)
 
+    # FX code generation prints device constants as ``device(...)`` and binds
+    # that name only to the object registered as its ``device`` builtin, which
+    # must be the current torch.device. Registering the wrapper updates both
+    # the builtin table and the reserved-name table FX uses to assign names.
+    try:
+        from torch.fx.graph import _register_custom_builtin
+
+        _register_custom_builtin("device", "from torch import device", DeviceFactoryWrapper)
+    except (AttributeError, ImportError, TypeError):
+        logger.debug("Unable to register torch.device as an FX builtin", exc_info=True)
+
 
 # Store original torch.Generator for patching
 _original_torch_generator = None

src/torchada/csrc/ops.h

@@ -40,7 +40,7 @@
 namespace torchada {
 
 // Version information
-constexpr const char* VERSION = "0.1.93";
+constexpr const char* VERSION = "0.1.94";
 
 // Check if operator override is enabled via environment variable
 inline bool is_override_enabled(const char* op_name) {

tests/test_fx_device_builtin.py

@@ -0,0 +1,293 @@
+"""
+Tests for FX code generation of torch.device constants under the torch.device patch.
+
+FX prints a device constant as ``device(type=...)`` and binds ``device`` to its
+registered builtin. With torch.device replaced by the translating factory, the
+builtin must be that factory, or GraphModules holding a device constant fail
+with ``NameError: name 'device' is not defined`` (Dynamo's eager and aot_eager
+backends run such GraphModules directly).
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+import torchada
+
+BACKENDS = ("eager", "aot_eager")
+DYNAMIC = (False, True)
+
+# Applying the device patch on a non-MUSA platform replaces torch.device for the
+# whole process, so the CPU cases run in a child interpreter.
+_CPU_SCRIPT = r"""
+import json
+import sys
+
+sys.path.insert(0, sys.argv[1])
+
+import torch
+# torchada loads Dynamo (through Inductor) before it replaces torch.device, so
+# Dynamo's constant types hold the original class; keep that order here.
+import torch._dynamo
+import torch.fx
+import torch.fx.graph as fx_graph
+
+
+class _ZerosOnCpu(torch.nn.Module):
+    def forward(self, x):
+        cpu = torch.device("cpu")
+        return x.to(cpu) + torch.zeros_like(x, device=cpu)
+
+
+# Source and import block of GraphModules that hold a device constant.
+def generated_code():
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    empty = graph.call_function(
+        torch.ops.aten.empty.memory_format, ([2, 3],), {"device": torch.empty(0).device}
+    )
+    graph.output((empty, x))
+    modules = (
+        torch.fx.GraphModule(torch.nn.Module(), graph),
+        torch.fx.symbolic_trace(_ZerosOnCpu()),
+    )
+    # The import block is the part of GraphModule.__reduce__ that FX graph caches hash.
+    return [(gm.code, gm.__reduce__()[1][1]) for gm in modules]
+
+
+# A torchada imported at interpreter start (e.g. by a sitecustomize) would make
+# this capture a patched one, so the comparisons below report SKIP instead.
+torchada_preloaded = "torchada" in sys.modules
+unpatched_code = generated_code()
+
+from torchada import _patch
+
+if _patch._original_torch_device is None:
+    _patch._patch_torch_device()
+original_device = _patch._original_torch_device
+results = {}
+
+
+class Skip(Exception):
+    pass
+
+
+def record(name, check):
+    try:
+        check()
+        results[name] = "ok"
+    except Skip as exc:
+        results[name] = "SKIP: {}".format(exc)
+    except Exception as exc:
+        detail = (str(exc).splitlines() or [""])[0]
+        results[name] = "FAILED: {}: {}".format(type(exc).__name__, detail)
+
+
+def check_patched():
+    assert torch.device is _patch.DeviceFactoryWrapper
+    assert isinstance(torch.device("cpu"), torch.device)
+    assert type(torch.device("cpu")) is original_device
+
+
+def check_builtin():
+    assert fx_graph._custom_builtins["device"].obj is torch.device
+    assert fx_graph._illegal_names["device"] is torch.device
+
+
+def check_graph():
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    empty = graph.call_function(
+        torch.ops.aten.empty.memory_format, ([2, 3],), {"device": original_device("cpu")}
+    )
+    from_original = graph.call_function(original_device, ("cpu",))
+    from_factory = graph.call_function(torch.device, ("cpu",))
+    graph.output((empty, from_original, from_factory, x))
+    gm = torch.fx.GraphModule(torch.nn.Module(), graph)
+    assert "device(type='cpu')" in gm.code
+    out, dev_original, dev_factory, _ = gm(torch.ones(1))
+    assert out.device == dev_original == dev_factory == original_device("cpu")
+    assert type(dev_original) is original_device
+    assert type(dev_factory) is original_device
+
+
+def check_code_identical():
+    if torchada_preloaded:
+        raise Skip("torchada was imported before the unpatched capture")
+    patched_code = generated_code()
+    for (code, _), (patched, _) in zip(unpatched_code, patched_code):
+        assert code == patched, (code, patched)
+
+
+def check_import_block_identical():
+    if torchada_preloaded:
+        raise Skip("torchada was imported before the unpatched capture")
+    patched_code = generated_code()
+    for (_, imports), (_, patched) in zip(unpatched_code, patched_code):
+        assert imports == patched, (imports, patched)
+
+
+def check_symbolic_trace():
+    gm = torch.fx.symbolic_trace(_ZerosOnCpu())
+    assert "device(type='cpu')" in gm.code
+    x = torch.ones(3, 2)
+    assert torch.equal(gm(x), _ZerosOnCpu()(x))
+
+
+def fill_rows(x):
+    return torch.empty(x.shape[0], 32, device=x.device).fill_(1.0) + x.sum()
+
+
+def scale_rows(x):
+    return (x * torch.ones(x.shape[0], 1, device=x.device)).sum()
+
+
+def check_compile(backend, dynamic):
+    torch._dynamo.reset()
+    compiled = torch.compile(fill_rows, backend=backend, dynamic=dynamic, fullgraph=True)
+    for rows in (4, 7):
+        x = torch.randn(rows, 3)
+        assert torch.equal(compiled(x), fill_rows(x))
+
+
+def zeros_from_factory(x):
+    return torch.zeros(2, device=torch.device("cpu")) + x[0, :2]
+
+
+def check_factory_in_region(backend):
+    torch._dynamo.reset()
+    compiled = torch.compile(zeros_from_factory, backend=backend, fullgraph=True)
+    x = torch.randn(3, 4)
+    assert torch.equal(compiled(x), zeros_from_factory(x))
+
+
+def check_backward(dynamic):
+    torch._dynamo.reset()
+    compiled = torch.compile(scale_rows, backend="aot_eager", dynamic=dynamic, fullgraph=True)
+    x = torch.randn(5, 3, requires_grad=True)
+    compiled(x).backward()
+    expected = x.detach().clone().requires_grad_(True)
+    scale_rows(expected).backward()
+    assert torch.equal(x.grad, expected.grad)
+
+
+record("patched", check_patched)
+record("builtin", check_builtin)
+record("graph", check_graph)
+record("code_identical", check_code_identical)
+record("import_block_identical", check_import_block_identical)
+record("symbolic_trace", check_symbolic_trace)
+for backend in ("eager", "aot_eager"):
+    for dynamic in (False, True):
+        record("{}-dynamic={}".format(backend, dynamic), lambda: check_compile(backend, dynamic))
+    record("factory-{}".format(backend), lambda: check_factory_in_region(backend))
+for dynamic in (False, True):
+    record("aot_eager-backward-dynamic={}".format(dynamic), lambda: check_backward(dynamic))
+print("RESULTS " + json.dumps(results))
+"""
+
+
+@pytest.fixture(scope="module")
+def patched_cpu_results(tmp_path_factory):
+    src_root = str(Path(torchada.__file__).resolve().parents[1])
+    env = dict(os.environ)
+    env["TORCHINDUCTOR_CACHE_DIR"] = str(tmp_path_factory.mktemp("inductor-cache"))
+    proc = subprocess.run(
+        [sys.executable, "-c", _CPU_SCRIPT, src_root],
+        capture_output=True,
+        text=True,
+        timeout=600,
+        env=env,
+    )
+    lines = [line for line in proc.stdout.splitlines() if line.startswith("RESULTS ")]
+    assert proc.returncode == 0 and lines, proc.stdout[-2000:] + proc.stderr[-4000:]
+    return json.loads(lines[-1][len("RESULTS ") :])
+
+
+@pytest.mark.parametrize(
+    "case",
+    ["patched", "builtin", "graph", "code_identical", "import_block_identical", "symbolic_trace"],
+)
+def test_fx_codegen_with_patched_device(patched_cpu_results, case):
+    result = patched_cpu_results[case]
+    if result.startswith("SKIP"):
+        pytest.skip(result)
+    assert result == "ok"
+
+
+@pytest.mark.parametrize("dynamic", DYNAMIC)
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_compile_device_constant_with_patched_device(patched_cpu_results, backend, dynamic):
+    assert patched_cpu_results[f"{backend}-dynamic={dynamic}"] == "ok"
+
+
+@pytest.mark.parametrize("dynamic", DYNAMIC)
+def test_aot_eager_backward_with_patched_device(patched_cpu_results, dynamic):
+    assert patched_cpu_results[f"aot_eager-backward-dynamic={dynamic}"] == "ok"
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_device_factory_in_compiled_region_with_patched_device(patched_cpu_results, backend):
+    assert patched_cpu_results[f"factory-{backend}"] == "ok"
+
+
+def _fill_rows(x):
+    return torch.empty(x.shape[0], 32, device=x.device).fill_(1.0) + x.sum()
+
+
+@pytest.mark.musa
+def test_fx_device_builtin_is_patched_device():
+    import torch.fx.graph as fx_graph
+
+    assert fx_graph._custom_builtins["device"].obj is torch.device
+    assert fx_graph._illegal_names["device"] is torch.device
+
+
+@pytest.mark.musa
+def test_dynamo_constant_types_hold_original_device():
+    # Dynamo records torch.device as a constant type when it is first imported;
+    # torchada loads Dynamo before replacing torch.device so the original class is kept.
+    import torch._dynamo.utils as dynamo_utils
+
+    from torchada import _patch
+
+    assert _patch._original_torch_device in dynamo_utils.common_constant_types
+
+
+@pytest.mark.musa
+@pytest.mark.parametrize("dynamic", DYNAMIC)
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_compile_device_constant_on_musa(backend, dynamic):
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(_fill_rows, backend=backend, dynamic=dynamic, fullgraph=True)
+        for rows in (4, 7):
+            x = torch.randn(rows, 3, device="cuda")
+            out = compiled(x)
+            assert out.device.type == "musa"
+            assert torch.equal(out, _fill_rows(x))
+    finally:
+        torch._dynamo.reset()
+
+
+def _zeros_from_factory(x):
+    return torch.zeros(2, device=torch.device("cuda")) + x[0, :2]
+
+
+@pytest.mark.musa
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_device_factory_in_compiled_region_on_musa(backend):
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(_zeros_from_factory, backend=backend, fullgraph=True)
+        x = torch.randn(3, 4, device="cuda")
+        out = compiled(x)
+        assert out.device.type == "musa"
+        assert torch.equal(out, _zeros_from_factory(x))
+    finally:
+        torch._dynamo.reset()

tests/test_platform.py

@@ -59,7 +59,7 @@ def test_get_version(self):
 
         version = torchada.get_version()
         assert version == torchada.__version__
-        assert version == "0.1.93"
+        assert version == "0.1.94"
         assert isinstance(version, str)
 
     def test_project_version_matches_runtime_version(self):