diff --git a/README.md b/README.md index f3d1df7..62ec77b 100644 --- a/README.md +++ b/README.md @@ -392,7 +392,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain. ``` # pyproject.toml or requirements.txt -torchada>=0.1.79 +torchada>=0.1.80 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index c3844b9..7ff295f 100644 --- a/README_CN.md +++ b/README_CN.md @@ -375,7 +375,7 @@ if torchada.is_gpu_device(device): # 在 CUDA 和 MUSA 上都能工作 ``` # pyproject.toml 或 requirements.txt -torchada>=0.1.79 +torchada>=0.1.80 ``` ### 步骤 2：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index d81bbd9..b980dc1 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.79", + "version": "0.1.80", "date": "2026-01-29", "platform": "MUSA", "pytorch_version": "2.7.1", diff --git a/pyproject.toml b/pyproject.toml index a47cfaf..15668c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchada" -version = "0.1.79" +version = "0.1.80" 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 7135c6a..53e74d6 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.79" +__version__ = "0.1.80" from . import cuda, utils diff --git a/src/torchada/_mappings/libtorch_stable.py b/src/torchada/_mappings/libtorch_stable.py index f55b49a..e6a9ac5 100644 --- a/src/torchada/_mappings/libtorch_stable.py +++ b/src/torchada/_mappings/libtorch_stable.py @@ -3,5 +3,7 @@ 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', + 'torch_set_current_cuda_stream': 'torch_set_current_musa_stream', + 'torch_get_cuda_stream_from_pool': 'torch_get_musa_stream_from_pool', + 'torch_cuda_stream_synchronize': 'torch_musa_stream_synchronize', } diff --git a/src/torchada/utils/cpp_extension.py b/src/torchada/utils/cpp_extension.py index 5453840..1dbee36 100644 --- a/src/torchada/utils/cpp_extension.py +++ b/src/torchada/utils/cpp_extension.py @@ -377,6 +377,20 @@ def _narrow_cuda_header_mapping(mapping_rule): _INCLUDE_DIRECTIVE_RE = re.compile(r'^\s*#\s*include\s*[<"](?P
[^>"]+)[>"]') _NVJPEG_PREFIX_RULES = frozenset(("nvjpeg", "NVJPEG")) +_STABLE_LIBRARY_IMPL_CUDA_RE = re.compile( + r"(\bSTABLE_TORCH_LIBRARY_IMPL\s*\(\s*[A-Za-z_]\w*\s*,\s*)CUDA(\s*,)" +) + + +def _rekey_stable_library_impl(source_code): + """Register stable-ABI CUDA implementations on MUSA's dispatch key. + + The dispatch key is a bare macro argument, so mappings such as + ``torch::kCUDA -> torch::kPrivateUse1`` cannot match it. Handle every + valid library namespace and whitespace layout, including multiline macro + invocations, instead of maintaining project-specific literal mappings. + """ + return _STABLE_LIBRARY_IMPL_CUDA_RE.sub(r"\1PrivateUse1\2", source_code) def _replace_porting_line(line, mapping_rule): @@ -487,7 +501,8 @@ def modify_file(self, cuda_filepath, musa_filepath): 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() + source_code = _rekey_stable_library_impl(f.read()) + lines = source_code.splitlines(keepends=True) def port_line(line): if line.startswith("*") or line.startswith("/") or line == "": @@ -813,6 +828,8 @@ def _port_cuda_source(source_code: str, mapping_rules: Optional[Dict[str, str]] if mapping_rules is None: mapping_rules = _MAPPING_RULE + source_code = _rekey_stable_library_impl(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) return "".join( diff --git a/tests/test_inplace_porting.py b/tests/test_inplace_porting.py index 09274bd..ca3c5d6 100644 --- a/tests/test_inplace_porting.py +++ b/tests/test_inplace_porting.py @@ -23,6 +23,33 @@ def _build_extension_command(): return BuildExtension(Distribution()) +def test_stable_library_impl_cuda_key_is_rekeyed_for_any_namespace(tmp_path): + source_dir = tmp_path / "csrc" + source_dir.mkdir() + source = source_dir / "registration.cpp" + source.write_text( + "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CUDA, m) {}\n" + "STABLE_TORCH_LIBRARY_IMPL(\n" + " third_party_ops,\n" + " CUDA,\n" + " m) {}\n" + "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CPU, m) {}\n", + encoding="utf-8", + ) + + command = _build_extension_command() + command._port_directory(str(source_dir), {}) + + assert source.read_text(encoding="utf-8") == ( + "STABLE_TORCH_LIBRARY_IMPL(custom_ops, PrivateUse1, m) {}\n" + "STABLE_TORCH_LIBRARY_IMPL(\n" + " third_party_ops,\n" + " PrivateUse1,\n" + " m) {}\n" + "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CPU, m) {}\n" + ) + + def test_existing_mirror_sibling_is_preserved(tmp_path): source_dir = tmp_path / "csrc" mirror_dir = tmp_path / "csrc_musa" diff --git a/tests/test_platform.py b/tests/test_platform.py index b0987b1..4fc93c7 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.79" + assert version == "0.1.80" assert isinstance(version, str) def test_project_version_matches_runtime_version(self): diff --git a/tests/test_stable_abi_mappings.py b/tests/test_stable_abi_mappings.py index a8afbec..03bdfd9 100644 --- a/tests/test_stable_abi_mappings.py +++ b/tests/test_stable_abi_mappings.py @@ -28,15 +28,18 @@ def test_stable_blas_handle_symbol_rule(self): == "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.""" + @pytest.mark.parametrize( + ("cuda_symbol", "musa_symbol"), + [ + ("torch_set_current_cuda_stream", "torch_set_current_musa_stream"), + ("torch_get_cuda_stream_from_pool", "torch_get_musa_stream_from_pool"), + ("torch_cuda_stream_synchronize", "torch_musa_stream_synchronize"), + ], + ) + def test_stable_stream_symbol_rules(self, cuda_symbol, musa_symbol): from torchada._mapping import _MAPPING_RULE - assert ( - _MAPPING_RULE["STABLE_TORCH_LIBRARY_IMPL(_C, CUDA"] - == "STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1" - ) + assert _MAPPING_RULE[cuda_symbol] == musa_symbol def test_cuda_header_rules_present(self): from torchada._mapping import _MAPPING_RULE @@ -65,6 +68,17 @@ def test_port_rewrites_stable_blas_handle(self): assert "torch_get_current_musa_blas_handle" in ported assert "torch_get_current_cuda_blas_handle" not in ported + def test_port_rewrites_stable_stream_symbols(self): + ported = self._port( + "torch_set_current_cuda_stream(stream, device);\n" + "torch_get_cuda_stream_from_pool(false, device, &stream);\n" + "torch_cuda_stream_synchronize(stream, device);\n" + ) + assert "torch_set_current_musa_stream" in ported + assert "torch_get_musa_stream_from_pool" in ported + assert "torch_musa_stream_synchronize" in ported + assert "cuda_stream" 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 @@ -82,10 +96,25 @@ def test_ported_blas_handle_has_torch_29_fallback(self): 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): + @pytest.mark.parametrize("namespace", ["_C", "custom_ops", "third_party_ops"]) + def test_port_rekeys_stable_impl_block_for_any_namespace(self, namespace): ported = self._port( - "STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl(\"x\", f); }") - assert "STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1" in ported + f"STABLE_TORCH_LIBRARY_IMPL({namespace}, CUDA, ops) {{ ops.impl(\"x\", f); }}") + assert f"STABLE_TORCH_LIBRARY_IMPL({namespace}, PrivateUse1" in ported + + def test_port_rekeys_multiline_stable_impl_block(self): + ported = self._port( + "STABLE_TORCH_LIBRARY_IMPL(\n" + " third_party_ops,\n" + " CUDA,\n" + " ops) { ops.impl(\"x\", f); }" + ) + assert "third_party_ops,\n PrivateUse1," in ported + + @pytest.mark.parametrize("dispatch_key", ["CPU", "CompositeExplicitAutograd", "PrivateUse1"]) + def test_port_preserves_non_cuda_stable_impl_key(self, dispatch_key): + source = f"STABLE_TORCH_LIBRARY_IMPL(custom_ops, {dispatch_key}, ops) {{}}" + assert self._port(source) == source class TestStableCompatHeadersShipped: