diff --git a/README.md b/README.md index aaf5757..b64bd5c 100644 --- a/README.md +++ b/README.md @@ -298,7 +298,7 @@ See `src/torchada/_mapping.py` for the complete mapping table (380+ mappings). ``` # pyproject.toml or requirements.txt -torchada>=0.1.69 +torchada>=0.1.70 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index 1055358..7cbb646 100644 --- a/README_CN.md +++ b/README_CN.md @@ -293,7 +293,7 @@ if torchada.is_gpu_device(device): # 在 CUDA 和 MUSA 上都能工作 ``` # pyproject.toml 或 requirements.txt -torchada>=0.1.69 +torchada>=0.1.70 ``` ### 步骤 2：条件导入 diff --git a/pyproject.toml b/pyproject.toml index 7602e96..c04dd82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchada" -version = "0.1.69" +version = "0.1.70" 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 36088e6..59cd793 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.69" +__version__ = "0.1.70" from . import cuda, utils diff --git a/src/torchada/_mapping.py b/src/torchada/_mapping.py index 95d8327..f6fe01c 100644 --- a/src/torchada/_mapping.py +++ b/src/torchada/_mapping.py @@ -2,11 +2,14 @@ from ._mappings import MAPPING_RULE as _MAPPING_RULE # noqa: F401 -# Extension file suffix mappings: convert .cu/.cuh to .mu/.muh so torch_musa's -# musa_compile rule (which only adds -x musa for .mu/.muh) treats them correctly. +# Keep source extensions unchanged during porting. Sources are ported in place +# (no _musa mirror), and .cu/.cuh compile as MUSA via the patched +# _is_musa_file (which selects the -x musa rule), so renaming to .mu/.muh is +# unnecessary -- and avoiding it removes any need for include/source-path +# rewriting. EXT_REPLACED_MAPPING = { - 'cuh': 'muh', - 'cu': 'mu', + 'cuh': 'cuh', + 'cu': 'cu', 'cc': 'cc', 'cpp': 'cpp', 'cxx': 'cxx', diff --git a/src/torchada/_mappings/cuda_runtime.py b/src/torchada/_mappings/cuda_runtime.py index 7f5b757..ee176cb 100644 --- a/src/torchada/_mappings/cuda_runtime.py +++ b/src/torchada/_mappings/cuda_runtime.py @@ -1,8 +1,6 @@ """cuda_runtime CUDA->MUSA porting rules.""" MAPPING = { - '.cuh"': '.muh"', - '.cuh>': '.muh>', 'cudaMalloc': 'musaMalloc', 'cudaFree': 'musaFree', 'cudaMemcpy': 'musaMemcpy', diff --git a/src/torchada/_mappings/data_types.py b/src/torchada/_mappings/data_types.py index 42837f2..3a9dc1a 100644 --- a/src/torchada/_mappings/data_types.py +++ b/src/torchada/_mappings/data_types.py @@ -30,4 +30,11 @@ 'CUDA_R_8F_E5M2': 'MUSA_R_8F_E5M2', 'cuda_fp16.h': 'musa_fp16.h', 'cuda_bf16.h': 'musa_bf16.h', + # torch_musa's torch::headeronly snapshot omits these scalar-type headers; + # the c10 variants carry the same definitions. Only the dtype headers move + # (Exception.h / ScalarType.h under torch/headeronly resolve on torch_musa). + '#include ': '#include ', + '#include ': '#include ', + '#include ': '#include ', + '#include ': '#include ', } diff --git a/src/torchada/utils/cpp_extension.py b/src/torchada/utils/cpp_extension.py index eb07254..06aa51a 100644 --- a/src/torchada/utils/cpp_extension.py +++ b/src/torchada/utils/cpp_extension.py @@ -25,9 +25,12 @@ ) """ +import functools +import inspect import logging import os import re +import tempfile from typing import Any, Dict, List, Optional, Tuple from .._mapping import _MAPPING_RULE, EXT_REPLACED_MAPPING @@ -44,6 +47,41 @@ # _ensure_stable_headers_patched. _stable_headers_patched = False +# CUDA/C++ translation units and headers whose CONTENT is CUDA→MUSA ported in +# place. SimplePorting.run() walks every file in a directory; restricting the +# substitution to these extensions keeps in-place porting from rewriting build +# scripts, templates, docs and configs (.py / .jinja / .cmake / .md / .json / +# .gitignore ...) onto themselves. See _patch_simple_porting_modify_file. +_PORTABLE_SOURCE_EXTS = ( + ".cu", + ".cuh", + ".cc", + ".cpp", + ".cxx", + ".c", + ".h", + ".hpp", + ".hh", + ".hxx", + ".inl", + ".inc", + ".ipp", + ".tpp", + ".txx", + ".ixx", + ".cppm", +) + +# MUSA-native sources are already MUSA and must NOT be CUDA→MUSA-substituted: +# the mapping would corrupt valid MUSA constructs (e.g. the asm-volatile +# neutralization rule, which disables un-assemblable CUDA PTX, would disable a +# hand-written MUSA inline-asm block and leave its output uninitialized). They +# compile as-is and are left byte-identical. +_MUSA_NATIVE_EXTS = ( + ".mu", + ".muh", +) + def _get_cuda_home() -> Optional[str]: """ @@ -104,6 +142,105 @@ def _is_musa_file(path: str) -> bool: return ext in [".cu", ".cuh", ".mu", ".muh"] +def _with_explicit_musa_language(flags): + """Return MUSA compiler flags that force filename-independent parsing.""" + result = list(flags or []) + for index, flag in enumerate(result): + if flag == "-x=musa" or flag == "-xmusa": + return result + if flag == "-x" and index + 1 < len(result) and result[index + 1] == "musa": + return result + # Keep this as the final pre-source language option so it overrides any + # earlier ``-x`` supplied by generic CUDA-oriented build configuration. + return [*result, "-x", "musa"] + + +def _patch_musa_ninja_language(musa_ext): + """Force ``mcc`` to parse identity-named ``.cu`` files as MUSA. + + torch_musa chooses its ``musa_compile`` Ninja rule via ``_is_musa_file``, + but that rule historically relied on the ``.mu`` suffix to select the MUSA + language. In-place porting deliberately keeps ``.cu`` names, so inject the + explicit language flag into ``musa_cflags`` (which appears before the input + path) for both setuptools and JIT extension builds. + """ + original = getattr(musa_ext, "_write_ninja_file", None) + if original is None or getattr(original, "_torchada_explicit_musa_language", False): + return + + parameters = list(inspect.signature(original).parameters) + if "musa_cflags" not in parameters: + return + musa_cflags_index = parameters.index("musa_cflags") + + @functools.wraps(original) + def patched_write_ninja_file(*args, **kwargs): + args = list(args) + if "musa_cflags" in kwargs: + kwargs["musa_cflags"] = _with_explicit_musa_language(kwargs["musa_cflags"]) + elif musa_cflags_index < len(args): + args[musa_cflags_index] = _with_explicit_musa_language(args[musa_cflags_index]) + else: + kwargs["musa_cflags"] = _with_explicit_musa_language(None) + return original(*args, **kwargs) + + patched_write_ninja_file._torchada_explicit_musa_language = True + musa_ext._write_ninja_file = patched_write_ninja_file + + +def _path_is_within(path: str, root: str) -> bool: + """Return whether ``path`` resolves to ``root`` or one of its descendants.""" + try: + return os.path.commonpath( + (os.path.realpath(path), os.path.realpath(root)) + ) == os.path.realpath(root) + except ValueError: + return False + + +def _coalesce_port_roots(paths): + """Canonicalize roots and retain only their shallowest non-overlapping set.""" + roots = sorted( + {os.path.realpath(os.path.abspath(path)) for path in paths}, + key=lambda path: (len(path.split(os.sep)), path), + ) + result = [] + for root in roots: + if not any(_path_is_within(root, ancestor) for ancestor in result): + result.append(root) + return result + + +def _validate_portable_symlinks(source_dir: str) -> None: + """Reject portable symlink files before upstream ``realpath`` can escape. + + SimplePorting resolves each filename before writing. Under in-place porting + that would modify a symlink target, potentially outside the project tree. + Failing explicitly preserves both the link and its target. + """ + for root, _dirs, files in os.walk(source_dir): + for name in files: + path = os.path.join(root, name) + if os.path.islink(path) and os.path.splitext(name)[1].lower() in _PORTABLE_SOURCE_EXTS: + raise RuntimeError(f"Refusing to port symlinked CUDA/C++ source in place: {path}") + + +def _create_in_place_porter(musa_sp, source_dir: str, mapping_rule): + """Construct SimplePorting without exposing ``_musa`` to its init. + + Upstream initialization unconditionally removes and recreates the computed + mirror directory. Seed it with a disposable directory, then redirect both + input and output to the real source only after construction has completed. + """ + with tempfile.TemporaryDirectory(prefix="torchada-port-init-") as temp_dir: + seed_dir = os.path.join(temp_dir, "source") + os.makedirs(seed_dir) + porter = musa_sp.SimplePorting(cuda_dir_path=seed_dir, mapping_rule=mapping_rule) + porter.cuda_dir_path = source_dir + porter.musa_dir_path = source_dir + return porter + + def _patch_simple_porting_load_replaced_mapping(musa_sp): """ Patch SimplePorting.load_replaced_mapping to suppress unwanted print output. @@ -154,6 +291,73 @@ def open_with_surrogateescape(file, mode="r", *args, **kwargs): musa_sp.open = open_with_surrogateescape +def _patch_simple_porting_modify_file(musa_sp): + """Make ``SimplePorting.modify_file`` (a) only rewrite compiled C/C++/CUDA + files and (b) read the whole source before writing, so porting a file **in + place** (destination path == source path) is safe. + + Two problems with the stock method under in-place porting: + + 1. ``SimplePorting.run`` walks *every* file in the tree and calls + ``modify_file`` on each, regardless of extension. In the legacy + ``_musa`` mirror mode that only wrote CUDA→MUSA-substituted copies + into the throwaway mirror. In place (dst == src) it would rewrite build + scripts, Jinja templates, docs and configs (``.py`` / ``.jinja`` / + ``.md`` / ``.gitignore`` ...) onto themselves, corrupting committed + tooling (e.g. a codegen ``generate.py`` gets ``cutlass``→``mutlass`` + applied to its imports). So the content substitution is gated to the + CUDA/C++ translation-unit / header extensions in + ``_PORTABLE_SOURCE_EXTS``; any other file is left byte-for-byte + untouched (or copied verbatim if a distinct mirror destination is still + in use). MUSA-native ``.mu``/``.muh`` sources (``_MUSA_NATIVE_EXTS``) are + likewise left untouched: they are already MUSA, so substitution only + corrupts them (e.g. the asm-volatile neutralization rule disables a + hand-written MUSA inline-asm block, leaving its result uninitialized). + + 2. The stock method opens the destination with ``"w"`` (truncating) while + the source handle is open, zeroing the file when src == dst. Reading all + lines up front before opening the destination makes the in-place write + safe. + """ + + def modify_file(self, cuda_filepath, musa_filepath): + ext = os.path.splitext(cuda_filepath)[1].lower() + in_place = os.path.realpath(self.cuda_dir_path) == os.path.realpath(self.musa_dir_path) + if ext in _MUSA_NATIVE_EXTS or ext not in _PORTABLE_SOURCE_EXTS: + # Leave byte-identical: either MUSA-native (.mu/.muh — already MUSA, + # substitution would corrupt it) or not a compiled source at all + # (.py/.jinja/.md/...). In place (dst == src) leave it alone; for a + # distinct destination (legacy mirror mode) copy it verbatim so the + # mirror stays complete. + # SimplePorting.change_filename turns a dotless name such as + # ``Makefile`` into ``.Makefile``. When the roots are in-place, do + # not copy to that synthetic path (or overwrite a real hidden file). + if not in_place and os.path.abspath(cuda_filepath) != os.path.abspath(musa_filepath): + import shutil + + shutil.copyfile(cuda_filepath, musa_filepath) + return + if in_place and not _path_is_within(cuda_filepath, self.cuda_dir_path): + raise RuntimeError( + f"Refusing to port CUDA/C++ source outside the in-place root: {cuda_filepath}" + ) + with open(cuda_filepath, encoding="utf-8", errors="surrogateescape") as f: + lines = f.readlines() + out = [] + for line in lines: + if line.startswith("*") or line.startswith("/") or line == "": + out.append(line) + continue + for k, v in self.mapping_rule: + if "cub/" not in line: + line = line.replace(k, v) + out.append(line) + with open(musa_filepath, "w", encoding="utf-8", errors="surrogateescape") as f_musa: + f_musa.writelines(out) + + musa_sp.SimplePorting.modify_file = modify_file + + # Anchor for the accessor injection: the ``numel()`` definition that stock # torch_musa 2.9 already ships. The backported block is spliced in just before # it. ``mutable_data_ptr`` is the idempotency sentinel — absent in stock @@ -211,8 +415,7 @@ def open_with_surrogateescape(file, mode="r", *args, **kwargs): # keeps already-inline / template / comment lines untouched, which makes the # inline rewrite idempotent. _TENSOR_INL_DEF_RE = re.compile( - r"^(?!\s*(?:inline|template|//|\*))" - r"([A-Za-z_][\w:<>,\s\*&]*?\bTensor::[A-Za-z_]\w*\s*\()" + r"^(?!\s*(?:inline|template|//|\*))" r"([A-Za-z_][\w:<>,\s\*&]*?\bTensor::[A-Za-z_]\w*\s*\()" ) @@ -290,8 +493,14 @@ def _patch_torch_musa_stable_headers() -> None: try: import torch_musa - roots.append(os.path.join(os.path.dirname(torch_musa.__file__), - "share", "generated_cuda_compatible", "include")) + roots.append( + os.path.join( + os.path.dirname(torch_musa.__file__), + "share", + "generated_cuda_compatible", + "include", + ) + ) except ImportError: pass @@ -313,7 +522,9 @@ def _patch_torch_musa_stable_headers() -> None: logger.warning( "torchada: could not backport torch::stable::Tensor " "accessors into %s (numel() anchor not found); " - "libtorch-stable kernels may fail to compile", ts) + "libtorch-stable kernels may fail to compile", + ts, + ) if os.path.exists(ti): with open(ti, encoding="utf-8") as f: contents = f.read() @@ -355,11 +566,16 @@ def _apply_musa_patches(): Apply patches to torch_musa modules for CUDA compatibility. This function patches: - 1. musa_ext._is_musa_file - to recognize .cu/.cuh files - 2. musa_sp.EXT_REPLACED_MAPPING - to convert .cu/.cuh to .mu/.muh - 3. musa_sp._MAPPING_RULE - to apply CUDA->MUSA symbol mapping - - These patches are required to compile .cu files on MUSA platform. + 1. musa_ext._is_musa_file - to recognize .cu/.cuh files as MUSA sources + 2. musa_ext._write_ninja_file - to pass ``-x musa`` before identity-named + .cu inputs instead of relying on the legacy .mu suffix + 3. musa_sp.EXT_REPLACED_MAPPING - an identity map so porting keeps the + original .cu/.cuh names (no .mu/.muh rename), so it can run in place + 4. musa_sp._MAPPING_RULE - to apply CUDA->MUSA symbol mapping + 5. musa_sp.SimplePorting.modify_file - restrict content substitution to + compiled sources/headers and make in-place (dst == src) writes safe + + These patches are required to compile .cu files in place on MUSA platform. """ global _musa_patches_applied @@ -375,8 +591,10 @@ def _apply_musa_patches(): # Patch _is_musa_file to recognize .cu/.cuh files musa_ext._is_musa_file = _is_musa_file + _patch_musa_ninja_language(musa_ext) - # Patch EXT_REPLACED_MAPPING to convert .cu/.cuh to .mu/.muh + # Patch EXT_REPLACED_MAPPING to an identity map: porting keeps the + # original .cu/.cuh names so it runs in place (no .mu/.muh rename) musa_sp.EXT_REPLACED_MAPPING = EXT_REPLACED_MAPPING # Patch _MAPPING_RULE with our comprehensive CUDA->MUSA mappings @@ -390,6 +608,8 @@ def _apply_musa_patches(): # Patch simple_porting.open to tolerate non-UTF-8 source files # This preserves SimplePorting's original logic while allowing undecodable bytes to round-trip _patch_simple_porting_open(musa_sp) + # Patch modify_file to read-all-before-write so in-place porting (dst == src) is safe. + _patch_simple_porting_modify_file(musa_sp) # NOTE: the libtorch-stable header backport is intentionally NOT applied # here. It writes into the torch / torch_musa site-packages headers, so @@ -683,7 +903,8 @@ def _create_musa_extension(name: str, sources: List[str], *args, **kwargs): The patches applied by _apply_musa_patches() make MUSAExtension accept .cu/.cuh files directly by: 1. Patching musa_ext._is_musa_file to recognize .cu/.cuh as valid MUSA files - 2. Patching musa_sp.EXT_REPLACED_MAPPING to convert .cu/.cuh to .mu/.muh + 2. Patching musa_sp.EXT_REPLACED_MAPPING to an identity map so .cu/.cuh keep + their names (no .mu/.muh rename) 3. Patching musa_sp._MAPPING_RULE to convert CUDA symbols to MUSA in source code 4. Translating 'nvcc' compile args key to 'mcc' for MUSA compiler """ @@ -713,12 +934,14 @@ def _get_build_extension_class(): Get the BuildExtension class for the current platform. On MUSA platform, returns a custom class that: - 1. Uses SimplePorting to convert CUDA sources to MUSA in run() (like torch's HIPIFY) + 1. Uses SimplePorting to convert CUDA sources to MUSA in place in run() 2. Registers .cu/.cuh as valid source extensions in build_extensions() 3. Provides extensible mapping rules via get_mapping_rule() method - The porting process is automatic and transparent - developers use csrc/*.cu paths - and the build system handles conversion to csrc_musa/*.cu internally. + The porting is automatic and transparent: developers list csrc/*.cu source + paths and the build ports each project-local include root in place (no + _musa mirror, no .cu->.mu rename), so original #include paths and + source paths resolve as-is and nothing downstream needs rewriting. """ platform = detect_platform() @@ -748,10 +971,12 @@ class _MUSABuildExtension(musa_ext.BuildExtension): """ Custom BuildExtension that handles CUDA->MUSA source porting. - This class works like torch's HIPIFY for ROCm: - - run(): Automatically ports CUDA sources to MUSA using SimplePorting - - build_extensions(): Registers .cu/.cuh as valid extensions - - get_mapping_rule(): Returns mapping rules (override in subclass to extend) + - run(): ports each project-local include root's CUDA sources to + MUSA in place (no _musa mirror, no .cu->.mu rename), so + original includes and source paths resolve as-is without + per-file rewriting and no header is reachable through two trees + - build_extensions(): registers .cu/.cuh as valid extensions + - get_mapping_rule(): returns mapping rules (override to extend) Subclasses can override get_mapping_rule() to add project-specific mappings: @@ -764,7 +989,6 @@ def get_mapping_rule(self): } """ - # Track directories that have been ported (class-level for persistence) _ported_dirs = set() def get_mapping_rule(self): @@ -788,172 +1012,93 @@ def build_extensions(self): super().build_extensions() def _port_directory(self, source_dir, mapping_rule=None): - """ - Port a directory containing CUDA sources to MUSA. - - When both .cu and .mu files exist with the same base name, - the .mu file takes precedence (it's the hand-written MUSA version). - - Args: - source_dir: Path to directory containing CUDA sources - mapping_rule: Optional custom mapping rules (uses get_mapping_rule() if None) - - Returns: - str: Path to the ported directory (source_dir + "_musa") + """Port a directory's CUDA sources to MUSA **in place** (no + ``_musa`` mirror): SimplePorting rewrites each file's + content and, with the identity extension map, keeps its name. + Original ``#include`` paths and source paths therefore stay + valid, so nothing downstream needs rewriting. Idempotent per + process via ``_ported_dirs``. """ if mapping_rule is None: mapping_rule = self.get_mapping_rule() - source_dir = os.path.abspath(source_dir) - musa_dir = source_dir + "_musa" - - if source_dir not in self._ported_dirs: - musa_sp.LOGGER.setLevel(logging.ERROR) - musa_sp.SimplePorting( - cuda_dir_path=source_dir, mapping_rule=mapping_rule - ).run() - self._ported_dirs.add(source_dir) - - # After porting, copy any original .mu/.muh files to ensure - # hand-written MUSA files take precedence over ported .cu files. - # This fixes the case where both foo.cu and foo.mu exist - - # SimplePorting processes both, but order is non-deterministic. - import shutil - - for root, _, files in os.walk(source_dir): - for file in files: - ext = os.path.splitext(file)[1].lower() - if ext in [".mu", ".muh"]: - src_path = os.path.join(root, file) - rel_path = os.path.relpath(root, source_dir) - dst_dir = os.path.join(musa_dir, rel_path) - dst_path = os.path.join(dst_dir, file) - os.makedirs(dst_dir, exist_ok=True) - shutil.copy2(src_path, dst_path) - - return musa_dir - - def _convert_source_path(self, source): - """ - Convert a CUDA source path to its ported MUSA equivalent. - - Args: - source: Original source file path (e.g., "csrc/kernel.cu") - - Returns: - tuple: (converted_path, needs_porting) - - converted_path: Path to ported file (e.g., "csrc_musa/kernel.mu") - - needs_porting: True if the source directory needs porting - """ - source_path = os.path.abspath(source) - source_dir = os.path.dirname(source_path) - source_file = os.path.basename(source_path) - base_name, ext_name = os.path.splitext(source_file) - ext_name_lower = ext_name.lower() - - # Port all source files that may contain CUDA references: - # - .cu/.cuh: CUDA source/header files - # - .cc/.cpp/.cxx: C++ files that may reference CUDA symbols - if ext_name_lower in [".cu", ".cuh", ".cc", ".cpp", ".cxx"]: - # Get the ported extension (kept same with our EXT_REPLACED_MAPPING) - new_ext = EXT_REPLACED_MAPPING.get(ext_name_lower[1:], ext_name_lower[1:]) - musa_dir = source_dir + "_musa" - new_source = os.path.join(musa_dir, base_name + "." + new_ext) - return new_source, True - # Handle .mu/.muh files that don't exist at original path - # Try to find them in the _musa directory (already ported files) - elif ext_name_lower in [".mu", ".muh"]: - if not os.path.exists(source_path): - musa_dir = source_dir + "_musa" - musa_source = os.path.join(musa_dir, source_file) - if os.path.exists(musa_source): - # File exists in _musa dir, use it (no porting needed) - return musa_source, False - # File exists at original path, use it as-is - return source, False - else: - return source, False + source_dir = os.path.realpath(os.path.abspath(source_dir)) + if source_dir in self._ported_dirs: + return source_dir + + _validate_portable_symlinks(source_dir) + musa_sp.LOGGER.setLevel(logging.ERROR) + sp = _create_in_place_porter(musa_sp, source_dir, mapping_rule) + sp.run() + + self._ported_dirs.add(source_dir) + return source_dir + + @staticmethod + def _dir_has_portable_sources(path): + """True if ``path`` recursively holds any file the porter would + rewrite — any extension in ``_PORTABLE_SOURCE_EXTS``. Kept in + sync with the porting allowlist so a directory of only + ``.cc``/``.cpp`` sources (no ``.h``/``.cu``) is not skipped. A + directory of only MUSA-native ``.mu``/``.muh`` sources has + nothing to port and is intentionally not a porting target.""" + try: + for _root, _dirs, files in os.walk(path): + for f in files: + if os.path.splitext(f)[1].lower() in _PORTABLE_SOURCE_EXTS: + return True + except OSError: + pass + return False + + @staticmethod + def _is_system_include_dir(path): + return ( + path.startswith("/usr/") + or path.startswith("/opt/") + or "site-packages" in path + or "dist-packages" in path + ) def run(self): - """ - Run the build process with automatic CUDA->MUSA porting. - - This method: - 1. Identifies CUDA source directories from extension sources - 2. Ports each directory using SimplePorting (like torch's HIPIFY) - 3. Updates source paths to point to ported files - 4. Ports include directories as well - 5. Calls parent run() to perform actual compilation + """Port each project-local include root's CUDA sources to MUSA + **in place** (no ``_musa`` mirror, no ``.cu``->``.mu`` + rename), then compile. + + Because nothing moves or is renamed, every source's original + ``#include`` directives -- relative (``../x``) and root-relative + (``dir/x``) alike -- still resolve against the same include roots, + and the extension's source paths stay valid. So there is no + per-file include rewriting, no source-path remapping, and no + second mirror to cause cross-tree ODR. ``.cu``/``.cuh`` compile as + MUSA via the patched ``_is_musa_file`` and explicit + ``-x musa`` compiler flag. """ mapping_rule = self.get_mapping_rule() - + self._ported_dirs = set() + candidate_dirs = [] for ext in self.extensions: - new_sources = [] - dirs_to_port = set() - - # First pass: identify directories that need porting - for source in ext.sources: - ( - new_source, - needs_porting, - ) = self._convert_source_path(source) - new_sources.append(new_source) - if needs_porting: - source_dir = os.path.dirname(os.path.abspath(source)) - dirs_to_port.add(source_dir) - # Sort directories by depth (deepest first) to ensure proper porting order - dirs_to_port = sorted( - dirs_to_port, key=lambda p: p.count("/"), reverse=True - ) - # Port each unique directory - for cuda_dir in dirs_to_port: - self._port_directory(cuda_dir, mapping_rule) - - # Update extension sources to point to ported files - ext.sources = new_sources - - # Port include directories and update include_dirs - # Only port project-local directories, not system paths - if hasattr(ext, "include_dirs") and ext.include_dirs: - new_include_dirs = [] - for inc_dir in ext.include_dirs: - inc_dir_abs = os.path.abspath(inc_dir) - # Skip system directories - only port project-local dirs - # System dirs typically start with /usr, /opt, or site-packages - is_system_dir = ( - inc_dir_abs.startswith("/usr/") - or inc_dir_abs.startswith("/opt/") - or "site-packages" in inc_dir_abs - or "dist-packages" in inc_dir_abs - ) - if os.path.isdir(inc_dir_abs) and not is_system_dir: - # Check if directory might contain CUDA headers (recursively) - has_cuda_headers = False - try: - for root, dirs, files in os.walk(inc_dir_abs): - for f in files: - if f.endswith( - ( - ".h", - ".hpp", - ".cuh", - ".cu", - ) - ): - has_cuda_headers = True - break - if has_cuda_headers: - break - except OSError: - pass - - if has_cuda_headers: - ported_dir = self._port_directory(inc_dir_abs, mapping_rule) - # Add ported dir first so ported headers take precedence - if os.path.isdir(ported_dir): - new_include_dirs.append(ported_dir) - new_include_dirs.append(inc_dir_abs) - ext.include_dirs = new_include_dirs + # System filtering applies only to include roots. An + # explicit source parent is always project-owned even + # when the checkout lives under /opt, /usr/src, or a + # site-packages tree. + for include_dir in list(getattr(ext, "include_dirs", None) or []): + root = os.path.realpath(os.path.abspath(include_dir)) + if not self._is_system_include_dir(root): + candidate_dirs.append(root) + for src in list(getattr(ext, "sources", None) or []): + candidate_dirs.append(os.path.dirname(os.path.abspath(src))) + + # ext.sources and ext.include_dirs are intentionally left + # unchanged: the tree is ported in place, so the original + # paths and includes resolve as-is. + + # Coalesce all extensions at once so a child include root is + # never ported and then recursively ported again through a + # later source-parent ancestor. + for root in _coalesce_port_roots(candidate_dirs): + if os.path.isdir(root) and self._dir_has_portable_sources(root): + self._port_directory(root, mapping_rule) super().run() diff --git a/tests/test_cpp_extension.py b/tests/test_cpp_extension.py index 358f4b8..f347a2c 100644 --- a/tests/test_cpp_extension.py +++ b/tests/test_cpp_extension.py @@ -7,6 +7,7 @@ """ import os +from types import SimpleNamespace # Import torchada first to apply patches import torchada # noqa: F401 @@ -138,9 +139,10 @@ def test_ext_replaced_mapping(self): if torchada.is_musa_platform(): import torch_musa.utils.simple_porting as musa_sp - # Extensions are converted: .cu -> .mu, .cuh -> .muh for mcc compiler - assert musa_sp.EXT_REPLACED_MAPPING["cu"] == "mu" - assert musa_sp.EXT_REPLACED_MAPPING["cuh"] == "muh" + # In-place porting retains filenames; mcc receives an explicit + # ``-x musa`` flag instead of relying on .mu/.muh suffixes. + assert musa_sp.EXT_REPLACED_MAPPING["cu"] == "cu" + assert musa_sp.EXT_REPLACED_MAPPING["cuh"] == "cuh" def test_mapping_rule_exists(self): """Test _MAPPING_RULE is set.""" @@ -167,3 +169,42 @@ def test_mapping_rule_has_expected_entries(self): assert rules.get("cudaStream_t") == "musaStream_t" assert rules.get("at::cuda") == "at::musa" assert rules.get("c10::cuda") == "c10::musa" + + +class TestMusaNinjaLanguagePatch: + """Test the filename-independent mcc language selection patch.""" + + def test_adds_explicit_language_once(self): + from torchada.utils.cpp_extension import _with_explicit_musa_language + + assert _with_explicit_musa_language(["-O3"]) == ["-O3", "-x", "musa"] + assert _with_explicit_musa_language(["-x", "musa", "-O3"]) == [ + "-x", + "musa", + "-O3", + ] + assert _with_explicit_musa_language(["-x=musa", "-O3"]) == ["-x=musa", "-O3"] + assert _with_explicit_musa_language(["-x", "cuda"]) == [ + "-x", + "cuda", + "-x", + "musa", + ] + + def test_patches_positional_and_keyword_musa_cflags(self): + from torchada.utils.cpp_extension import _patch_musa_ninja_language + + calls = [] + + def write_ninja(path, musa_cflags): + calls.append((path, musa_cflags)) + + module = SimpleNamespace(_write_ninja_file=write_ninja) + _patch_musa_ninja_language(module) + module._write_ninja_file("positional", ["-O2"]) + module._write_ninja_file(path="keyword", musa_cflags=["-g"]) + + assert calls == [ + ("positional", ["-O2", "-x", "musa"]), + ("keyword", ["-g", "-x", "musa"]), + ] diff --git a/tests/test_extension_build.py b/tests/test_extension_build.py index 70d8000..1387fee 100644 --- a/tests/test_extension_build.py +++ b/tests/test_extension_build.py @@ -221,8 +221,8 @@ class TestMixedSourcesBuild: """ Test building extensions with mixed source types (.cu, .cuh, .mu, .muh, .cpp). - This tests the fix for the issue where .mu files required manually specifying - the ported path (e.g., csrc_musa/foo.mu instead of csrc/foo.mu). + In-place porting keeps every source at its original path (no _musa + mirror, no .cu -> .mu rename), so a .mu source is listed as csrc/foo.mu. """ def test_mixed_sources_dir_exists(self): @@ -246,12 +246,12 @@ def test_build_mixed_sources_extension(self): """ Test building an extension with mixed .cu/.cuh/.mu/.muh/.cpp sources. - This is the key e2e test that verifies: - 1. .cu files are ported to .mu in the _musa directory - 2. .cuh files are ported to .muh in the _musa directory - 3. .mu files that don't exist at original path are found in _musa directory - 4. .muh files are handled correctly - 5. .cpp files are ported for CUDA symbol translation + This is the key e2e test that verifies in-place porting: + 1. .cu / .cuh files are ported in place (content substituted, names and + locations unchanged — no .mu/.muh rename) + 2. .mu / .muh files already at their original path resolve directly + 3. .cpp files are ported for CUDA symbol translation + 4. no _musa mirror is created """ if not _is_gpu_available(): pytest.skip("CUDA/MUSA not available") @@ -264,9 +264,8 @@ def test_build_mixed_sources_extension(self): src_dir = os.path.join(tmpdir, "csrc") shutil.copytree(MIXED_SOURCES_DIR, src_dir) - # Create setup.py that uses CUDA-style paths for .mu files - # This tests the fix: users can specify csrc/mul_kernel.mu - # and torchada will find it in csrc_musa/mul_kernel.mu after porting + # setup.py lists sources at their real paths; in-place porting keeps + # them there (no csrc_musa mirror, no .cu -> .mu rename). setup_content = """ import torchada # noqa: F401 - Apply MUSA patches from setuptools import setup @@ -279,8 +278,8 @@ def test_build_mixed_sources_extension(self): name="test_mixed_sources", sources=[ "csrc/bindings.cpp", # C++ file with CUDA symbols - "csrc/add_kernel.cu", # CUDA kernel -> ported to csrc_musa/add_kernel.mu - "csrc/mul_kernel.mu", # MUSA kernel -> found in csrc_musa/mul_kernel.mu + "csrc/add_kernel.cu", # CUDA kernel -> ported in place, name kept + "csrc/mul_kernel.mu", # MUSA-native kernel -> left as-is ], include_dirs=["csrc"], # Include dir for headers ) @@ -311,17 +310,17 @@ def test_build_mixed_sources_extension(self): ext_files = [f for f in os.listdir(tmpdir) if f.endswith(".so") or f.endswith(".pyd")] assert len(ext_files) > 0, "No extension file was built" - # Verify ported directory was created - ported_dir = os.path.join(tmpdir, "csrc_musa") - assert os.path.isdir(ported_dir), "Ported directory csrc_musa was not created" - - # Verify ported files exist + # In-place porting: no _musa mirror, no .cu -> .mu rename. + assert not os.path.isdir( + os.path.join(tmpdir, "csrc_musa") + ), "in-place porting must not create a csrc_musa mirror" + # Sources keep their original names and locations. assert os.path.exists( - os.path.join(ported_dir, "add_kernel.mu") - ), "add_kernel.cu was not ported to add_kernel.mu" + os.path.join(src_dir, "add_kernel.cu") + ), "add_kernel.cu should be ported in place (name/location kept)" assert os.path.exists( - os.path.join(ported_dir, "utils.muh") - ), "utils.cuh was not ported to utils.muh" + os.path.join(src_dir, "mul_kernel.mu") + ), "mul_kernel.mu should resolve at its original path" def test_run_mixed_sources_extension(self): """Test running the mixed sources extension after building.""" @@ -410,14 +409,14 @@ def test_run_mixed_sources_extension(self): ) class TestSameNameFilePrecedence: """ - Test that .mu files take precedence over .cu files when both exist. - - This tests the fix for the issue where SimplePorting's file processing - order is non-deterministic, causing either the ported .cu or original .mu - to end up in the _musa directory depending on which is processed last. - - The fix ensures that original .mu/.muh files are always copied after - porting, so they take precedence over auto-ported .cu/.cuh files. + Test that a .cu and a same-named .mu coexist under in-place porting. + + With in-place porting there is no _musa mirror and no .cu -> .mu + rename, so kernel.cu and kernel.mu stay distinct files at their original + paths. The source listed in the build is the one compiled; a sibling .mu is + left in place. (The old mirror behavior — a .mu silently winning over a + same-named .cu inside the generated mirror — no longer applies; list the + .mu explicitly to build it.) """ def test_same_name_dir_exists(self): @@ -431,27 +430,24 @@ def test_both_cu_and_mu_exist(self): assert os.path.exists(cu_path), f"kernel.cu not found: {cu_path}" assert os.path.exists(mu_path), f"kernel.mu not found: {mu_path}" - def test_mu_file_takes_precedence(self): + def test_cu_and_mu_coexist_in_place(self): """ - Test that .mu file takes precedence when both .cu and .mu exist. - - This is the key test: after porting, the csrc_musa/kernel.mu file - should contain the content from the original kernel.mu (magic number 123), - not the ported kernel.cu (magic number 42). + Build listing kernel.cu (kernel.mu also present): in-place porting keeps + both distinct files at their original paths, creates no csrc_musa mirror, + and builds the listed .cu (ported in place). The sibling .mu is untouched. """ if not _is_gpu_available(): pytest.skip("CUDA/MUSA not available") if not torchada.is_musa_platform(): - pytest.skip("Same name precedence test only applicable on MUSA platform") + pytest.skip("In-place coexistence test only applicable on MUSA platform") with tempfile.TemporaryDirectory() as tmpdir: # Create csrc directory and copy source files csrc_dir = os.path.join(tmpdir, "csrc") shutil.copytree(SAME_NAME_DIR, csrc_dir) - # Create setup.py that only specifies the .cu file - # (simulating user who wants to build CUDA code) + # setup.py lists the .cu (a same-named .mu also exists in csrc/) setup_content = f""" import torchada # noqa: F401 - Apply MUSA patches from setuptools import setup @@ -464,7 +460,7 @@ def test_mu_file_takes_precedence(self): name="test_same_name", sources=[ "csrc/bindings.cpp", - "csrc/kernel.cu", # .mu file exists too - should take precedence + "csrc/kernel.cu", # kernel.mu also present; both coexist in place ], ) ], @@ -485,17 +481,21 @@ def test_mu_file_takes_precedence(self): ) assert result.returncode == 0, f"Build failed: {result.stderr}" - # Check that csrc_musa/kernel.mu contains the MUSA version (magic 123) - ported_mu_path = os.path.join(tmpdir, "csrc_musa", "kernel.mu") - assert os.path.exists(ported_mu_path), f"Ported file not found: {ported_mu_path}" - - with open(ported_mu_path, "r") as f: - content = f.read() - # The MUSA version returns 123, CUDA version returns 42 - assert "return 123" in content, ( - f".mu file should take precedence but found ported .cu content. " - f"Expected 'return 123' but got:\n{content}" - ) - assert "return 42" not in content, ( - f"Found ported .cu content instead of original .mu. " f"Content:\n{content}" - ) + # In-place: no mirror, no rename, both files still present in place. + assert not os.path.exists( + os.path.join(tmpdir, "csrc_musa") + ), "in-place porting must not create a csrc_musa mirror" + assert os.path.exists( + os.path.join(csrc_dir, "kernel.cu") + ), "kernel.cu must remain in place" + assert os.path.exists( + os.path.join(csrc_dir, "kernel.mu") + ), "sibling kernel.mu must be left in place" + + # The listed .cu was ported in place and is the built source. + with open(os.path.join(csrc_dir, "kernel.cu"), "r") as f: + cu_content = f.read() + assert "return 42" in cu_content, ( + "the listed kernel.cu (magic 42) should be the built source; " + f"got:\n{cu_content}" + ) diff --git a/tests/test_inplace_porting.py b/tests/test_inplace_porting.py new file mode 100644 index 0000000..5c563fe --- /dev/null +++ b/tests/test_inplace_porting.py @@ -0,0 +1,138 @@ +"""Regression tests for safe, deterministic in-place CUDA source porting.""" + +import os + +import pytest +from setuptools import Distribution, Extension + +import torchada + +pytestmark = pytest.mark.skipif( + not torchada.is_musa_platform(), + reason="In-place SimplePorting integration requires torch_musa", +) + +TOKEN = "TORCHADA_CUDA_TOKEN" +PORTED_TOKEN = "TORCHADA_MUSA_TOKEN" +MAPPING = {TOKEN: PORTED_TOKEN} + + +def _build_extension_command(): + from torchada.utils.cpp_extension import BuildExtension + + return BuildExtension(Distribution()) + + +def test_existing_mirror_sibling_is_preserved(tmp_path): + source_dir = tmp_path / "csrc" + mirror_dir = tmp_path / "csrc_musa" + source_dir.mkdir() + mirror_dir.mkdir() + source = source_dir / "kernel.cu" + sentinel = mirror_dir / "DO_NOT_DELETE.txt" + source.write_text(f"{TOKEN}\n", encoding="utf-8") + sentinel.write_text("user-managed content\n", encoding="utf-8") + + command = _build_extension_command() + command._port_directory(str(source_dir), MAPPING) + + assert sentinel.read_text(encoding="utf-8") == "user-managed content\n" + assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n" + + +def test_portable_symlink_fails_without_modifying_target(tmp_path): + source_dir = tmp_path / "csrc" + source_dir.mkdir() + external = tmp_path / "user_owned.h" + external.write_text(f"{TOKEN}\n", encoding="utf-8") + link = source_dir / "linked.h" + link.symlink_to(external) + + command = _build_extension_command() + with pytest.raises(RuntimeError, match="symlinked CUDA/C\\+\\+ source"): + command._port_directory(str(source_dir), MAPPING) + + assert link.is_symlink() + assert external.read_text(encoding="utf-8") == f"{TOKEN}\n" + + +def test_dotless_files_do_not_create_or_overwrite_hidden_siblings(tmp_path): + source_dir = tmp_path / "csrc" + source_dir.mkdir() + (source_dir / "kernel.cu").write_text(f"{TOKEN}\n", encoding="utf-8") + makefile = source_dir / "Makefile" + hidden_makefile = source_dir / ".Makefile" + makefile.write_text("public makefile\n", encoding="utf-8") + hidden_makefile.write_text("private sentinel\n", encoding="utf-8") + + command = _build_extension_command() + command._port_directory(str(source_dir), MAPPING) + + assert makefile.read_text(encoding="utf-8") == "public makefile\n" + assert hidden_makefile.read_text(encoding="utf-8") == "private sentinel\n" + + +@pytest.mark.parametrize("suffix", [".ipp", ".tpp", ".txx", ".ixx", ".cppm"]) +def test_common_cpp_fragments_are_ported(tmp_path, suffix): + source_dir = tmp_path / "csrc" + source_dir.mkdir() + fragment = source_dir / f"fragment{suffix}" + fragment.write_text(f"{TOKEN}\n", encoding="utf-8") + + command = _build_extension_command() + command._port_directory(str(source_dir), MAPPING) + + assert fragment.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n" + + +def test_overlapping_roots_are_ported_once(tmp_path, monkeypatch): + from torchada.utils.cpp_extension import BuildExtension + + source_dir = tmp_path / "src" + include_dir = source_dir / "include" + include_dir.mkdir(parents=True) + source = source_dir / "kernel.cu" + header = include_dir / "header.h" + source.write_text(f"{TOKEN}\n", encoding="utf-8") + header.write_text(f"{TOKEN}\n", encoding="utf-8") + + class CustomBuildExtension(BuildExtension): + def get_mapping_rule(self): + return {TOKEN: TOKEN + "X"} + + command = CustomBuildExtension(Distribution()) + command.extensions = [ + Extension("test_overlap", sources=[str(source)], include_dirs=[str(include_dir)]) + ] + monkeypatch.setattr(BuildExtension.__mro__[1], "run", lambda self: None) + + command.run() + + assert source.read_text(encoding="utf-8") == f"{TOKEN}X\n" + assert header.read_text(encoding="utf-8") == f"{TOKEN}X\n" + assert command._ported_dirs == {os.path.realpath(source_dir)} + + +def test_source_parent_is_not_filtered_as_system_include(tmp_path, monkeypatch): + from torchada.utils.cpp_extension import BuildExtension + + source_dir = tmp_path / "opt-like-project" + source_dir.mkdir() + source = source_dir / "kernel.cu" + source.write_text(f"{TOKEN}\n", encoding="utf-8") + + class CustomBuildExtension(BuildExtension): + def get_mapping_rule(self): + return MAPPING + + command = CustomBuildExtension(Distribution()) + command.extensions = [ + Extension("test_system_source", sources=[str(source)], include_dirs=[str(source_dir)]) + ] + monkeypatch.setattr(BuildExtension, "_is_system_include_dir", staticmethod(lambda _path: True)) + monkeypatch.setattr(BuildExtension.__mro__[1], "run", lambda self: None) + + command.run() + + assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n" + assert command._ported_dirs == {os.path.realpath(source_dir)} diff --git a/tests/test_mappings.py b/tests/test_mappings.py index b4a4414..76778b1 100644 --- a/tests/test_mappings.py +++ b/tests/test_mappings.py @@ -532,9 +532,10 @@ def test_mapping_count(self): def test_ext_replaced_mapping(self): from torchada._mapping import EXT_REPLACED_MAPPING - # Extensions are converted: .cu -> .mu, .cuh -> .muh for mcc compiler - assert EXT_REPLACED_MAPPING["cu"] == "mu" - assert EXT_REPLACED_MAPPING["cuh"] == "muh" + # Identity map: porting keeps the original .cu/.cuh names (no .mu/.muh + # rename) so it runs in place and source paths need no rewriting. + assert EXT_REPLACED_MAPPING["cu"] == "cu" + assert EXT_REPLACED_MAPPING["cuh"] == "cuh" class TestMappingRobustness: @@ -617,13 +618,15 @@ def test_port_cuda_source_basic(self): assert "at::cuda" not in result def test_port_cuda_source_includes(self): - """Test include statement porting.""" + """CUDA runtime headers are renamed, but a project .cuh include keeps + its extension: porting is in place, so there is no .cuh -> .muh rename.""" from torchada.utils.cpp_extension import _port_cuda_source source = '#include \n#include "my_file.cuh"' result = _port_cuda_source(source) assert "musa_runtime.h" in result - assert 'my_file.muh"' in result + assert 'my_file.cuh"' in result + assert 'my_file.muh"' not in result def test_port_cuda_source_types(self): """Test type replacement.""" @@ -1171,7 +1174,12 @@ def test_run_device_type_extension(self): del sys.modules["test_device_type"] def test_ported_source_contains_privateuse1(self): - """Verify the ported source code contains PrivateUse1, not MUSA.""" + """Verify porting rewrites CUDA device types to PrivateUse1 in place. + + Porting is in place: the source keeps its original .cu name and + location (no _musa mirror, no .cu -> .mu rename), and its CUDA + device types are rewritten to PrivateUse1. + """ if not torchada.is_musa_platform(): pytest.skip("Source porting test only applicable on MUSA platform") @@ -1200,7 +1208,7 @@ def test_ported_source_contains_privateuse1(self): with open(setup_path, "w") as f: f.write(setup_content) - # Build the extension (this triggers source porting) + # Build the extension (this triggers in-place source porting) result = subprocess.run( [sys.executable, "setup.py", "build_ext", "--inplace"], cwd=tmpdir, @@ -1208,31 +1216,22 @@ def test_ported_source_contains_privateuse1(self): text=True, ) - # Find all ported .mu files - # The ported files are in {tmpdir}_musa directory (created by torchada) - ported_files = [] - musa_dir = f"{tmpdir}_musa" - if os.path.exists(musa_dir): - for root, dirs, files in os.walk(musa_dir): - for f in files: - if f.endswith(".mu"): - ported_files.append(os.path.join(root, f)) - - # Also check inside tmpdir in case porting puts files there - for root, dirs, files in os.walk(tmpdir): - for f in files: - if f.endswith(".mu"): - ported_files.append(os.path.join(root, f)) - - assert ( - len(ported_files) > 0 - ), f"No .mu file found after porting. Build output:\n{result.stdout}" - - # Read ported content - ported_content = "" - for pf in ported_files: - with open(pf, "r") as f: - ported_content += f.read() + # In-place porting must not create a _musa mirror ... + assert not os.path.exists( + f"{tmpdir}_musa" + ), "in-place porting must not create a _musa mirror" + # ... nor rename .cu -> .mu: the source keeps its name and location. + assert not any( + f.endswith(".mu") for f in os.listdir(tmpdir) + ), "in-place porting must not rename .cu -> .mu" + + ported_cu = os.path.join(tmpdir, "device_type_test.cu") + assert os.path.exists(ported_cu), ( + "ported source not found in place. Build output:\n" + f"{result.stdout}\n{result.stderr}" + ) + with open(ported_cu, "r") as f: + ported_content = f.read() # Verify the mappings were applied correctly assert ( diff --git a/tests/test_platform.py b/tests/test_platform.py index f46d974..bfb9b7a 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -2,6 +2,9 @@ Tests for platform detection and initialization. """ +import re +from pathlib import Path + class TestPlatformDetection: """Test platform detection functionality.""" @@ -56,8 +59,21 @@ def test_get_version(self): version = torchada.get_version() assert version == torchada.__version__ + assert version == "0.1.70" assert isinstance(version, str) + def test_project_version_matches_runtime_version(self): + """The package metadata and runtime version must stay in sync.""" + import torchada + + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + match = re.search( + r'^version = "([^"]+)"$', pyproject.read_text(encoding="utf-8"), re.MULTILINE + ) + + assert match is not None + assert match.group(1) == torchada.__version__ + def test_get_platform(self): """Test get_platform helper.""" import torchada diff --git a/tests/test_python_compat.py b/tests/test_python_compat.py index cc77874..4f6ad6d 100644 --- a/tests/test_python_compat.py +++ b/tests/test_python_compat.py @@ -39,7 +39,7 @@ def test_all_files_parse(self): except SyntaxError as e: errors.append(f"{filepath}: {e}") - assert not errors, f"Syntax errors found:\n" + "\n".join(errors) + assert not errors, "Syntax errors found:\n" + "\n".join(errors) def test_all_modules_importable(self): """Test that all torchada modules can be imported.""" @@ -103,27 +103,36 @@ def test_no_union_type_operator(self): In Python 3.10+, you can use `int | str` for union types. In Python 3.8/3.9, you must use Union[int, str] from typing. """ - # Pattern for union type operator in type hints - # Matches things like `: int | str` or `-> str | None` - pattern = r"(?::|->)\s*\w+\s*\|\s*\w+" - errors = [] for filepath in get_python_files(): with open(filepath, "r", encoding="utf-8") as f: - content = f.read() - - for i, line in enumerate(content.split("\n"), 1): - # Skip comments - if line.strip().startswith("#"): - continue - # Skip bitwise OR operations (likely not type hints) - if "==" in line or "!=" in line: - continue - - if re.search(pattern, line): - # Exclude dict merge operations and other valid uses - if " | {" not in line and "} | " not in line: - errors.append(f"{filepath}:{i}: {line.strip()}") + source = f.read() + tree = ast.parse(source, filename=str(filepath)) + + annotations = [] + for node in ast.walk(tree): + if isinstance(node, ast.AnnAssign): + annotations.append(node.annotation) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.returns is not None: + annotations.append(node.returns) + annotations.extend( + arg.annotation + for arg in [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs] + if arg.annotation is not None + ) + if node.args.vararg is not None and node.args.vararg.annotation is not None: + annotations.append(node.args.vararg.annotation) + if node.args.kwarg is not None and node.args.kwarg.annotation is not None: + annotations.append(node.args.kwarg.annotation) + + for annotation in annotations: + if any( + isinstance(child, ast.BinOp) and isinstance(child.op, ast.BitOr) + for child in ast.walk(annotation) + ): + text = ast.get_source_segment(source, annotation) or "" + errors.append(f"{filepath}:{annotation.lineno}: {text}") assert not errors, ( "Found Python 3.10+ union type operator. "