diff --git a/README.md b/README.md index e59d032..036ba96 100644 --- a/README.md +++ b/README.md @@ -202,9 +202,18 @@ the compatibility include directory on MUSA. Custom stable-ABI builds should add `stable_compat_include_dir()` explicitly, and kernels that use `TORCH_BOX` must force-include the path returned by `stable_compat_box_header()`. Both helpers are available from `torchada.utils.cpp_extension`. The torch_musa 2.9 -header backport runs lazily and best-effort at extension-build time; read-only -headers are left unchanged. A plain `import torchada` does not modify PyTorch or -torch_musa headers. +header backport runs lazily and best-effort at MUSA extension +build time on torch 2.9; read-only headers are left unchanged. Torch 2.11 and +newer provide the stable ABI directly, so the backport is skipped. A plain +`import torchada` does not modify PyTorch or torch_musa headers. + +In-place CUDA-to-MUSA porting protects system include roots by default. Set +`TORCHADA_EXCLUDE_DIRS` to add excluded roots for your environment. Entries may +be directory paths or directory names and are separated by the platform path +separator; commas are also accepted. A name matches a complete path component, +so `TORCHADA_EXCLUDE_DIRS=torch_musa` excludes `/home/torch_musa` without the +full path. Explicit source directories remain eligible for porting even when +they are below an excluded root. ### Custom Ops @@ -383,7 +392,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain. ``` # pyproject.toml or requirements.txt -torchada>=0.1.77 +torchada>=0.1.78 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index 62590d6..09ab843 100644 --- a/README_CN.md +++ b/README_CN.md @@ -192,9 +192,16 @@ torchada 还为近期 vLLM 和 SGLang 在 torch_musa 2.9 上使用的 libtorch `torch.utils.cpp_extension.include_paths()` 会返回该兼容 include 目录。自定义 稳定 ABI 构建应显式加入 `stable_compat_include_dir()`；使用 `TORCH_BOX` 的内核 还必须通过编译器强制 include `stable_compat_box_header()` 返回的头文件。这两个 -辅助函数都位于 `torchada.utils.cpp_extension`。torch_musa 2.9 头文件回补只会在 -扩展构建时延迟、尽力执行；只读头文件保持不变。单纯 `import torchada` 不会修改 -PyTorch 或 torch_musa 头文件。 +辅助函数都位于 `torchada.utils.cpp_extension`。自定义 stable ABI 扩展应通过 +上述方式显式加入兼容头。torch_musa 2.9 头文件回补会在 MUSA 扩展构建时延迟、 +尽力执行；torch 2.11 及更新版本已经原生提供 stable ABI，因此会跳过回补。单纯 +`import torchada` 不会修改 PyTorch 或 torch_musa 头文件。 + +原地 CUDA 到 MUSA 的转换会继续按原有规则保护系统 include 目录。可以通过环境变量 +`TORCHADA_EXCLUDE_DIRS` 额外配置要排除的目录；每一项可以是目录路径，也可以是目录 +名称，使用平台路径分隔符，也支持逗号分隔。名称会按完整路径组件匹配，因此 +`TORCHADA_EXCLUDE_DIRS=torch_musa` 可以直接排除 `/home/torch_musa`，无需填写完整路径。 +即使源目录位于排除目录下，扩展显式提供的源目录仍会执行转换。 ### 自定义算子 @@ -368,7 +375,7 @@ if torchada.is_gpu_device(device): # 在 CUDA 和 MUSA 上都能工作 ``` # pyproject.toml 或 requirements.txt -torchada>=0.1.77 +torchada>=0.1.78 ``` ### 步骤 2：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index 378b43e..2b7361b 100644 --- a/benchmarks/benchmark_history.json +++ b/benchmarks/benchmark_history.json @@ -3,7 +3,7 @@ "description": "Historical benchmark results for torchada performance tracking", "results": [ { - "version": "0.1.77", + "version": "0.1.78", "date": "2026-01-29", "platform": "MUSA", "pytorch_version": "2.7.1", diff --git a/pyproject.toml b/pyproject.toml index 813a009..897743e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchada" -version = "0.1.77" +version = "0.1.78" description = "Adapter package for torch_musa to act exactly like PyTorch CUDA" readme = "README.md" license = {text = "MIT"} diff --git a/src/torchada/__init__.py b/src/torchada/__init__.py index 7b9dcf7..55c45fd 100644 --- a/src/torchada/__init__.py +++ b/src/torchada/__init__.py @@ -24,7 +24,7 @@ from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME """ -__version__ = "0.1.77" +__version__ = "0.1.78" from . import cuda, utils diff --git a/src/torchada/_mappings/libtorch_stable.py b/src/torchada/_mappings/libtorch_stable.py index 690629d..f55b49a 100644 --- a/src/torchada/_mappings/libtorch_stable.py +++ b/src/torchada/_mappings/libtorch_stable.py @@ -2,5 +2,6 @@ MAPPING = { 'aoti_torch_get_current_cuda_stream': 'aoti_torch_get_current_musa_stream', + 'torch_get_current_cuda_blas_handle': 'torch_get_current_musa_blas_handle', 'STABLE_TORCH_LIBRARY_IMPL(_C, CUDA': 'STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1', } diff --git a/src/torchada/csrc/stable_compat/torch/csrc/stable/device.h b/src/torchada/csrc/stable_compat/torch/csrc/stable/device.h index ce0bf7f..72a873f 100644 --- a/src/torchada/csrc/stable_compat/torch/csrc/stable/device.h +++ b/src/torchada/csrc/stable_compat/torch/csrc/stable/device.h @@ -3,6 +3,19 @@ // include it for torch::stable::Device / DeviceType. Provide them on the AOTI // C-shim. The force-included box header pulls this in before so the // patched tensor_struct.h device() accessor can return it. +// +// Newer torch_musa provides a native device.h after this compatibility include +// directory in the search path. Forward to that complete wrapper (including +// device_inl.h), rather than including device_struct.h alone and silently +// dropping inline definitions such as the string constructor. +#if defined(__has_include_next) +#if __has_include_next() +#include_next +#define TORCHADA_HAS_NATIVE_STABLE_DEVICE 1 +#endif +#endif + +#ifndef TORCHADA_HAS_NATIVE_STABLE_DEVICE #include #include namespace torch { @@ -40,3 +53,4 @@ struct Device { }; } // namespace stable } // namespace torch +#endif // TORCHADA_HAS_NATIVE_STABLE_DEVICE diff --git a/src/torchada/csrc/stable_compat/torch/headeronly/core/Dispatch.h b/src/torchada/csrc/stable_compat/torch/headeronly/core/Dispatch.h index e84c1a2..1f5e3b3 100644 --- a/src/torchada/csrc/stable_compat/torch/headeronly/core/Dispatch.h +++ b/src/torchada/csrc/stable_compat/torch/headeronly/core/Dispatch.h @@ -11,8 +11,14 @@ #include #include #include // STD_TORCH_CHECK +#include #include +// torch 2.11+ provides these through ScalarType.h. Keep only the dispatch +// macros below on that version: torch_musa's generated compatible include tree +// may still omit Dispatch.h itself. +#if TORCH_VERSION_MAJOR < 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 11) namespace torch { namespace headeronly { namespace impl { @@ -25,6 +31,7 @@ inline const char* toString(torch::headeronly::ScalarType t) { } } // namespace headeronly } // namespace torch +#endif // torch < 2.11 #define THO_PRIVATE_CASE_TYPE_USING_HINT_TMPL(PRELUDE, enum_type, HINT, ...) \ case enum_type: { \ diff --git a/src/torchada/csrc/stable_compat/torchada_stable_box.h b/src/torchada/csrc/stable_compat/torchada_stable_box.h index 13f1e00..00be265 100644 --- a/src/torchada/csrc/stable_compat/torchada_stable_box.h +++ b/src/torchada/csrc/stable_compat/torchada_stable_box.h @@ -1,4 +1,104 @@ #pragma once +#include +#include +#include +#include +#include + +// This header can be force-included unconditionally by downstream build files. +// PyTorch 2.11+ already provides the boxer, Device, HeaderOnlyArrayRef and the +// stable free functions below, so compiling the backport there would redefine +// native stable-ABI symbols. + +// --- Always-on shims (needed on every torch_musa version) --- +// CUDA_VERSION guards in the kernels: define low so the sm100/Blackwell fast +// paths (also gated on cc_major>=10) compile out on MUSA. torchada also passes +// -DCUDA_VERSION=0 via the build, but define here for direct/JIT use. +#ifndef CUDA_VERSION +#define CUDA_VERSION 0 +#endif + +// Stable-ABI runtime-error check used by some libtorch-stable kernels (e.g. +// minimax_reduce_rms_kernel). The ported kernel calls musa* runtime APIs that +// return a musaError_t (0 == success); wrap them in STD_TORCH_CHECK. +// mcc translates STD_CUDA_CHECK -> STD_MUSA_CHECK at the source level, so both +// names are guarded here. +#ifndef STD_CUDA_CHECK +#define STD_CUDA_CHECK(EXPR) \ + do { \ + auto _musa_err = (EXPR); \ + STD_TORCH_CHECK(_musa_err == 0, "MUSA runtime error: ", \ + static_cast(_musa_err)); \ + } while (0) +#endif +#ifndef STD_MUSA_CHECK +#define STD_MUSA_CHECK(EXPR) STD_CUDA_CHECK(EXPR) +#endif + +// Kernel-launch error check used by some libtorch-stable kernels after a <<<>>> +// launch (e.g. selective_scan). Checks the last runtime error via the same +// STD_CUDA_CHECK path. +#ifndef STD_CUDA_KERNEL_LAUNCH_CHECK +#define STD_CUDA_KERNEL_LAUNCH_CHECK() STD_CUDA_CHECK(musaGetLastError()) +#endif +#ifndef STD_MUSA_KERNEL_LAUNCH_CHECK +#define STD_MUSA_KERNEL_LAUNCH_CHECK() STD_CUDA_KERNEL_LAUNCH_CHECK() +#endif + +// Newer torch_musa exposes the BLAS handle through its stable C shim. Keep a +// CUDA-named forwarding wrapper for downstream sources that bypass source +// porting; normally the mapping rule rewrites the call directly to the MUSA +// name. This avoids depending on torch_musa's private C++ handle-pool ABI. +#if TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 11) +#include +static inline AOTITorchError torch_get_current_cuda_blas_handle(void** ret) { + return torch_get_current_musa_blas_handle(ret); +} +#else +// torch_musa 2.9 has no stable C shim for the BLAS handle. Provide the MUSA +// name produced by source porting through the legacy C++ handle-pool API, then +// keep the CUDA name as a forwarding wrapper for unported sources. +struct _mublasHandle_t; +typedef struct _mublasHandle_t* mublasHandle_t; +namespace at { +namespace musa { +mublasHandle_t getCurrentMUSABlasHandle(); +} +} +static inline AOTITorchError torch_get_current_musa_blas_handle(void** ret) { + auto handle = at::musa::getCurrentMUSABlasHandle(); + *ret = reinterpret_cast(handle); + return handle ? 0 : 1; // fail fast on a null handle instead of crashing in muBLAS +} +static inline AOTITorchError torch_get_current_cuda_blas_handle(void** ret) { + return torch_get_current_musa_blas_handle(ret); +} +#endif + +// torch_musa's AOTI C-shim exposes aoti_torch_get_current_musa_stream, not the +// upstream _cuda_ spelling that libtorch-stable kernels (torch_utils.h) call. +// Forward so those kernels compile against the upstream name. +#ifndef TORCHADA_HAVE_AOTI_CUDA_STREAM +#define TORCHADA_HAVE_AOTI_CUDA_STREAM 1 +static inline AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, + void** ret) { + return aoti_torch_get_current_musa_stream(device_index, ret); +} +#endif + +// Some libtorch-stable kernels reference TORCH_UTILS_CHECK, which torch_musa's +// stable headers do not define; alias it to the stable check macro. +#ifndef TORCH_UTILS_CHECK +#define TORCH_UTILS_CHECK STD_TORCH_CHECK +#endif + +// --- torch < 2.11 backport: boxer, Device, stable free functions --- +// PyTorch 2.11+ already provides the boxer, Device, HeaderOnlyArrayRef and the +// stable free functions below, so compiling the backport there would redefine +// native stable-ABI symbols. +#if TORCH_VERSION_MAJOR < 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 11) // torchada stable-ABI compat: TORCH_BOX + small helpers. // // torch_musa 2.9.0's torch::stable runtime predates the TORCH_BOX boxer family @@ -16,7 +116,6 @@ // kernels call. Define the stable Device here -- BEFORE -- so the // torchada-patched tensor_struct.h accessors (gated on TORCHADA_STABLE_ACCESSORS) // can reference it; empty/from_blob are defined after below. -#include #include #include #include @@ -123,69 +222,6 @@ inline Tensor from_blob(void* data, c10::IntArrayRef sizes, } // namespace stable } // namespace torch -// CUDA_VERSION guards in the kernels: define low so the sm100/Blackwell fast -// paths (also gated on cc_major>=10) compile out on MUSA. torchada also passes -// -DCUDA_VERSION=0 via the build, but define here for direct/JIT use. -#ifndef CUDA_VERSION -#define CUDA_VERSION 0 -#endif - -// Stable-ABI runtime-error check used by some libtorch-stable kernels (e.g. -// minimax_reduce_rms_kernel). The ported kernel calls musa* runtime APIs that -// return a musaError_t (0 == success); wrap them in STD_TORCH_CHECK. -#ifndef STD_CUDA_CHECK -#define STD_CUDA_CHECK(EXPR) \ - do { \ - auto _musa_err = (EXPR); \ - STD_TORCH_CHECK(_musa_err == 0, "MUSA runtime error: ", \ - static_cast(_musa_err)); \ - } while (0) -#endif - -// Kernel-launch error check used by some libtorch-stable kernels after a <<<>>> -// launch (e.g. selective_scan). Checks the last runtime error via the same -// STD_CUDA_CHECK path. -#ifndef STD_CUDA_KERNEL_LAUNCH_CHECK -#define STD_CUDA_KERNEL_LAUNCH_CHECK() STD_CUDA_CHECK(musaGetLastError()) -#endif - -// torch_get_current_cuda_blas_handle has no AOTI stable-ABI shim on torch_musa, -// but torch_musa exposes the stream-bound current handle through its handle pool -// (at::musa::getCurrentMUSABlasHandle). Forward-declare it (resolved at import -// time from the already-loaded libtorch_musa) and return the muBLAS handle so the -// gptq cuBLAS->muBLAS GEMM path works. mublasHandle_t per -// /usr/local/musa/include/internal/mublas_types.h; the typedef is harmless if a -// later (via the cublas_v2.h->mublas.h mapping) repeats it. -struct _mublasHandle_t; -typedef struct _mublasHandle_t* mublasHandle_t; -namespace at { -namespace musa { -mublasHandle_t getCurrentMUSABlasHandle(); -} -} // namespace at -static inline AOTITorchError torch_get_current_cuda_blas_handle(void** ret) { - auto handle = at::musa::getCurrentMUSABlasHandle(); - *ret = reinterpret_cast(handle); - return handle ? 0 : 1; // fail fast on a null handle instead of crashing in muBLAS -} - -// torch_musa's AOTI C-shim exposes aoti_torch_get_current_musa_stream, not the -// upstream _cuda_ spelling that libtorch-stable kernels (torch_utils.h) call. -// Forward so those kernels compile against the upstream name. -#ifndef TORCHADA_HAVE_AOTI_CUDA_STREAM -#define TORCHADA_HAVE_AOTI_CUDA_STREAM 1 -static inline AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, - void** ret) { - return aoti_torch_get_current_musa_stream(device_index, ret); -} -#endif - -// Some libtorch-stable kernels reference TORCH_UTILS_CHECK, which torch_musa's -// stable headers do not define; alias it to the stable check macro. -#ifndef TORCH_UTILS_CHECK -#define TORCH_UTILS_CHECK STD_TORCH_CHECK -#endif - namespace torchada_stable { template @@ -249,3 +285,5 @@ inline void boxed(StableIValue* stack, uint64_t /*nargs*/, uint64_t /*nout*/) { #ifndef TORCH_BOX #define TORCH_BOX(func) (&::torchada_stable::boxed) #endif + +#endif // torch < 2.11 diff --git a/src/torchada/utils/cpp_extension.py b/src/torchada/utils/cpp_extension.py index fcf5b63..5453840 100644 --- a/src/torchada/utils/cpp_extension.py +++ b/src/torchada/utils/cpp_extension.py @@ -26,6 +26,7 @@ """ import functools +import importlib import inspect import logging import os @@ -243,6 +244,61 @@ def _coalesce_port_roots(paths): return result +def _configured_exclusions() -> Tuple[List[str], List[str]]: + """Resolve extra path roots and directory names excluded from porting. + + ``TORCHADA_EXCLUDE_DIRS`` is a path-list (``os.pathsep`` separated; commas + are accepted as well). A bare entry such as ``torch_musa`` matches that + directory name anywhere in an include path. If it is also importable, its + package root is protected as well. The value adds to the existing system + directory rules and is read at build time. + """ + value = os.environ.get("TORCHADA_EXCLUDE_DIRS", "") + separator_pattern = rf"[{re.escape(os.pathsep)},]" + roots = [] + names = [] + for item in (part.strip() for part in re.split(separator_pattern, value)): + if not item: + continue + + expanded = os.path.expanduser(os.path.expandvars(item)) + if os.path.isabs(expanded) or os.sep in expanded: + roots.append(os.path.realpath(os.path.abspath(expanded))) + continue + + names.append(item) + try: + module = importlib.import_module(item) + except (ImportError, AttributeError, OSError): + continue + module_file = getattr(module, "__file__", None) + if module_file: + roots.append(os.path.realpath(os.path.dirname(module_file))) + + return _coalesce_port_roots(roots), names + + +def _configured_exclude_dirs() -> List[str]: + """Return path roots configured through ``TORCHADA_EXCLUDE_DIRS``.""" + return _configured_exclusions()[0] + + +def _is_configured_exclude_dir(path: str) -> bool: + """Return whether ``path`` matches a configured root or directory name.""" + roots, names = _configured_exclusions() + if _path_overlaps_any(path, roots): + return True + path_parts = os.path.realpath(path).split(os.sep) + return any(name in path_parts for name in names) + + +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 + ) + + def _validate_portable_symlinks(source_dir: str) -> None: """Reject portable symlink files before upstream ``realpath`` can escape. @@ -926,6 +982,27 @@ def library_paths(cuda: Optional[bool] = None, device_type: Optional[str] = None return torch_library_paths(cuda=include_device) +def _stable_header_backport_required() -> bool: + """Return whether this torch version needs the stable-ABI header backport. + + torch 2.11 and newer provide the stable ABI directly. The supported torch + 2.9 line still needs torchada's compatibility backport. Parse only the + major/minor prefix so vendor and development suffixes do not affect the + decision. + """ + import torch + + match = re.match(r"^(\d+)\.(\d+)", str(torch.__version__)) + if match is None: + logger.warning( + "Unable to determine whether torch %r needs the stable header backport; " + "applying it for compatibility", + torch.__version__, + ) + return True + return (int(match.group(1)), int(match.group(2))) < (2, 11) + + class CUDAExtension: """ A wrapper that creates either a torch CUDAExtension or MUSA MUSAExtension. @@ -945,11 +1022,9 @@ def __new__(cls, name: str, sources: List[str], *args, **kwargs): **kwargs: Additional keyword arguments """ platform = detect_platform() - if platform == Platform.MUSA: return _create_musa_extension(name, sources, *args, **kwargs) - else: - return _create_cuda_extension(name, sources, *args, **kwargs) + return _create_cuda_extension(name, sources, *args, **kwargs) class CppExtension: @@ -1036,9 +1111,10 @@ def _create_musa_extension(name: str, sources: List[str], *args, **kwargs): """ # Ensure patches are applied _apply_musa_patches() - # Building a MUSA extension: apply the (on-disk) libtorch-stable header - # backport now so csrc/libtorch_stable/*.cu can compile. - _ensure_stable_headers_patched() + # torch 2.9 needs the compatibility backport; torch 2.11+ provides the + # stable ABI directly. + if _stable_header_backport_required(): + _ensure_stable_headers_patched() # Translate CUDA compiler, library, and feature-macro names to MUSA. kwargs = _translate_compile_args(kwargs) @@ -1131,9 +1207,10 @@ def get_mapping_rule(self): return _MAPPING_RULE.copy() def build_extensions(self): - # Building now: apply the (on-disk) libtorch-stable header - # backport before the compiler reads the torch headers. - _ensure_stable_headers_patched() + # torch 2.9 needs the compatibility backport; torch 2.11+ + # provides the stable ABI directly. + if _stable_header_backport_required(): + _ensure_stable_headers_patched() # Register .cu, .cuh as valid source extensions self.compiler.src_extensions += [".cu", ".cuh"] super().build_extensions() @@ -1180,12 +1257,19 @@ def _dir_has_portable_sources(path): @staticmethod def _is_system_include_dir(path): - return ( + is_system_path = ( path.startswith("/usr/") or path.startswith("/opt/") or "site-packages" in path or "dist-packages" in path ) + if is_system_path: + return True + + # Additional dependency roots can be supplied as paths or + # package names, including editable installs such as + # TORCHADA_EXCLUDE_DIRS=torch_musa. + return _is_configured_exclude_dir(path) def run(self): """Port each project-local include root's CUDA sources to MUSA diff --git a/tests/csrc/stable_abi_ops.cu b/tests/csrc/stable_abi_ops.cu index ae204eb..f7595bf 100644 --- a/tests/csrc/stable_abi_ops.cu +++ b/tests/csrc/stable_abi_ops.cu @@ -18,6 +18,7 @@ // ATen headers here (ATen/Dispatch.h also defines ::detail::scalar_type and // would clash with the headeronly Dispatch shim). Launch on the default stream. #include +#include #include #include // torchada THO_DISPATCH shim #include @@ -70,11 +71,23 @@ torch::stable::Tensor passthrough(torch::stable::Tensor& input) { return input; // scalar int return (boxer must from) int64_t numel_of(torch::stable::Tensor& input) { return input.numel(); } +torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& input) { +#if defined(USE_MUSA) + STD_TORCH_CHECK(input.device().is_privateuseone(), "Tensor must be on MUSA device"); +#else + STD_TORCH_CHECK(input.device().is_cuda(), "Tensor must be on CUDA device"); +#endif + return torch::stable::from_blob(input.mutable_data_ptr(), input.sizes(), + input.strides(), input.device(), + input.scalar_type()); +} + STABLE_TORCH_LIBRARY(torchada_stable_test, m) { m.def("negate(Tensor! out, Tensor input) -> ()"); m.def("scale(Tensor! out, Tensor input, float s) -> ()"); m.def("passthrough(Tensor input) -> Tensor"); m.def("numel_of(Tensor input) -> int"); + m.def("weak_ref_tensor(Tensor input) -> Tensor"); } // MUSA tensors are PrivateUse1; torchada's _mapping.py rewrites the upstream @@ -84,4 +97,5 @@ STABLE_TORCH_LIBRARY_IMPL(torchada_stable_test, PrivateUse1, m) { m.impl("scale", TORCH_BOX(&scale)); m.impl("passthrough", TORCH_BOX(&passthrough)); m.impl("numel_of", TORCH_BOX(&numel_of)); + m.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); } diff --git a/tests/test_cpp_extension.py b/tests/test_cpp_extension.py index c60ca0b..1e130fa 100644 --- a/tests/test_cpp_extension.py +++ b/tests/test_cpp_extension.py @@ -145,6 +145,88 @@ def test_porting_translates_torch_cuda_header(self): assert _replace_porting_line(project_include, rules) == project_include +class TestStableHeaderBackportVersionSelection: + """Only torch versions predating 2.11 need the compatibility backport.""" + + def test_torch_29_requires_backport(self, monkeypatch): + import torch + + from torchada.utils.cpp_extension import _stable_header_backport_required + + for version in ("2.9.0", "2.9.0+mtgpu"): + monkeypatch.setattr(torch, "__version__", version) + assert _stable_header_backport_required() is True + + def test_torch_211_has_native_stable_abi(self, monkeypatch): + import torch + + from torchada.utils.cpp_extension import _stable_header_backport_required + + for version in ("2.11.0", "2.11.0.dev20260730+mtgpu", "3.0.0"): + monkeypatch.setattr(torch, "__version__", version) + assert _stable_header_backport_required() is False + + +class TestDependencyIncludeProtection: + """Editable dependency paths are never treated as project porting roots.""" + + def test_dependency_package_and_checkout_paths_overlap(self): + from torchada.utils.cpp_extension import _path_overlaps_any + + dependency_package = "/home/torch_musa/torch_musa" + + assert _path_overlaps_any("/home/torch_musa", [dependency_package]) + assert _path_overlaps_any( + "/home/torch_musa/torch_musa/share/generated_cuda_compatible", + [dependency_package], + ) + assert not _path_overlaps_any("/home/torchvision", [dependency_package]) + + def test_exclude_dirs_can_be_configured_by_environment(self, monkeypatch, tmp_path): + from torchada.utils import cpp_extension + + excluded = tmp_path / "dependency" + excluded.mkdir() + monkeypatch.setenv("TORCHADA_EXCLUDE_DIRS", str(excluded)) + + assert cpp_extension._configured_exclude_dirs() == [str(excluded)] + + def test_exclude_dirs_accepts_package_names(self, monkeypatch, tmp_path): + from torchada.utils import cpp_extension + + package = tmp_path / "torch_musa" + package.mkdir() + module_file = package / "__init__.py" + module_file.write_text("", encoding="utf-8") + monkeypatch.setattr( + cpp_extension.importlib, + "import_module", + lambda name: SimpleNamespace(__file__=str(module_file)), + ) + monkeypatch.setenv("TORCHADA_EXCLUDE_DIRS", "torch_musa") + + assert cpp_extension._configured_exclude_dirs() == [str(package)] + + def test_exclude_dir_name_matches_path_component(self, monkeypatch): + from torchada.utils import cpp_extension + + def unavailable(_name): + raise ImportError + + monkeypatch.setattr(cpp_extension.importlib, "import_module", unavailable) + monkeypatch.setenv("TORCHADA_EXCLUDE_DIRS", "torch_musa") + + assert cpp_extension._is_configured_exclude_dir("/home/torch_musa") + assert cpp_extension._is_configured_exclude_dir("/home/torch_musa/include") + assert not cpp_extension._is_configured_exclude_dir("/home/torch_musa_extra") + + def test_empty_exclude_dirs_adds_nothing(self, monkeypatch): + from torchada.utils import cpp_extension + + monkeypatch.setenv("TORCHADA_EXCLUDE_DIRS", "") + assert cpp_extension._configured_exclude_dirs() == [] + + class TestMusaPatches: """Test patches applied to torch_musa for extension building.""" diff --git a/tests/test_cuda_patching.py b/tests/test_cuda_patching.py index 20753c0..6ad5436 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -3180,6 +3180,8 @@ def test_device_shim_header_present(self): p = os.path.join(stable_compat_include_dir(), "torch", "csrc", "stable", "device.h") assert os.path.isfile(p), f"device.h shim missing: {p}" text = open(p, encoding="utf-8").read() + assert "#include_next " in text + assert "TORCHADA_HAS_NATIVE_STABLE_DEVICE" in text assert "struct Device" in text assert "enum class DeviceType" in text # Predicates vLLM stable kernels call on a Device. @@ -3220,6 +3222,21 @@ def test_box_header_defines_torch_box(self): assert "#define TORCH_BOX(func)" in text assert "namespace torchada_stable" in text + def test_box_header_disables_backport_on_torch_211(self): + """An unconditional downstream -include must not redefine native 2.11 ABI.""" + from torchada.utils.cpp_extension import stable_compat_box_header + + text = open(stable_compat_box_header(), encoding="utf-8").read() + assert "TORCH_VERSION_MINOR < 11" in text + assert "#endif // torch < 2.11" in text + + def test_box_header_uses_stable_musa_blas_shim_on_torch_211(self): + from torchada.utils.cpp_extension import stable_compat_box_header + + text = open(stable_compat_box_header(), encoding="utf-8").read() + assert "#include " in text + assert "return torch_get_current_musa_blas_handle(ret);" in text + def test_box_header_backports_free_functions(self): """The box header supplies the torch::stable free functions torch_musa 2.9 lacks. diff --git a/tests/test_inplace_porting.py b/tests/test_inplace_porting.py index 8c5e590..09274bd 100644 --- a/tests/test_inplace_porting.py +++ b/tests/test_inplace_porting.py @@ -244,3 +244,52 @@ def get_mapping_rule(self): assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n" assert command._ported_dirs == {os.path.realpath(source_dir)} + + +def test_unset_exclude_dirs_keeps_original_include_filter(monkeypatch): + from torchada.utils.cpp_extension import BuildExtension + + monkeypatch.delenv("TORCHADA_EXCLUDE_DIRS", raising=False) + + assert BuildExtension._is_system_include_dir("/usr/include") + assert BuildExtension._is_system_include_dir("/opt/dependency") + assert BuildExtension._is_system_include_dir("/venv/site-packages/dependency") + assert not BuildExtension._is_system_include_dir("/home/project/include") + + +def test_editable_dependency_include_root_is_not_ported(tmp_path, monkeypatch): + """An editable torch_musa checkout stays an include path, not a port root.""" + from torchada.utils.cpp_extension import BuildExtension + + project_dir = tmp_path / "vision" + dependency_dir = tmp_path / "torch_musa" + project_dir.mkdir() + dependency_package = dependency_dir / "torch_musa" + dependency_package.mkdir(parents=True) + + source = project_dir / "kernel.cu" + dependency_header = dependency_package / "dependency.h" + source.write_text(f"{TOKEN}\n", encoding="utf-8") + dependency_header.write_text(f"{TOKEN}\n", encoding="utf-8") + + monkeypatch.setenv("TORCHADA_EXCLUDE_DIRS", "torch_musa") + + class CustomBuildExtension(BuildExtension): + def get_mapping_rule(self): + return MAPPING + + command = CustomBuildExtension(Distribution()) + command.extensions = [ + Extension( + "test_editable_dependency", + sources=[str(source)], + include_dirs=[str(dependency_dir)], + ) + ] + monkeypatch.setattr(BuildExtension.__mro__[1], "run", lambda self: None) + + command.run() + + assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n" + assert dependency_header.read_text(encoding="utf-8") == f"{TOKEN}\n" + assert command._ported_dirs == {os.path.realpath(project_dir)} diff --git a/tests/test_platform.py b/tests/test_platform.py index e843dbc..0676ca2 100644 --- a/tests/test_platform.py +++ b/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.77" + assert version == "0.1.78" assert isinstance(version, str) def test_project_version_matches_runtime_version(self): diff --git a/tests/test_stable_abi.py b/tests/test_stable_abi.py index 6142eb2..ffdf480 100644 --- a/tests/test_stable_abi.py +++ b/tests/test_stable_abi.py @@ -32,7 +32,23 @@ def _gpu_available() -> bool: _SETUP_TEMPLATE = """\ +import os +import pathlib import torchada # noqa: F401 +import torchada.utils.cpp_extension as cpp_extension +os.environ["TORCHADA_EXCLUDE_DIRS"] = {exclude_dir!r} + +if {expect_backport!r}: + _original_backport = cpp_extension._ensure_stable_headers_patched + def _record_backport(): + pathlib.Path("backport-called").write_text("1") + _original_backport() + cpp_extension._ensure_stable_headers_patched = _record_backport +else: + def _unexpected_backport(): + raise AssertionError("torchada patched stable headers on torch >= 2.11") + cpp_extension._ensure_stable_headers_patched = _unexpected_backport + from setuptools import setup from torch.utils.cpp_extension import CUDAExtension, BuildExtension @@ -41,10 +57,10 @@ def _gpu_available() -> bool: ext_modules=[CUDAExtension( name="torchada_stable_test", sources=["stable_abi_ops.cu"], - include_dirs=[{inc!r}], + include_dirs={include_dirs!r}, extra_compile_args={{ - "nvcc": ["-DCUDA_VERSION=0", "-DENABLE_FP8", "-include", {box!r}], - "cxx": ["force_mcc", "-include", {box!r}], + "nvcc": {nvcc_args!r}, + "cxx": {cxx_args!r}, }}, # AOTI stable-ABI shims live in libtorch_cpu.so. libraries=["c10", "torch", "torch_cpu", "torch_python", "musart"], @@ -54,22 +70,29 @@ def _gpu_available() -> bool: """ -@pytest.mark.musa -@pytest.mark.gpu -@pytest.mark.slow -@pytest.mark.skipif(not _gpu_available(), reason="requires a MUSA GPU") -def test_stable_abi_shim_general_signatures(tmp_path): +def _run_stable_abi_shim_test(tmp_path, expect_backport, use_compat_headers): import torch shutil.copy(OPS_CU, tmp_path) (tmp_path / "setup.py").write_text( - _SETUP_TEMPLATE.format(inc=stable_compat_include_dir(), - box=stable_compat_box_header())) + _SETUP_TEMPLATE.format( + expect_backport=expect_backport, + exclude_dir=stable_compat_include_dir(), + include_dirs=[stable_compat_include_dir()] if use_compat_headers else [], + nvcc_args=["-DUSE_MUSA"] + ([ + "-DCUDA_VERSION=0", "-DENABLE_FP8", "-include", + stable_compat_box_header() + ] if use_compat_headers else []), + cxx_args=["force_mcc", "-include", stable_compat_box_header()] + if use_compat_headers else ["force_mcc"], + )) res = subprocess.run( [sys.executable, "setup.py", "build_ext", "--inplace"], cwd=tmp_path, capture_output=True, text=True) assert res.returncode == 0, f"build failed:\n{res.stderr[-4000:]}" + marker = tmp_path / "backport-called" + assert marker.exists() is expect_backport so = glob.glob(str(tmp_path / "torchada_stable_test*.so")) assert so, "extension .so not produced" @@ -99,3 +122,36 @@ def test_stable_abi_shim_general_signatures(tmp_path): assert ns.passthrough(x).data_ptr() == x.data_ptr() # (4) scalar int return -> from assert int(ns.numel_of(x)) == 8 + + x = torch.arange(24, device=dev).reshape(4, 6)[:, ::2] + alias = ns.weak_ref_tensor(x) + assert alias.data_ptr() == x.data_ptr() + assert alias.shape == x.shape + assert alias.stride() == x.stride() + assert alias.dtype == x.dtype + alias.add_(1) + assert torch.equal(alias, x) + + +def _torch_minor() -> int: + import torch + + return int(torch.__version__.split(".")[1]) + + +@pytest.mark.musa +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.skipif(not _gpu_available(), reason="requires a MUSA GPU") +@pytest.mark.skipif(_torch_minor() != 9, reason="requires torch 2.9") +def test_stable_abi_ops_with_torch29_backport(tmp_path): + _run_stable_abi_shim_test(tmp_path, expect_backport=True, use_compat_headers=True) + + +@pytest.mark.musa +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.skipif(not _gpu_available(), reason="requires a MUSA GPU") +@pytest.mark.skipif(_torch_minor() < 11, reason="requires torch 2.11 or newer") +def test_stable_abi_ops_with_native_torch211_abi(tmp_path): + _run_stable_abi_shim_test(tmp_path, expect_backport=False, use_compat_headers=False) diff --git a/tests/test_stable_abi_mappings.py b/tests/test_stable_abi_mappings.py index 2d948e9..a8afbec 100644 --- a/tests/test_stable_abi_mappings.py +++ b/tests/test_stable_abi_mappings.py @@ -20,6 +20,14 @@ def test_aoti_stream_symbol_rule(self): == "aoti_torch_get_current_musa_stream" ) + def test_stable_blas_handle_symbol_rule(self): + from torchada._mapping import _MAPPING_RULE + + assert ( + _MAPPING_RULE["torch_get_current_cuda_blas_handle"] + == "torch_get_current_musa_blas_handle" + ) + def test_stable_impl_dispatch_key_rekey(self): """STABLE_TORCH_LIBRARY_IMPL registers under the literal dispatch-key token; MUSA tensors are PrivateUse1, so the block must be re-keyed.""" @@ -51,6 +59,29 @@ def test_port_rewrites_aoti_stream(self): assert "aoti_torch_get_current_musa_stream" in ported assert "aoti_torch_get_current_cuda_stream" not in ported + def test_port_rewrites_stable_blas_handle(self): + ported = self._port( + "auto e = torch_get_current_cuda_blas_handle(&handle);") + assert "torch_get_current_musa_blas_handle" in ported + assert "torch_get_current_cuda_blas_handle" not in ported + + def test_ported_blas_handle_has_torch_29_fallback(self): + """The symbol emitted by porting must exist on the torch 2.9 path.""" + from torchada.utils.cpp_extension import stable_compat_box_header + + ported = self._port( + "auto e = torch_get_current_cuda_blas_handle(&handle);") + symbol = "torch_get_current_musa_blas_handle" + assert symbol in ported + + with open(stable_compat_box_header(), encoding="utf-8") as f: + header = f.read() + torch_29_branch = header.split( + "// torch_musa 2.9 has no stable C shim", 1 + )[1].split("#endif", 1)[0] + assert f"static inline AOTITorchError {symbol}(void** ret)" in torch_29_branch + assert f"return {symbol}(ret);" in torch_29_branch + def test_port_rekeys_stable_impl_block(self): ported = self._port( "STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl(\"x\", f); }")