[FEAT] Add Ascend stream support to the Torch C-DLPack addon

pyproject.toml

@@ -286,6 +286,8 @@ allowed-unresolved-imports = [
   "torch",
   "torch.*",
   "torch.utils.*",
+  "torch_npu",
+  "torch_npu.*",
   "my_ffi_extension",
   "my_ffi_extension.*",
   "_pytest.*",

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)

python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py

@@ -50,6 +50,9 @@
 #include <c10/hip/HIPStream.h>
 #include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
 #endif
+#ifdef BUILD_WITH_TORCH_NPU
+#include <torch_npu/csrc/core/npu/NPUStream.h>
+#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:

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 <torch/extension.h>
+    #include <dlpack/dlpack.h>
+
+    void assert_current_work_stream(int64_t api_ptr_int,
+                                    int32_t device_id,
+                                    int64_t expected_stream) {
+        DLPackExchangeAPI* api = reinterpret_cast<DLPackExchangeAPI*>(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, &current_stream);
+        TORCH_CHECK(result == 0, "current_work_stream(kDLExtDev) failed");
+        TORCH_CHECK(reinterpret_cast<int64_t>(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)