zihaomu opened a new pull request, #585: URL: https://github.com/apache/tvm-ffi/pull/585
This PR fixes ROCm support for the optional Torch C DLPack extension. Previously, the extension backend was selected using `torch.cuda.is_available()`. This is not enough because PyTorch on ROCm also uses the `torch.cuda` API. As a result, ROCm environments could be misdetected as CUDA and try to load CUDA-linked libraries, leading to import errors. This PR updates the logic to explicitly distinguish CUDA, ROCm, and CPU using `torch.version.cuda` and `torch.version.hip`. It also falls back to local JIT build when an installed prebuilt extension is not compatible with the current backend.Regression tests are added for backend detection and GPU tensor metadata through DLPack. -- 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]
