From 1403e7faf8412d48ab6c0701b213a132cbef1bc1 Mon Sep 17 00:00:00 2001 From: Yuzhi Zhang <87814399+voidLitchi@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:40:48 +0800 Subject: [PATCH] [FEAT] Add torch_npu stream support to the Torch C-DLPack addon ## Motivation PyTorch exposes Ascend tensors through torch_npu's PrivateUse1 backend and exports them as kDLExtDev through DLPack. The existing Torch C-DLPack addon can convert these tensors, but its current_work_stream callback only handles CUDA and ROCm. TVM-FFI therefore receives a null stream for torch_npu and may launch a kernel on the default NPU stream instead of the caller's current torch.npu stream. ## Changes - Detect torch_npu after CUDA and ROCm on Linux, and use a torch_npu-specific addon even when PyTorch already exposes a DLPack C exchange API. - Preserve the upstream PyTorch API preference for CUDA and ROCm. - Add a mutually exclusive --build-with-torch-npu option and reject it on Windows and macOS. - Discover torch_npu headers and libraries from the installed package, define BUILD_WITH_TORCH_NPU, and link libtorch_npu. - Return c10_npu::getCurrentNPUStream(device_id).stream() for kDLExtDev from the addon's current_work_stream callback. - Reuse the upstream ABI-aware addon cache key and Torch-version-based C++ standard selection. ## Testing - Add unit coverage for backend selection, platform fallback, existing API preference, and mutually exclusive build options. - Add a torch_npu runtime test that compares current_work_stream against a non-default torch.npu stream. - Verify the modified files with the CI-pinned Ruff and ty versions, ASF header and file-type checks, version consistency, Python compilation, and standalone selection and argument-parsing checks on Windows. The torch_npu runtime test requires a Linux host with torch_npu and Ascend hardware. --- pyproject.toml | 2 + python/tvm_ffi/_optional_torch_c_dlpack.py | 25 ++- .../utils/_build_optional_torch_c_dlpack.py | 72 ++++++-- .../test_current_work_stream_torch_npu.py | 80 +++++++++ tests/python/test_optional_torch_c_dlpack.py | 170 ++++++++++++++++-- 5 files changed, 319 insertions(+), 30 deletions(-) create mode 100644 tests/python/test_current_work_stream_torch_npu.py diff --git a/pyproject.toml b/pyproject.toml index 7a05abedd..f00081d59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -286,6 +286,8 @@ allowed-unresolved-imports = [ "torch", "torch.*", "torch.utils.*", + "torch_npu", + "torch_npu.*", "my_ffi_extension", "my_ffi_extension.*", "_pytest.*", diff --git a/python/tvm_ffi/_optional_torch_c_dlpack.py b/python/tvm_ffi/_optional_torch_c_dlpack.py index b50cc8bf0..ad1b4f6d8 100644 --- a/python/tvm_ffi/_optional_torch_c_dlpack.py +++ b/python/tvm_ffi/_optional_torch_c_dlpack.py @@ -44,6 +44,9 @@ logger = logging.getLogger(__name__) +IS_WINDOWS = sys.platform == "win32" +IS_DARWIN = sys.platform == "darwin" + def _torch_extension_device(torch_module: Any) -> str: """Return the torch backend name used in the optional extension library name.""" @@ -53,6 +56,13 @@ def _torch_extension_device(torch_module: Any) -> str: if getattr(torch_module.version, "hip", None) is not None: return "rocm" return "cuda" + if ( + not IS_WINDOWS + and not IS_DARWIN + and hasattr(torch_module, "npu") + and torch_module.npu.is_available() + ): + return "torch_npu" return "cpu" @@ -101,12 +111,21 @@ def _check_and_update_dlpack_c_exchange_api(tensor_cls: object) -> bool: return False +def _should_use_existing_torch_dlpack_api(torch_module: Any) -> bool: + """Return whether Torch's existing DLPack API can be used without an addon.""" + prefer_torch_npu_override = _torch_extension_device(torch_module) == "torch_npu" + return ( + _check_and_update_dlpack_c_exchange_api(torch_module.Tensor) + and not prefer_torch_npu_override + ) + + def load_torch_c_dlpack_extension() -> Any: # noqa: PLR0912, PLR0915 try: import torch # noqa: PLC0415 import torch.version # noqa: PLC0415 - if _check_and_update_dlpack_c_exchange_api(torch.Tensor): + if _should_use_existing_torch_dlpack_api(torch): # skip loading the extension if the __dlpack_c_exchange_api__ # attribute is already set so we don't have to do it in # newer version of PyTorch @@ -118,7 +137,7 @@ def load_torch_c_dlpack_extension() -> Any: # noqa: PLR0912, PLR0915 try: import torch_c_dlpack_ext # noqa: PLC0415, F401 - if _check_and_update_dlpack_c_exchange_api(torch.Tensor): + if _should_use_existing_torch_dlpack_api(torch): return None except ImportError: pass @@ -164,6 +183,8 @@ def load_torch_c_dlpack_extension() -> Any: # noqa: PLR0912, PLR0915 args.append("--build-with-cuda") elif device == "rocm": args.append("--build-with-rocm") + elif device == "torch_npu": + args.append("--build-with-torch-npu") # use capture_output to reduce noise when building the torch c dlpack addon result = subprocess.run(args, check=False, capture_output=True) diff --git a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py index 9e4f9a166..8d20e5dd6 100644 --- a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py +++ b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py @@ -50,6 +50,9 @@ #include #include #endif +#ifdef BUILD_WITH_TORCH_NPU +#include +#endif using namespace std; namespace at { @@ -504,7 +507,7 @@ } } - // Get current CUDA/ROCm work stream + // Get current CUDA/ROCm/torch_npu work stream static int CurrentWorkStream(DLDeviceType device_type, int32_t device_id, void** out_stream) { try { #ifdef BUILD_WITH_ROCM @@ -518,6 +521,15 @@ *out_stream = at::cuda::getCurrentCUDAStream(device_id).stream(); return 0; } +#endif +#ifdef BUILD_WITH_TORCH_NPU + // torch_npu exposes Ascend as kDLExtDev. Return its current stream so + // that kernels launched through tvm-ffi stay aligned with the caller's + // torch.npu stream. + if (device_type == kDLExtDev) { + *out_stream = c10_npu::getCurrentNPUStream(device_id).stream(); + return 0; + } #endif // For CPU and other devices, return NULL (no stream concept) *out_stream = nullptr; @@ -744,12 +756,8 @@ def get_torch_include_paths(build_with_cuda: bool) -> Sequence[str]: return torch.utils.cpp_extension.include_paths(cuda=build_with_cuda) # ty: ignore[unknown-argument] -def main() -> None: # noqa: PLR0912, PLR0915 - """Build the torch c dlpack extension.""" - # we need to set the following env to avoid tvm_ffi to build the torch c-dlpack addon during importing - os.environ["TVM_FFI_DISABLE_TORCH_C_DLPACK"] = "1" - from tvm_ffi.utils.lockfile import FileLock # noqa: PLC0415 - +def _parse_args(args: Sequence[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments for the Torch C DLPack addon build.""" parser = argparse.ArgumentParser( description="Build the torch c dlpack extension. After building, a shared library will be placed in the output directory.", ) @@ -766,26 +774,44 @@ def main() -> None: # noqa: PLR0912, PLR0915 default=str(Path(os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")).expanduser()), help="Directory to store the built extension library. If not specified, the default cache directory of tvm-ffi will be used.", ) - parser.add_argument( + device_group = parser.add_mutually_exclusive_group() + device_group.add_argument( "--build-with-cuda", action="store_true", help="Build with CUDA support.", ) - parser.add_argument( + device_group.add_argument( "--build-with-rocm", action="store_true", help="Build with ROCm support.", ) + device_group.add_argument( + "--build-with-torch-npu", + action="store_true", + help="Build with torch_npu support.", + ) parser.add_argument( "--libname", type=str, default="auto", - help="The name of the generated library. It can be a name 'auto' to auto-generate a name following 'libtorch_c_dlpack_addon_torch{version.major}{version.minor}-cpu/cuda.{extension}'.", + help="The name of the generated library. It can be 'auto' to generate " + "'libtorch_c_dlpack_addon_torch{version.major}{version.minor}-" + "cpu/cuda/rocm/torch_npu.{extension}'.", ) - args = parser.parse_args() - if args.build_with_cuda and args.build_with_rocm: - raise ValueError("Cannot enable both CUDA and ROCm at the same time.") + parsed_args = parser.parse_args(args) + if parsed_args.build_with_torch_npu and (IS_WINDOWS or IS_DARWIN): + parser.error("--build-with-torch-npu is not supported on Windows or macOS.") + return parsed_args + + +def main() -> None: # noqa: PLR0912, PLR0915 + """Build the torch c dlpack extension.""" + # we need to set the following env to avoid tvm_ffi to build the torch c-dlpack addon during importing + os.environ["TVM_FFI_DISABLE_TORCH_C_DLPACK"] = "1" + from tvm_ffi.utils.lockfile import FileLock # noqa: PLC0415 + + args = _parse_args() # resolve build directory if args.build_dir is None: @@ -803,6 +829,8 @@ def main() -> None: # noqa: PLR0912, PLR0915 device = "cuda" elif args.build_with_rocm: device = "rocm" + elif args.build_with_torch_npu: + device = "torch_npu" else: device = "cpu" suffix = ".dll" if IS_WINDOWS else ".so" @@ -837,8 +865,22 @@ def main() -> None: # noqa: PLR0912, PLR0915 elif args.build_with_rocm: cflags.extend(torch.utils.cpp_extension.COMMON_HIP_FLAGS) cflags.append("-DBUILD_WITH_ROCM") + elif args.build_with_torch_npu: + cflags.append("-DBUILD_WITH_TORCH_NPU") include_paths.extend(get_torch_include_paths(args.build_with_cuda or args.build_with_rocm)) + # torch_npu ships headers and libraries under its own package directory; + # add both include and lib paths so the linker can find libtorch_npu. + # Note: c10_npu symbols are compiled into libtorch_npu itself; there is + # no separate libc10_npu to link against. + if args.build_with_torch_npu: + import torch_npu # noqa: PLC0415 + + torch_npu_path = Path(torch_npu.__file__).parent + include_paths.append(str(torch_npu_path / "include")) + torch_npu_lib_dir = str(torch_npu_path / "lib") + ldflags.extend(["-L", torch_npu_lib_dir]) + # use CXX11 ABI if torch.compiled_with_cxx11_abi(): cflags.append("-D_GLIBCXX_USE_CXX11_ABI=1") @@ -862,6 +904,10 @@ def main() -> None: # noqa: PLR0912, PLR0915 ldflags.extend(["-lc10", "-ltorch", "-ltorch_cpu", "-ltorch_python"]) if args.build_with_cuda: ldflags.extend(["-ltorch_cuda", "-lc10_cuda"]) + if args.build_with_torch_npu: + # c10_npu symbols are exported from libtorch_npu.so; there is + # no separate libc10_npu to link against. + ldflags.extend(["-ltorch_npu"]) # Add Python library linking if IS_WINDOWS: diff --git a/tests/python/test_current_work_stream_torch_npu.py b/tests/python/test_current_work_stream_torch_npu.py new file mode 100644 index 000000000..f5ae1d33e --- /dev/null +++ b/tests/python/test_current_work_stream_torch_npu.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import ctypes + +import pytest + +try: + import torch + import torch_npu # noqa: F401 + from torch.utils import cpp_extension + from tvm_ffi import _optional_torch_c_dlpack, libinfo +except ImportError: + torch = None # ty: ignore[invalid-assignment] + +_HAS_TORCH_NPU = bool(torch is not None and hasattr(torch, "npu") and torch.npu.is_available()) + + +@pytest.mark.skipif(not _HAS_TORCH_NPU, reason="Requires torch_npu and an Ascend runtime") +def test_current_work_stream_matches_torch_npu_stream() -> None: + assert torch is not None + addon_lib = getattr(_optional_torch_c_dlpack, "_LIB", None) + assert addon_lib is not None, "torch_npu DLPack addon was not loaded" + assert hasattr(torch.Tensor, "__dlpack_c_exchange_api__") + api_attr = torch.Tensor.__dlpack_c_exchange_api__ # ty: ignore[unresolved-attribute] + + pythonapi = ctypes.pythonapi + pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t + pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] + api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, b"dlpack_exchange_api") + assert api_ptr != 0 + + source = r""" + #include + #include + + void assert_current_work_stream(int64_t api_ptr_int, + int32_t device_id, + int64_t expected_stream) { + DLPackExchangeAPI* api = reinterpret_cast(api_ptr_int); + TORCH_CHECK(api != nullptr, "API pointer is NULL"); + TORCH_CHECK(api->current_work_stream != nullptr, "current_work_stream is NULL"); + + void* current_stream = nullptr; + int result = api->current_work_stream(kDLExtDev, device_id, ¤t_stream); + TORCH_CHECK(result == 0, "current_work_stream(kDLExtDev) failed"); + TORCH_CHECK(reinterpret_cast(current_stream) == expected_stream, + "kDLExtDev stream mismatch"); + } + """ + + mod = cpp_extension.load_inline( + name="test_current_work_stream_torch_npu_ext", + cpp_sources=[source], + functions=["assert_current_work_stream"], + with_cuda=False, + extra_include_paths=libinfo.include_paths(), + ) + + device_id = torch.npu.current_device() + stream = torch.npu.Stream(device=device_id) + with torch.npu.stream(stream): + expected_stream = int(stream.npu_stream) + mod.assert_current_work_stream(api_ptr, device_id, expected_stream) diff --git a/tests/python/test_optional_torch_c_dlpack.py b/tests/python/test_optional_torch_c_dlpack.py index 398afdbb3..24bab0f5b 100644 --- a/tests/python/test_optional_torch_c_dlpack.py +++ b/tests/python/test_optional_torch_c_dlpack.py @@ -40,11 +40,118 @@ IS_WINDOWS = sys.platform.startswith("win") +@pytest.mark.parametrize( + ("has_existing_api", "torch_npu_available", "expected"), + [ + pytest.param(False, None, False, id="missing-api"), + pytest.param(True, None, True, id="existing-api"), + pytest.param(True, True, False, id="torch-npu-override"), + ], +) +def test_should_use_existing_torch_dlpack_api( + monkeypatch: pytest.MonkeyPatch, + has_existing_api: bool, + torch_npu_available: bool | None, + expected: bool, +) -> None: + tensor_cls = type("Tensor", (), {}) + if has_existing_api: + setattr(tensor_cls, "__dlpack_c_exchange_api__", object()) + torch_module = SimpleNamespace( + Tensor=tensor_cls, + cuda=SimpleNamespace(is_available=lambda: False), + version=SimpleNamespace(cuda=None, hip=None), + ) + if torch_npu_available is not None: + torch_module.npu = SimpleNamespace(is_available=lambda: torch_npu_available) + monkeypatch.setattr(_optional_torch_c_dlpack, "IS_WINDOWS", False) + monkeypatch.setattr(_optional_torch_c_dlpack, "IS_DARWIN", False) + + assert _optional_torch_c_dlpack._should_use_existing_torch_dlpack_api(torch_module) is expected + + +@pytest.mark.skipif(torch is None, reason="torch is not installed") +@pytest.mark.parametrize( + ("device_options", "expected"), + [ + pytest.param((), (False, False, False), id="cpu"), + pytest.param(("--build-with-cuda",), (True, False, False), id="cuda"), + pytest.param(("--build-with-rocm",), (False, True, False), id="rocm"), + pytest.param(("--build-with-torch-npu",), (False, False, True), id="torch-npu"), + ], +) +def test_build_device_option_selection( + monkeypatch: pytest.MonkeyPatch, + device_options: tuple[str, ...], + expected: tuple[bool, bool, bool], +) -> None: + from tvm_ffi.utils import _build_optional_torch_c_dlpack # noqa: PLC0415 + + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_WINDOWS", False) + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_DARWIN", False) + args = _build_optional_torch_c_dlpack._parse_args(device_options) + + assert ( + args.build_with_cuda, + args.build_with_rocm, + args.build_with_torch_npu, + ) == expected + + +@pytest.mark.skipif(torch is None, reason="torch is not installed") +@pytest.mark.parametrize( + "device_options", + [ + pytest.param(("--build-with-cuda", "--build-with-rocm"), id="cuda-rocm"), + pytest.param(("--build-with-cuda", "--build-with-torch-npu"), id="cuda-torch-npu"), + pytest.param(("--build-with-rocm", "--build-with-torch-npu"), id="rocm-torch-npu"), + ], +) +def test_build_device_options_are_mutually_exclusive( + monkeypatch: pytest.MonkeyPatch, + device_options: tuple[str, ...], +) -> None: + from tvm_ffi.utils import _build_optional_torch_c_dlpack # noqa: PLC0415 + + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_WINDOWS", False) + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_DARWIN", False) + + with pytest.raises(SystemExit) as exc_info: + _build_optional_torch_c_dlpack._parse_args(device_options) + + assert exc_info.value.code == 2 + + +@pytest.mark.skipif(torch is None, reason="torch is not installed") +@pytest.mark.parametrize( + ("is_windows", "is_darwin"), + [ + pytest.param(True, False, id="windows"), + pytest.param(False, True, id="macos"), + ], +) +def test_torch_npu_build_is_rejected_on_unsupported_platforms( + monkeypatch: pytest.MonkeyPatch, + is_windows: bool, + is_darwin: bool, +) -> None: + from tvm_ffi.utils import _build_optional_torch_c_dlpack # noqa: PLC0415 + + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_WINDOWS", is_windows) + monkeypatch.setattr(_build_optional_torch_c_dlpack, "IS_DARWIN", is_darwin) + + with pytest.raises(SystemExit) as exc_info: + _build_optional_torch_c_dlpack._parse_args(["--build-with-torch-npu"]) + + assert exc_info.value.code == 2 + + def _fake_torch_module( *, cuda_available: bool, cuda_version: str | None = None, hip_version: str | None = None, + torch_npu_available: bool | None = None, include_cuda_attr: bool = True, include_hip_attr: bool = True, ) -> Any: @@ -53,10 +160,13 @@ def _fake_torch_module( version.cuda = cuda_version if include_hip_attr: version.hip = hip_version - return SimpleNamespace( + torch_module = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: cuda_available), version=version, ) + if torch_npu_available is not None: + torch_module.npu = SimpleNamespace(is_available=lambda: torch_npu_available) + return torch_module def test_torch_extension_device() -> None: @@ -90,6 +200,37 @@ def test_torch_extension_device() -> None: ) +@pytest.mark.parametrize( + ("cuda_available", "cuda_version", "hip_version", "platform", "expected"), + [ + pytest.param(False, None, None, (False, False), "torch_npu", id="torch-npu"), + pytest.param(True, "12.8", None, (False, False), "cuda", id="cuda-before-torch-npu"), + pytest.param(True, None, "7.2", (False, False), "rocm", id="rocm-before-torch-npu"), + pytest.param(False, None, None, (True, False), "cpu", id="windows-fallback"), + pytest.param(False, None, None, (False, True), "cpu", id="macos-fallback"), + ], +) +def test_torch_extension_device_with_torch_npu( + monkeypatch: pytest.MonkeyPatch, + cuda_available: bool, + cuda_version: str | None, + hip_version: str | None, + platform: tuple[bool, bool], + expected: str, +) -> None: + is_windows, is_darwin = platform + monkeypatch.setattr(_optional_torch_c_dlpack, "IS_WINDOWS", is_windows) + monkeypatch.setattr(_optional_torch_c_dlpack, "IS_DARWIN", is_darwin) + + torch_module = _fake_torch_module( + cuda_available=cuda_available, + cuda_version=cuda_version, + hip_version=hip_version, + torch_npu_available=True, + ) + assert _optional_torch_c_dlpack._torch_extension_device(torch_module) == expected + + def test_existing_torch_dlpack_api_is_preferred_on_rocm(monkeypatch: pytest.MonkeyPatch) -> None: torch_module = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: True), @@ -125,30 +266,29 @@ def _run_build(args: list[str]) -> None: @pytest.mark.skipif(torch is None, reason="torch is not installed") -def test_build_torch_c_dlpack_extension() -> None: +def test_build_torch_c_dlpack_extension(tmp_path: Path) -> None: assert torch is not None build_script = Path(tvm_ffi.__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py" + output_dir = tmp_path / "output-dir" + libname = "libtorch_c_dlpack_addon_test.so" args = [ sys.executable, str(build_script), "--output-dir", - "./output-dir", + str(output_dir), "--libname", - "libtorch_c_dlpack_addon_test.so", + libname, ] - # First use "torch.cuda.is_available()" to check whether GPU environment - # is available. Then determine the GPU type. - if torch.cuda.is_available(): - device = _optional_torch_c_dlpack._torch_extension_device(torch) - if device == "cuda": - args.append("--build-with-cuda") - elif device == "rocm": - args.append("--build-with-rocm") - else: - raise ValueError("Cannot determine whether to build with CUDA or ROCm.") + device = _optional_torch_c_dlpack._torch_extension_device(torch) + if device == "cuda": + args.append("--build-with-cuda") + elif device == "rocm": + args.append("--build-with-rocm") + elif device == "torch_npu": + args.append("--build-with-torch-npu") _run_build(args) - lib_path = str(Path("./output-dir/libtorch_c_dlpack_addon_test.so").resolve()) + lib_path = str((output_dir / libname).resolve()) assert Path(lib_path).exists() lib = ctypes.CDLL(lib_path)