diff --git a/src/torchada/utils/cpp_extension.py b/src/torchada/utils/cpp_extension.py index ec7cbdc..5cedd61 100644 --- a/src/torchada/utils/cpp_extension.py +++ b/src/torchada/utils/cpp_extension.py @@ -294,9 +294,7 @@ def _is_configured_exclude_dir(path: str) -> bool: def _path_overlaps_any(path: str, roots: List[str]) -> bool: """Return whether ``path`` contains or is contained by a protected root.""" - return any( - _path_is_within(path, root) or _path_is_within(root, path) for root in roots - ) + return any(_path_is_within(path, root) or _path_is_within(root, path) for root in roots) def _validate_portable_symlinks(source_dir: str) -> None: @@ -366,9 +364,9 @@ def _narrow_cuda_header_mapping(mapping_rule): narrowed = [(key, value) for key, value in mapping_rule if key != "cuda.h"] narrowed.extend( [ - ('#include ', '#include '), + ("#include ", "#include "), ('#include "cuda.h"', '#include "musa.h"'), - ('#include ', '#include '), + ("#include ", "#include "), ('#include "torch/cuda.h"', '#include "torch/musa.h"'), ] ) @@ -405,11 +403,7 @@ def _replace_porting_line(line, mapping_rule): header = line[start:end] if key == "nvjpeg" and header == "nvjpeg.h": header = f"{value}.h" - line = ( - line[:start].replace(key, value) - + header - + line[end:].replace(key, value) - ) + line = line[:start].replace(key, value) + header + line[end:].replace(key, value) continue line = line.replace(key, value) return line @@ -833,79 +827,88 @@ def _port_cuda_source(source_code: str, mapping_rules: Optional[Dict[str, str]] # Sort rules by length (longest first) to avoid partial replacements sorted_rules = sorted(mapping_rules.items(), key=lambda x: len(x[0]), reverse=True) return "".join( - _replace_porting_line(line, sorted_rules) - for line in source_code.splitlines(keepends=True) + _replace_porting_line(line, sorted_rules) for line in source_code.splitlines(keepends=True) ) -def include_paths(cuda: Optional[bool] = None, device_type: Optional[str] = None) -> List[str]: - """ - Get include paths for compiling extensions. - - Supports both PyTorch < 2.6 (cuda=True) and PyTorch 2.6+ (device_type="cuda") - signatures for compatibility. - - Args: - cuda: (PyTorch < 2.6) Whether to include CUDA/MUSA paths. Deprecated in 2.6+. - device_type: (PyTorch 2.6+) Device type string, e.g. "cuda", "cpu", "musa". - - Returns: - List of include paths - """ - # Handle both old (cuda=bool) and new (device_type=str) signatures +def _path_device_type(device_type, cuda: Optional[bool]) -> str: + """Normalize both positional device names and legacy CUDA booleans.""" + if isinstance(device_type, bool): + return "cuda" if device_type else "cpu" if device_type is not None: - # PyTorch 2.6+ style: device_type="cuda" or "cpu" - # Translate "cuda" to MUSA include paths on MUSA platform - include_device = device_type.lower() in ("cuda", "musa") - elif cuda is not None: - include_device = cuda - else: - # Default: include device paths - include_device = True + return device_type.lower() + return "cuda" if cuda is None or cuda else "cpu" + + +def _path_query_options(device_type, torch_include_dirs, cuda): + # The former torchada signature also allowed (cuda_bool, device_name). + if isinstance(torch_include_dirs, str) and ( + device_type is None or isinstance(device_type, bool) + ): + return torch_include_dirs.lower(), None + return _path_device_type(device_type, cuda), torch_include_dirs + + +def _forward_path_query(path_fn, device_type, torch_include_dirs, cross_target_platform=None): + """Forward supported options to the installed, non-MUSA Torch version.""" + import inspect + + parameters = inspect.signature(path_fn).parameters + device_key = "device_type" if "device_type" in parameters else "cuda" + device_value = device_type if device_key == "device_type" else device_type in ("cuda", "musa") + kwargs = {device_key: device_value} + if torch_include_dirs is not None and "torch_include_dirs" in parameters: + kwargs["torch_include_dirs"] = torch_include_dirs + if cross_target_platform is not None: + if "cross_target_platform" not in parameters: + raise NotImplementedError("Installed Torch does not support cross-target library paths") + kwargs["cross_target_platform"] = cross_target_platform + paths = path_fn(**kwargs) + if torch_include_dirs is False and "torch_include_dirs" not in parameters: + base = set(path_fn(**{device_key: "cpu" if device_key == "device_type" else False})) + paths = [path for path in paths if path not in base] + return paths + + +def include_paths( + device_type=None, + torch_include_dirs: Optional[bool] = None, + *, + cuda: Optional[bool] = None, +) -> List[str]: + """Get extension headers across old CUDA and modern device-type APIs. + + Supports both legacy `cuda=bool` / positional booleans and PyTorch 2.11's + `include_paths(device_type, torch_include_dirs)`. An explicit device name + takes precedence over the legacy CUDA keyword. + """ + normalized_device, torch_include_dirs = _path_query_options( + device_type, torch_include_dirs, cuda + ) + include_device = normalized_device in ("cuda", "musa") + if detect_platform() != Platform.MUSA: + from torch.utils.cpp_extension import include_paths as native_include_paths - platform = detect_platform() + return _forward_path_query(native_include_paths, normalized_device, torch_include_dirs) - if platform == Platform.MUSA: - paths: List[str] = [] - try: - import torch_musa.utils.musa_extension as musa_ext - - if hasattr(musa_ext, "include_paths"): - # musa_ext uses musa=bool parameter, not cuda= or device_type= - paths = list(musa_ext.include_paths(musa=include_device)) - except ImportError: - pass + paths: List[str] = [] + try: + import torch_musa.utils.musa_extension as musa_ext - if not paths: - # Fallback: construct paths manually - musa_home = _get_cuda_home() - if musa_home: - paths.append(os.path.join(musa_home, "include")) - - # Auto-append torchada's libtorch-stable ABI compat headers so - # libtorch-stable kernels (vLLM, SGLang, ...) resolve - # on MUSA. Appended LAST so a future - # torch_musa shipping the real header wins. - if include_device: - paths.append(stable_compat_include_dir()) - return paths - - else: - # Check which signature the torch version supports - import inspect - - from torch.utils.cpp_extension import include_paths as torch_include_paths - - sig = inspect.signature(torch_include_paths) - if "device_type" in sig.parameters: - # PyTorch 2.6+ - if device_type is not None: - return torch_include_paths(device_type=device_type) - else: - return torch_include_paths(device_type="cuda" if include_device else "cpu") - else: - # PyTorch < 2.6 - return torch_include_paths(cuda=include_device) + if hasattr(musa_ext, "include_paths"): + paths = list(musa_ext.include_paths(musa=include_device)) + if torch_include_dirs is False: + base = set(musa_ext.include_paths(musa=False)) + paths = [path for path in paths if path not in base] + except ImportError: + pass + if include_device and not paths: + musa_home = _get_cuda_home() + if musa_home: + paths.append(os.path.join(musa_home, "include")) + if include_device: + paths.append(stable_compat_include_dir()) + return paths def stable_compat_include_dir() -> str: @@ -933,70 +936,51 @@ def stable_compat_box_header() -> str: return os.path.join(stable_compat_include_dir(), "torchada_stable_box.h") -def library_paths(cuda: Optional[bool] = None, device_type: Optional[str] = None) -> List[str]: - """ - Get library paths for compiling extensions. - - Supports both PyTorch < 2.6 (cuda=True) and PyTorch 2.6+ (device_type="cuda") - signatures for compatibility. - - Args: - cuda: (PyTorch < 2.6) Whether to include CUDA/MUSA library paths. Deprecated in 2.6+. - device_type: (PyTorch 2.6+) Device type string, e.g. "cuda", "cpu", "musa". +def library_paths( + device_type=None, + torch_include_dirs: Optional[bool] = None, + cross_target_platform=None, + *, + cuda: Optional[bool] = None, +) -> List[str]: + """Get extension libraries, including the PyTorch 2.11 path options. - Returns: - List of library paths + None preserves torchada's legacy CPU result. Explicit torch_include_dirs=True + includes Torch libraries on CPU, as requested by modern Inductor. """ - # Handle both old (cuda=bool) and new (device_type=str) signatures - if device_type is not None: - # PyTorch 2.6+ style: device_type="cuda" or "cpu" - # Translate "cuda" to MUSA library paths on MUSA platform - include_device = device_type.lower() in ("cuda", "musa") - elif cuda is not None: - include_device = cuda - else: - # Default: include device paths - include_device = True - - platform = detect_platform() - - if platform == Platform.MUSA: - if not include_device: - return [] - - try: - import torch_musa.utils.musa_extension as musa_ext - - if hasattr(musa_ext, "library_paths"): - # musa_ext uses musa=bool parameter, not cuda= or device_type= - return musa_ext.library_paths(musa=include_device) - except ImportError: - pass - - # Fallback: construct paths manually - paths = [] - musa_home = _get_cuda_home() - if musa_home: - paths.append(os.path.join(musa_home, "lib")) - paths.append(os.path.join(musa_home, "lib64")) - return [p for p in paths if os.path.exists(p)] - - else: - # Check which signature the torch version supports - import inspect + normalized_device, torch_include_dirs = _path_query_options( + device_type, torch_include_dirs, cuda + ) + include_device = normalized_device in ("cuda", "musa") + if not include_device and torch_include_dirs is None: + return [] + if detect_platform() != Platform.MUSA: + from torch.utils.cpp_extension import library_paths as native_library_paths + + return _forward_path_query( + native_library_paths, normalized_device, torch_include_dirs, cross_target_platform + ) - from torch.utils.cpp_extension import library_paths as torch_library_paths + if cross_target_platform is not None: + raise NotImplementedError("MUSA cross-target library path discovery is not supported") + if not include_device and torch_include_dirs is not True: + return [] + try: + import torch_musa.utils.musa_extension as musa_ext - sig = inspect.signature(torch_library_paths) - if "device_type" in sig.parameters: - # PyTorch 2.6+ - if device_type is not None: - return torch_library_paths(device_type=device_type) - else: - return torch_library_paths(device_type="cuda" if include_device else "cpu") - else: - # PyTorch < 2.6 - return torch_library_paths(cuda=include_device) + if hasattr(musa_ext, "library_paths"): + paths = list(musa_ext.library_paths(musa=include_device)) + if torch_include_dirs is False: + base = set(musa_ext.library_paths(musa=False)) + paths = [path for path in paths if path not in base] + return paths + except ImportError: + pass + paths = [] + musa_home = _get_cuda_home() + if include_device and musa_home: + paths.extend([os.path.join(musa_home, "lib"), os.path.join(musa_home, "lib64")]) + return [path for path in paths if os.path.exists(path)] def _stable_header_backport_required() -> bool: @@ -1102,8 +1086,7 @@ def _translate_link_args(kwargs: Dict[str, Any]) -> Dict[str, Any]: if kwargs.get("libraries") is not None: new_kwargs["libraries"] = [ - "mtjpeg" if library == "nvjpeg" else library - for library in kwargs["libraries"] + "mtjpeg" if library == "nvjpeg" else library for library in kwargs["libraries"] ] if kwargs.get("define_macros") is not None: diff --git a/tests/test_cpp_extension_path_signatures.py b/tests/test_cpp_extension_path_signatures.py new file mode 100644 index 0000000..0585df1 --- /dev/null +++ b/tests/test_cpp_extension_path_signatures.py @@ -0,0 +1,115 @@ +"""CPU-only contracts for legacy and modern extension path discovery.""" + +import sys +import types + +import pytest + +from torchada import Platform +from torchada.utils import cpp_extension as ext + + +@pytest.fixture +def musa_paths(monkeypatch): + root = types.ModuleType("torch_musa") + utils = types.ModuleType("torch_musa.utils") + backend = types.ModuleType("torch_musa.utils.musa_extension") + backend.include_paths = lambda musa=False: ["/torch/include"] + ( + ["/sdk/include"] if musa else [] + ) + backend.library_paths = lambda musa=False: ["/torch/lib"] + (["/sdk/lib"] if musa else []) + root.utils = utils + utils.musa_extension = backend + for name, module in ( + ("torch_musa", root), + ("torch_musa.utils", utils), + ("torch_musa.utils.musa_extension", backend), + ): + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr(ext, "detect_platform", lambda: Platform.MUSA) + monkeypatch.setattr(ext, "_get_cuda_home", lambda: "/sdk") + return backend + + +@pytest.mark.parametrize( + "tested_device", + ["cpu", "cuda", "musa", False, True], + ids=["host-name", "mapped-gpu-name", "native-gpu-name", "old-host-flag", "old-gpu-flag"], +) +def test_positional_device_and_torch_flag(musa_paths, tested_device): + gpu = tested_device in ("cuda", "musa", True) + includes = ext.include_paths(tested_device, False) + libraries = ext.library_paths(tested_device, False) + assert "/torch/include" not in includes + assert "/torch/lib" not in libraries + assert ("/sdk/include" in includes) == gpu + assert ("/sdk/lib" in libraries) == gpu + + +@pytest.mark.parametrize("flag", [True, False]) +def test_legacy_cuda_keyword_and_positional_flag(musa_paths, flag): + assert ext.include_paths(flag) == ext.include_paths(cuda=flag) + assert ext.library_paths(flag) == ext.library_paths(cuda=flag) + + +def test_cpu_inductor_211_calls(musa_paths): + assert ext.include_paths("cpu", False) == [] + assert ext.library_paths("cpu", torch_include_dirs=True, cross_target_platform=None) == [ + "/torch/lib" + ] + assert ext.library_paths("cpu", torch_include_dirs=False, cross_target_platform=None) == [] + assert ext.library_paths("cpu", True, None) == ["/torch/lib"] + + +@pytest.mark.parametrize("flag, override", [(False, "musa"), (True, "cpu"), (None, "cpu")]) +def test_legacy_two_positional_arguments(musa_paths, flag, override): + assert ext.include_paths(flag, override) == ext.include_paths(cuda=flag, device_type=override) + assert ext.library_paths(flag, override) == ext.library_paths(cuda=flag, device_type=override) + + +def test_keep_legacy_defaults_and_device_precedence(musa_paths): + assert "/sdk/include" in ext.include_paths() + assert "/sdk/lib" in ext.library_paths() + assert ext.include_paths(cuda=True, device_type="cpu") == ["/torch/include"] + assert ext.library_paths(cuda=True, device_type="cpu") == [] + + +def test_musa_cross_target_is_explicitly_unsupported(musa_paths): + with pytest.raises(NotImplementedError, match="cross-target"): + ext.library_paths("musa", cross_target_platform="windows") + + +@pytest.mark.parametrize("kind", ["include_paths", "library_paths"]) +def test_modern_native_forwarding(monkeypatch, kind): + import torch.utils.cpp_extension as native + + calls = [] + + def paths(device_type="cuda", torch_include_dirs=True, cross_target_platform=None): + calls.append((device_type, torch_include_dirs, cross_target_platform)) + return (["/torch/base"] if torch_include_dirs else []) + ["/native/sdk"] + + monkeypatch.setattr(ext, "detect_platform", lambda: Platform.CUDA) + monkeypatch.setattr(native, kind, paths) + target = ext.include_paths if kind == "include_paths" else ext.library_paths + assert target("cuda", False) == ["/native/sdk"] + assert calls == [("cuda", False, None)] + if kind == "library_paths": + target("cuda", True, cross_target_platform="windows") + assert calls[-1] == ("cuda", True, "windows") + + +@pytest.mark.parametrize("kind", ["include_paths", "library_paths"]) +def test_legacy_native_forwarding(monkeypatch, kind): + import torch.utils.cpp_extension as native + + def paths(cuda=False): + return ["/torch/base"] + (["/native/sdk"] if cuda else []) + + monkeypatch.setattr(ext, "detect_platform", lambda: Platform.CUDA) + monkeypatch.setattr(native, kind, paths) + target = ext.include_paths if kind == "include_paths" else ext.library_paths + assert target(True, False) == ["/native/sdk"] + assert target(False, False) == [] + if kind == "library_paths": + assert target(cuda=False) == []