Add graph-safe ops and Inductor heuristics

README.md

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

README_CN.md

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

benchmarks/benchmark_history.json

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

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.84"
+version = "0.1.85"
 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.84"
+__version__ = "0.1.85"
 
 from . import cuda, utils
 

src/torchada/_cpp_ops.py

@@ -93,6 +93,15 @@ def load_cpp_ops(force_reload: bool = False) -> Optional[object]:
     if not is_musa_platform():
         return None
 
+    import torch
+
+    from ._patch import _is_pre_torch_musa_2_11_0_post2
+
+    musa_module = getattr(torch, "musa", None)
+    if not _is_pre_torch_musa_2_11_0_post2(getattr(musa_module, "__version__", None)):
+        for op_name in ("multinomial", "log", "log_"):
+            os.environ[f"TORCHADA_DISABLE_OP_OVERRIDE_{op_name}"] = "1"
+
     try:
         import os.path as osp
 

src/torchada/_patch.py

@@ -65,6 +65,14 @@ def _patch_something():
     return func
 
 
+@patch_function
+def _patch_visible_devices_env():
+    if "MUSA_VISIBLE_DEVICES" in os.environ:
+        os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["MUSA_VISIBLE_DEVICES"]
+    else:
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+
+
 def requires_import(*module_names: str) -> Callable[[Callable], Callable]:
     """
     Decorator to guard a patch function with import checks.
@@ -108,6 +116,50 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+@patch_function
+@requires_import("torch._inductor.template_heuristics.registry")
+def _patch_inductor_template_heuristics():
+    """Reuse CUDA Inductor template heuristics for CUDA-compatible MUSA templates."""
+    if not is_musa_platform():
+        return
+
+    musa_module = getattr(torch, "musa", None)
+    if not _is_pre_torch_musa_2_11_0_post2(getattr(musa_module, "__version__", None)):
+        return
+
+    from torch._inductor.codegen.common import init_backend_registration
+
+    # torch_musa registers its native MUSA heuristics lazily from this entry
+    # point. Run it before inspecting the registry so TorchAda only fills keys
+    # that remain absent after native backend registration.
+    init_backend_registration()
+
+    import torch._inductor.template_heuristics.registry as registry
+
+    heuristic_registry = getattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", None)
+    if not isinstance(heuristic_registry, dict):
+        return
+
+    changed = False
+    for key, heuristic_class in list(heuristic_registry.items()):
+        if not isinstance(key, tuple) or len(key) != 3:
+            continue
+        template_name, device_type, op_name = key
+        if device_type != "cuda":
+            continue
+        if not isinstance(template_name, str) or not template_name.startswith("triton::"):
+            continue
+        musa_key = (template_name, "musa", op_name)
+        if musa_key not in heuristic_registry:
+            heuristic_registry[musa_key] = heuristic_class
+            changed = True
+
+    if changed:
+        heuristic_cache = getattr(registry, "_HEURISTIC_CACHE", None)
+        if isinstance(heuristic_cache, dict):
+            heuristic_cache.clear()
+
+
 # Cache for translated device strings - avoids repeated string operations
 _device_str_cache = {}
 
@@ -1721,13 +1773,13 @@ def __getitem__(self, name: str):
 _original_ctypes_CDLL = None
 
 
-_TORCH_MUSA_ACCELERATOR_FIX_VERSION = "2.11.0.post2"
+_TORCH_MUSA_POST2_VERSION = "2.11.0.post2"
 
 
-def _musa_accelerator_overrides_required(version) -> bool:
-    """Return whether torch.accelerator still needs MUSA memory overrides.
+def _is_pre_torch_musa_2_11_0_post2(version) -> bool:
+    """Return whether the torch_musa version predates 2.11.0.post2.
 
-    torch_musa 2.11.0.post2 fixes the unified accelerator memory APIs.  Older
+    torch_musa 2.11.0.post2 fixes the unified accelerator memory APIs. Older
     releases still need torchada to force those calls through torch.musa.
     Ignore the local version suffix (for example ``+musa5.2.0``), because it
     identifies the MUSA stack build rather than the torch_musa fix level.
@@ -1745,19 +1797,19 @@ def _musa_accelerator_overrides_required(version) -> bool:
         from torch._vendor.packaging.version import InvalidVersion, Version
     except ImportError:
         # If the parser is unavailable, keep the workaround enabled: disabling
-        # it could re-expose the allocator failure this gate fixes.
+        # it could re-expose the failure this gate fixes.
         logger.warning(
-            "Unable to parse torch_musa version %r; retaining accelerator memory overrides",
+            "Unable to parse torch_musa version %r; retaining compatibility patches",
             version,
         )
         return True
     try:
-        return Version(public_version) < Version(_TORCH_MUSA_ACCELERATOR_FIX_VERSION)
+        return Version(public_version) < Version(_TORCH_MUSA_POST2_VERSION)
     except InvalidVersion:
         # An unknown or malformed version must keep the workaround enabled:
-        # disabling it could re-expose the allocator failure this gate fixes.
+        # disabling it could re-expose the failure this gate fixes.
         logger.warning(
-            "Unable to parse torch_musa version %r; retaining accelerator memory overrides",
+            "Unable to parse torch_musa version %r; retaining compatibility patches",
             version,
         )
         return True
@@ -1843,7 +1895,7 @@ def __init__(self, original_accel, musa_module):
         # torch._C._accelerator_* without dispatching to the MUSA allocator.
         # Newer versions provide working unified accelerator implementations,
         # so preserve those instead of forcing the torch.musa compatibility path.
-        if _musa_accelerator_overrides_required(getattr(musa_module, "__version__", None)):
+        if _is_pre_torch_musa_2_11_0_post2(getattr(musa_module, "__version__", None)):
             for name in self._MUSA_OVERRIDES:
                 musa_name = self._REMAP_ATTRS.get(name, name)
                 if hasattr(original_accel, name) and hasattr(musa_module, musa_name):
@@ -2122,6 +2174,7 @@ def apply_patches():
     - torch.cuda.nccl -> torch.musa.mccl
     - torch.amp.autocast(device_type='cuda') -> 'musa'
     - torch.utils.cpp_extension (CUDAExtension, BuildExtension) -> MUSA versions
+    - CUDA_VISIBLE_DEVICES -> MUSA_VISIBLE_DEVICES environment fallback
     - torch._inductor.autotune_process.CUDA_VISIBLE_DEVICES -> MUSA_VISIBLE_DEVICES
     - torch.accelerator.synchronize() -> torch.musa.synchronize()
     - torch.accelerator context managers (device_index, stream) for forward compatibility