voidLitchi opened a new pull request, #690: URL: https://github.com/apache/tvm-ffi/pull/690
## Motivation TVM-FFI uses the DLPack C exchange API’s `current_work_stream` callback to propagate the framework’s current execution stream into an FFI call. This allows a kernel launched through TVM-FFI to follow the caller’s stream ordering without requiring an explicit synchronization. On Ascend, `torch_npu` registers the NPU backend through PyTorch’s `PrivateUse1` mechanism. These tensors are currently represented as `kDLExtDev` when exchanged through DLPack. The optional Torch C-DLPack addon can already convert such tensors, but its `current_work_stream` implementation previously handled only CUDA and ROCm. An Ascend request therefore returned a null stream and caused the caller to fall back to the default NPU stream instead of using the active `torch.npu` stream. This PR extends the optional addon so that Ascend kernels launched through TVM-FFI receive the correct current NPU stream. ## Changes ### Backend selection and addon loading - Extend the Torch addon backend selector with an `ascend` result. - Keep the existing selection order: CUDA and ROCm are checked first, followed by Ascend, then CPU. - Detect Ascend only on Linux. Windows and macOS fall back to the CPU addon because neither `torch_npu` nor the Ascend runtime supports those platforms. - Add the backend name to the existing ABI-aware addon cache key, preventing an Ascend addon from colliding with a CPU addon built against the same PyTorch package. - Force the TVM-FFI addon to be used for Ascend even when `torch.Tensor` already provides `__dlpack_c_exchange_api__`. The upstream PyTorch API can convert the tensor, but it does not provide the required Ascend `current_work_stream` behavior. ### Ascend addon build support - Add a mutually exclusive `--build-with-ascend` build option alongside the existing CUDA and ROCm options. - Reject the Ascend build option on Windows and macOS with an explicit argument error. - Discover the installed `torch_npu` package and add its header and library directories to the build. - Define `BUILD_WITH_ASCEND` only for the Ascend addon. - Link against `libtorch_npu`, which exports the `c10_npu` stream symbols used by the addon. ### Current NPU stream propagation When built with Ascend support, the addon now handles `kDLExtDev` in its `current_work_stream` callback and returns: ``` c10_npu::getCurrentNPUStream(device_id).stream() ``` This keeps TVM-FFI kernel launches on the active `torch.npu` stream, including non-default streams selected through `torch.npu.stream(...)`. ## Why Ascend detection has a separate `torch.npu` branch ROCm follows PyTorch’s AMD compatibility model and intentionally reuses the `torch.cuda` Python API. CUDA and ROCm can therefore share the same initial availability check and be distinguished using `torch.version.cuda` and `torch.version.hip`, as documented in the [[PyTorch ROCm semantics](https://docs.pytorch.org/docs/stable/notes/hip.html)](https://docs.pytorch.org/docs/stable/notes/hip.html). `torch_npu` follows a different integration model. It is an out-of-tree extension that registers Ascend through `PrivateUse1` and exposes its runtime API as `torch.npu`, rather than reusing `torch.cuda`. This model is also described in PyTorch’s [[device-extension autoload RFC](https://github.com/pytorch/pytorch/issues/122468)](https://github.com/pytorch/pytorch/issues/122468). As a result, Ascend cannot be folded into the CUDA/ROCm branch without changing the supported PyTorch/extension compatibility model. A generic `torch.accelerator`-based implementation was considered, but adopting it here would broaden the change across existing backends and supported PyTorch versions. The guarded: ``` hasattr(torch, "npu") and torch.npu.is_available() ``` check is therefore kept as an Ascend-specific compatibility path. CUDA and ROCm retain priority when their shared `torch.cuda` interface is active. ## DLPack device type follow-up This PR intentionally continues to use `kDLExtDev` for Ascend. A dedicated `kDLAscend` device type will be proposed to the DLPack project so that Ascend memory is no longer represented by the generic extension-device value. That migration is not included here because it requires coordinated changes across DLPack, TVM-FFI, PyTorch/`torch_npu`, and downstream consumers. At the time of writing: - TVM-FFI pins the released [[DLPack v1.3](https://github.com/dmlc/dlpack/releases/tag/v1.3)](https://github.com/dmlc/dlpack/releases/tag/v1.3) commit `84d107b`. - DLPack `main` is three commits ahead at [`[77aafa4](https://github.com/dmlc/dlpack/commit/77aafa4d3b0f80feffce9ad4c718dd26751ee0e4)`](https://github.com/dmlc/dlpack/commit/77aafa4d3b0f80feffce9ad4c718dd26751ee0e4) and has already allocated values 19 and 20 to `kDLTPU` and `kDLTPUHost`. - Consequently, the numeric value and semantics of `kDLAscend` should be standardized in DLPack first instead of being provisionally assigned by TVM-FFI. After such an enum is accepted, TVM-FFI would need to update more than the vendored header. The required work includes: - Updating the DLPack submodule. - Extending C++ device parsing and formatting. - Extending the Python `DLDeviceType` enum, name mappings, type stubs, and device tests. - Updating the Torch addon’s `PrivateUse1` ↔ DLPack conversion logic. - Updating stream lookup to recognize the new type. - Providing a transition period in which both `kDLExtDev` and `kDLAscend` are accepted, because existing `torch_npu` versions will continue to produce `kDLExtDev`. Switching this PR directly to a new enum would therefore break compatibility with existing producers without completing the corresponding ecosystem migration. Keeping `kDLExtDev` here makes the stream fix usable immediately while leaving the enum standardization to a separate follow-up. ## Testing The PR adds coverage for: - CPU, CUDA, ROCm, and Ascend backend selection. - CUDA/ROCm precedence when an Ascend namespace is also present. - Windows and macOS fallback behavior. - Reuse of PyTorch’s existing exchange API on non-Ascend backends. - Forced addon selection on Ascend. - Mutually exclusive CUDA, ROCm, and Ascend build options. - Rejection of Ascend builds on unsupported platforms. - Building and loading the addon from a fresh per-test output directory. - An Ascend runtime test that extracts the installed DLPack exchange API and compares `current_work_stream(kDLExtDev, device_id)` against a non-default `torch.npu.Stream` handle. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
