spectrometerHBH opened a new pull request, #20222:
URL: https://github.com/apache/tvm/pull/20222
### Problem
`_compile_cuda_nvrtc` builds the architecture-specific CUDA include
directory as
`targets/{platform.machine()}-{platform.system().lower()}/include`, i.e.
`targets/aarch64-linux` on ARM64.
CUDA toolkits for ARM64 **servers** (SBSA — GH200/GB200 and similar) install
their headers under `targets/sbsa-linux`; `targets/aarch64-linux` only exists
in the embedded/L4T toolkits. On an SBSA machine the computed path therefore
does not exist, so both the architecture-specific include directory and the
CCCL include directory derived from it are silently dropped from the NVRTC
options. When the toolkit has no top-level `include/` either, compilation fails
with "Cannot find CUDA headers".
### Fix
Add `_find_cuda_target_include()`, which probes `sbsa-linux` first on ARM64
Linux (`aarch64`/`arm64`) and falls back to `aarch64-linux` for embedded/L4T
layouts. Other platforms keep the previous `{machine}-{system}` behavior. The
helper returns `None` when no `targets/` layout is present, which is also a
small cleanup of the CCCL lookup (it no longer re-runs `os.path.isdir` on the
same directory).
This mirrors the CUDA target selection used by other projects on SBSA hosts,
e.g. [Triton-distributed's
`build_scm.sh`](https://github.com/ByteDance-Seed/Triton-distributed/blob/8260bc34398c2b8f36dc840fd22f741ca9294584/build_scm.sh#L21).
### Testing
New `tests/python/support/test_nvcc.py` covers the SBSA, L4T, and x86_64
layouts plus the no-`targets/` case against a fake toolkit tree; it runs on any
host since the platform is monkeypatched.
```
python -m pytest tests/python/support/test_nvcc.py -q # 5 passed
pre-commit run --files python/tvm/support/nvcc.py
tests/python/support/test_nvcc.py
```
--
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]
