[Fix] Discover flat CUDA includes for NVRTC

[morluto]

Fixes #2828.

Problem

On Linux, the NVRTC backend always derived CCCL headers from $CUDA_HOME/targets/<arch>-linux/include/cccl. Flat pip CUDA installations instead provide $CUDA_HOME/include and $CUDA_HOME/include/cccl, so the generated NVRTC options omitted the existing CCCL tree and pointed at a non-existent target-specific path.

A synthetic flat tree reproduced the mismatch on the base revision:

flat_include=/var/tmp/.../include exists=True
flat_cccl=/var/tmp/.../include/cccl exists=True
old_cccl_option=/var/tmp/.../targets/x86_64-linux/include/cccl exists=False

Change

Discover existing flat and host-target include roots and their CCCL subdirectories, preserving their deterministic order. Installations that overlay both layouts include both. Windows retains its flat include behavior, and a missing Linux layout retains the previous fallback candidate list so NVRTC still owns the final missing-header diagnostic.

CUDA versions before 13 continue to receive the corresponding legacy cuda/std include for each discovered base include root.

Regression coverage

The new synthetic-layout tests construct temporary directory trees for:

• flat Linux pip installs;
• x86_64 target-specific system installs;
• combined flat and SBSA target-specific installs;
• missing Linux layouts and their legacy fallback;
• the existing flat Windows layout.

The tests never inspect the host CUDA installation and would fail on the base implementation because it has no flat CCCL discovery path.

Validation

• Focused pytest launcher removed the main checkout's editable _Finder and ScikitBuildRedirectingFinder entries from sys.meta_path, then printed tilelang.__file__ before calling pytest.main(["testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py", "-q"]). It resolved /home/morluto/dev/tilelang-agent-nvrtc/tilelang/__init__.py and reported 5 passed.
• python -m pre_commit run --files tilelang/jit/adapter/nvrtc/include_paths.py tilelang/jit/adapter/nvrtc/libgen.py testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py — all hooks passed.

Summary

• Added deterministic NVRTC CUDA include discovery for Linux flat pip layouts and target-specific toolkit layouts.
• Included existing cccl subdirectories for each discovered include root.
• Preserved Windows flat-layout handling and legacy fallback paths when Linux directories are missing.
• Preserved legacy cuda/std includes for CUDA versions before 13.
• Updated NVRTC compilation to use the discovered include paths.
• Added regression tests for flat, target-specific, combined, missing, and Windows layouts.
• Targeted tests and pre-commit hooks pass.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 7ea90f71-228c-4b43-b5ad-0e706689eb9c

📥 Commits

Reviewing files that changed from the base of the PR and between 171c51b and 998abf7.

📒 Files selected for processing (1)

• testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py

🚧 Files skipped from review as they are similar to previous changes (1)

• testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py

---

📝 Walkthrough

Walkthrough

The NVRTC backend now discovers CUDA include directories across flat, target-specific, mixed, missing-layout, and Windows installations. NVRTC generation uses the discovered directories and conditionally adds legacy cuda/std paths.

Changes

NVRTC include-path discovery

Layer / File(s)	Summary
Discover and validate CUDA include layouts tilelang/jit/adapter/nvrtc/include_paths.py, testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py	The helper selects sbsa-linux or x86_64-linux, handles Windows flat paths, discovers existing include and cccl directories, and preserves fallback paths. Tests cover all supported layouts.
Use discovered paths during NVRTC generation tilelang/jit/adapter/nvrtc/libgen.py	NVRTC generation replaces hard-coded platform paths with discovered directories and adds legacy cuda/std paths when the directory is not already under include/cccl.

Estimated code review effort: 3 (Moderate) | ~20 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the primary change: discovering flat CUDA include paths for NVRTC.
Linked Issues check	✅ Passed	The implementation and tests satisfy issue #2828 by supporting flat, target-specific, combined, Windows, and fallback CUDA include layouts.
Out of Scope Changes check	✅ Passed	All changes support NVRTC CUDA include discovery or provide regression coverage for the linked issue objectives.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py`:
- Around line 15-49: Update every discover_cuda_include_paths call in the four
tests to pass system="linux" explicitly, while preserving the existing machine
arguments, paths, and expected results.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: b34d6935-88e5-43b5-8dd1-811e5dc6d83e

📥 Commits

Reviewing files that changed from the base of the PR and between bdb769a and 171c51b.

📒 Files selected for processing (3)

• testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py
• tilelang/jit/adapter/nvrtc/include_paths.py
• tilelang/jit/adapter/nvrtc/libgen.py