Add transparent nvJPEG to MTJPEG porting

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,

src/torchada/_mappings/aten.py

@@ -3,6 +3,7 @@
 MAPPING = {
     '#include <ATen/cuda/Atomic.cuh>': '#include "torch_musa/share/generated_cuda_compatible/include/ATen/musa/Atomic.muh"',
     '#include <ATen/cuda/CUDAContext.h>': '#include "torch_musa/csrc/aten/musa/MUSAContext.h"',
+    '#include <ATen/cuda/CUDAEvent.h>': '#include "torch_musa/csrc/core/MUSAEvent.h"',
     '#include <ATen/cuda/CUDADataType.h>': '#include "torch_musa/csrc/aten/musa/MUSADtype.muh"',
     '#include <ATen/cuda/CUDAGeneratorImpl.h>': '#include "torch_musa/csrc/aten/musa/CUDAGeneratorImpl.h"',
     '#include <ATen/cuda/CUDAUtils.h>': '#include "torch_musa/share/generated_cuda_compatible/include/ATen/musa/MUSAUtils.h"',

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',
+}

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

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 <cuda.h>', '#include <musa.h>'),
+            ('#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<header>[^>"]+)[>"]')
+_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

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 <cuda.h>'] == '#include <musa.h>'
+        assert rules['#include "cuda.h"'] == '#include "musa.h"'
 
 class TestMusaPatches:
     """Test patches applied to torch_musa for extension building."""

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."""

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 <nvjpeg.h>\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 <mtjpeg.h>' 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 <ATen/cuda/CUDAEvent.h>"] == (
+            '#include "torch_musa/csrc/core/MUSAEvent.h"'
+        )
+
+
 class TestCUDARuntimeMappings:
     """Test CUDA runtime to MUSA runtime mappings."""