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.83
+torchada>=0.1.84
 ```
 
 ### 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.83
+torchada>=0.1.84
 ```
 
 ### 步骤 2：条件导入

pyproject.toml

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

src/torchada/_cpp_ops.py

@@ -93,6 +93,17 @@ def load_cpp_ops(force_reload: bool = False) -> Optional[object]:
     if not is_musa_platform():
         return None
 
+    # torch_musa 2.11.0.post2 contains the native graph-safe implementations.
+    # Do not build or load TorchAda's override extension on fixed releases;
+    # torch_musa/ATen must own dispatch directly.
+    import torch
+
+    from ._patch import _musa_accelerator_overrides_required
+
+    musa_module = getattr(torch, "musa", None)
+    if not _musa_accelerator_overrides_required(getattr(musa_module, "__version__", None)):
+        return None
+
     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,43 @@ 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 _musa_accelerator_overrides_required(getattr(musa_module, "__version__", None)):
+        return
+
+    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 = {}
 
@@ -1727,7 +1772,7 @@ def __getitem__(self, name: str):
 def _musa_accelerator_overrides_required(version) -> bool:
     """Return whether torch.accelerator still needs MUSA memory overrides.
 
-    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,7 +1790,7 @@ 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",
             version,
@@ -1755,7 +1800,7 @@ def _musa_accelerator_overrides_required(version) -> bool:
         return Version(public_version) < Version(_TORCH_MUSA_ACCELERATOR_FIX_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",
             version,
@@ -2122,6 +2167,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