diff --git a/src/torchada/_mappings/__init__.py b/src/torchada/_mappings/__init__.py index 9fe7b7a..c38761d 100644 --- a/src/torchada/_mappings/__init__.py +++ b/src/torchada/_mappings/__init__.py @@ -6,6 +6,7 @@ from . import curand as _curand from . import cudnn as _cudnn from . import nccl as _nccl +from . import nvjpeg as _nvjpeg from . import cusparse as _cusparse from . import cusolver as _cusolver from . import cufft as _cufft @@ -28,6 +29,7 @@ **_curand.MAPPING, **_cudnn.MAPPING, **_nccl.MAPPING, + **_nvjpeg.MAPPING, **_cusparse.MAPPING, **_cusolver.MAPPING, **_cufft.MAPPING, diff --git a/src/torchada/_mappings/aten.py b/src/torchada/_mappings/aten.py index 79900f3..ef19fa6 100644 --- a/src/torchada/_mappings/aten.py +++ b/src/torchada/_mappings/aten.py @@ -3,6 +3,7 @@ MAPPING = { '#include ': '#include "torch_musa/share/generated_cuda_compatible/include/ATen/musa/Atomic.muh"', '#include ': '#include "torch_musa/csrc/aten/musa/MUSAContext.h"', + '#include ': '#include "torch_musa/csrc/core/MUSAEvent.h"', '#include ': '#include "torch_musa/csrc/aten/musa/MUSADtype.muh"', '#include ': '#include "torch_musa/csrc/aten/musa/CUDAGeneratorImpl.h"', '#include ': '#include "torch_musa/share/generated_cuda_compatible/include/ATen/musa/MUSAUtils.h"', diff --git a/src/torchada/_mappings/nvjpeg.py b/src/torchada/_mappings/nvjpeg.py new file mode 100644 index 0000000..192b7be --- /dev/null +++ b/src/torchada/_mappings/nvjpeg.py @@ -0,0 +1,9 @@ +"""nvJPEG to MTJPEG porting rules.""" + +# These are the canonical lower- and upper-case spellings used by the C API; +# mixed-case variants are intentionally not inferred. cpp_extension protects +# project-local include paths before applying these prefix rules. +MAPPING = { + 'NVJPEG': 'MTJPEG', + 'nvjpeg': 'mtjpeg', +} diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index 44df92e..1ab38c3 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -23,6 +23,7 @@ import functools import inspect +import logging import os import sys import time @@ -35,6 +36,8 @@ from ._cpp_ops import get_module from ._platform import is_musa_platform +logger = logging.getLogger(__name__) + _patched = False _original_init_process_group = None @@ -320,6 +323,18 @@ def _patch_torch_device(): # Replace torch.device with our wrapper torch.device = DeviceFactoryWrapper + # TorchScript recognizes the original torch.device as an aten builtin. + # Register the wrapper under the same builtin so importing torchada does + # not make otherwise scriptable functions fail during compilation. + try: + from torch.jit._builtins import _register_builtin + + _register_builtin(DeviceFactoryWrapper, "aten::device") + except (AttributeError, ImportError, RuntimeError, TypeError): + logger.debug( + "Unable to register torch.device as a TorchScript builtin", exc_info=True + ) + # Store original torch.Generator for patching _original_torch_generator = None diff --git a/src/torchada/utils/cpp_extension.py b/src/torchada/utils/cpp_extension.py index 06aa51a..852bf54 100644 --- a/src/torchada/utils/cpp_extension.py +++ b/src/torchada/utils/cpp_extension.py @@ -260,13 +260,57 @@ def patched_load_replaced_mapping(self): old_stdout = sys.stdout sys.stdout = io.StringIO() try: - return original_method(self) + result = original_method(self) + # torch_musa's broad ``cuda.h`` rule also rewrites project-local + # headers such as ``decode_jpegs_cuda.h`` even though their file + # names are kept unchanged. Narrow it to actual CUDA header + # includes so relative project includes keep resolving. + self.mapping_rule = _narrow_cuda_header_mapping(self.mapping_rule) + return result finally: sys.stdout = old_stdout musa_sp.SimplePorting.load_replaced_mapping = patched_load_replaced_mapping +def _narrow_cuda_header_mapping(mapping_rule): + """Keep the cuda.h mapping from rewriting project-local header names.""" + narrowed = [(key, value) for key, value in mapping_rule if key != "cuda.h"] + narrowed.extend( + [ + ('#include ', '#include '), + ('#include "cuda.h"', '#include "musa.h"'), + ] + ) + return sorted(narrowed, key=lambda item: len(item[0]), reverse=True) + + +_INCLUDE_DIRECTIVE_RE = re.compile(r'^\s*#\s*include\s*[<"](?P
[^>"]+)[>"]') +_NVJPEG_PREFIX_RULES = frozenset(("nvjpeg", "NVJPEG")) + + +def _replace_porting_line(line, mapping_rule): + """Apply mappings without rewriting nvJPEG text in project header paths.""" + for key, value in mapping_rule: + if key == value: + continue + if key in _NVJPEG_PREFIX_RULES: + include = _INCLUDE_DIRECTIVE_RE.match(line) + if include: + start, end = include.span("header") + 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) + ) + continue + line = line.replace(key, value) + return line + + def _patch_simple_porting_open(musa_sp): """ Patch simple_porting.open to tolerate non-UTF-8 bytes in source files. @@ -348,9 +392,8 @@ def modify_file(self, cuda_filepath, musa_filepath): 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) + if "cub/" not in line: + line = _replace_porting_line(line, self.mapping_rule) out.append(line) with open(musa_filepath, "w", encoding="utf-8", errors="surrogateescape") as f_musa: f_musa.writelines(out) @@ -646,16 +689,12 @@ def _port_cuda_source(source_code: str, mapping_rules: Optional[Dict[str, str]] if mapping_rules is None: mapping_rules = _MAPPING_RULE - result = source_code - # Sort rules by length (longest first) to avoid partial replacements sorted_rules = sorted(mapping_rules.items(), key=lambda x: len(x[0]), reverse=True) - - for cuda_symbol, musa_symbol in sorted_rules: - if cuda_symbol != musa_symbol: - result = result.replace(cuda_symbol, musa_symbol) - - return result + return "".join( + _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]: @@ -897,6 +936,25 @@ def _translate_compile_args(kwargs: Dict[str, Any]) -> Dict[str, Any]: return new_kwargs +def _translate_link_args(kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Translate canonical CUDA library and feature names used by extensions.""" + new_kwargs = kwargs.copy() + + if kwargs.get("libraries") is not None: + new_kwargs["libraries"] = [ + "mtjpeg" if library == "nvjpeg" else library + for library in kwargs["libraries"] + ] + + if kwargs.get("define_macros") is not None: + new_kwargs["define_macros"] = [ + ("MTJPEG_FOUND" if name == "NVJPEG_FOUND" else name, value) + for name, value in kwargs["define_macros"] + ] + + return new_kwargs + + def _create_musa_extension(name: str, sources: List[str], *args, **kwargs): """Create a MUSA extension using torch_musa's MUSAExtension. @@ -914,8 +972,9 @@ def _create_musa_extension(name: str, sources: List[str], *args, **kwargs): # backport now so csrc/libtorch_stable/*.cu can compile. _ensure_stable_headers_patched() - # Translate 'nvcc' to 'mcc' in extra_compile_args + # Translate CUDA compiler, library, and feature-macro names to MUSA. kwargs = _translate_compile_args(kwargs) + kwargs = _translate_link_args(kwargs) try: import torch_musa.utils.musa_extension as musa_ext diff --git a/tests/test_cpp_extension.py b/tests/test_cpp_extension.py index f347a2c..5dbaaa1 100644 --- a/tests/test_cpp_extension.py +++ b/tests/test_cpp_extension.py @@ -99,6 +99,35 @@ def test_create_extension_with_extra_compile_args(self): ) assert ext.extra_compile_args is not None + def test_translate_nvjpeg_link_args(self): + from torchada.utils.cpp_extension import _translate_link_args + + kwargs = { + "libraries": ["jpeg", "nvjpeg"], + "define_macros": [("JPEG_FOUND", 1), ("NVJPEG_FOUND", 1)], + } + + translated = _translate_link_args(kwargs) + + assert translated["libraries"] == ["jpeg", "mtjpeg"] + assert translated["define_macros"] == [ + ("JPEG_FOUND", 1), + ("MTJPEG_FOUND", 1), + ] + assert kwargs["libraries"] == ["jpeg", "nvjpeg"] + + def test_porting_keeps_project_cuda_header_names(self): + from torchada.utils.cpp_extension import _narrow_cuda_header_mapping + + rules = dict( + _narrow_cuda_header_mapping( + [("cudaMalloc", "musaMalloc"), ("cuda.h", "musa.h")] + ) + ) + + assert "cuda.h" not in rules + assert rules['#include '] == '#include ' + assert rules['#include "cuda.h"'] == '#include "musa.h"' class TestMusaPatches: """Test patches applied to torch_musa for extension building.""" diff --git a/tests/test_device_strings.py b/tests/test_device_strings.py index 9cbbb1a..db683c1 100644 --- a/tests/test_device_strings.py +++ b/tests/test_device_strings.py @@ -3,6 +3,11 @@ """ import pytest +import torch + + +def _make_torch_device(device: str): + return torch.device(device) class TestTensorDevicePatching: @@ -391,6 +396,16 @@ def test_torch_device_type_keyword(self): assert device2.type == "cuda" assert device2.index == 0 + def test_torch_device_remains_a_torchscript_builtin(self): + """Test importing torchada keeps torch.device scriptable.""" + import torchada + + if not torchada.is_musa_platform(): + pytest.skip("MUSA platform required with patched torch.device") + + scripted = torch.jit.script(_make_torch_device) + assert scripted("musa") == torch.device("musa") + class TestDeviceContextManager: """Test device context manager (with torch.device(...):) works.""" diff --git a/tests/test_mappings.py b/tests/test_mappings.py index 76778b1..6b31493 100644 --- a/tests/test_mappings.py +++ b/tests/test_mappings.py @@ -167,6 +167,40 @@ def test_cudnn_functions(self): assert _MAPPING_RULE["cudnnSetStream"] == "mudnnSetStream" +class TestNVJPEGMappings: + """Test nvJPEG to MTJPEG mappings.""" + + def test_nvjpeg_prefixes(self): + from torchada._mapping import _MAPPING_RULE + + assert _MAPPING_RULE["nvjpeg"] == "mtjpeg" + assert _MAPPING_RULE["NVJPEG"] == "MTJPEG" + + def test_porting_preserves_project_header_names(self): + from torchada.utils.cpp_extension import _port_cuda_source + + source = ( + '#include "nvjpeg_helpers.h"\n' + '#include "nvjpeg.h"\n' + '# include \n' + 'nvjpegStatus_t status = NVJPEG_STATUS_SUCCESS;\n' + ) + + result = _port_cuda_source(source) + + assert '#include "nvjpeg_helpers.h"' in result + assert '#include "mtjpeg.h"' in result + assert '# include ' in result + assert "mtjpegStatus_t status = MTJPEG_STATUS_SUCCESS" in result + + def test_cuda_event_header(self): + from torchada._mapping import _MAPPING_RULE + + assert _MAPPING_RULE["#include "] == ( + '#include "torch_musa/csrc/core/MUSAEvent.h"' + ) + + class TestCUDARuntimeMappings: """Test CUDA runtime to MUSA runtime mappings."""