[Fix] Discover flat CUDA includes for NVRTC

testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+from tilelang.jit.adapter.nvrtc.include_paths import discover_cuda_include_paths
+
+
+def _make_include_tree(cuda_home: Path, relative_paths: list[str]) -> list[str]:
+    paths = []
+    for relative_path in relative_paths:
+        path = cuda_home / relative_path
+        path.mkdir(parents=True)
+        paths.append(str(path))
+    return paths
+
+
+def test_discovers_flat_pip_include_layout(tmp_path):
+    expected = _make_include_tree(tmp_path, ["include", "include/cccl"])
+
+    assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == expected
+
+
+def test_discovers_target_specific_system_include_layout(tmp_path):
+    expected = _make_include_tree(
+        tmp_path,
+        ["targets/x86_64-linux/include", "targets/x86_64-linux/include/cccl"],
+    )
+
+    assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == expected
+
+
+def test_discovers_flat_and_target_specific_include_layouts(tmp_path):
+    expected = _make_include_tree(
+        tmp_path,
+        [
+            "include",
+            "include/cccl",
+            "targets/sbsa-linux/include",
+            "targets/sbsa-linux/include/cccl",
+        ],
+    )
+
+    assert discover_cuda_include_paths(str(tmp_path), machine="aarch64", system="linux") == expected
+
+
+def test_preserves_legacy_paths_when_cuda_layout_is_missing(tmp_path):
+    assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == [
+        str(tmp_path / "include"),
+        str(tmp_path / "targets/x86_64-linux/include"),
+        str(tmp_path / "targets/x86_64-linux/include/cccl"),
+    ]
+
+
+def test_preserves_flat_windows_include_layout(tmp_path):
+    assert discover_cuda_include_paths(str(tmp_path), system="win32") == [
+        str(tmp_path / "include"),
+        str(tmp_path / "include/cccl"),
+    ]

tilelang/jit/adapter/nvrtc/include_paths.py

@@ -0,0 +1,43 @@
+"""CUDA include-path discovery for NVRTC compilation."""
+
+from __future__ import annotations
+
+import os.path as osp
+import platform
+import sys
+
+
+def _target_include_path(cuda_home: str, machine: str | None = None) -> str:
+    machine = machine or platform.machine()
+    target_arch = "sbsa-linux" if machine in ("aarch64", "arm64") else "x86_64-linux"
+    return osp.join(cuda_home, "targets", target_arch, "include")
+
+
+def discover_cuda_include_paths(cuda_home: str, machine: str | None = None, system: str | None = None) -> list[str]:
+    """Return NVRTC include paths supported by the CUDA layout at ``cuda_home``.
+
+    CUDA pip packages may install headers in a flat ``include`` tree, while
+    system toolkits commonly use ``targets/<arch>-linux/include``. Include both
+    layouts when present so split or overlaid installations continue to work.
+    """
+    flat_include = osp.join(cuda_home, "include")
+    system = system or sys.platform
+    if system.startswith("win32"):
+        return [flat_include, osp.join(flat_include, "cccl")]
+
+    target_include = _target_include_path(cuda_home, machine)
+    include_paths = []
+
+    for include_path in (flat_include, target_include):
+        if osp.isdir(include_path):
+            include_paths.append(include_path)
+            cccl_include = osp.join(include_path, "cccl")
+            if osp.isdir(cccl_include):
+                include_paths.append(cccl_include)
+
+    if include_paths:
+        return include_paths
+
+    # Preserve the previous Linux search paths when probing an incomplete or
+    # not-yet-mounted toolkit. NVRTC will then report the missing headers.
+    return [flat_include, target_include, osp.join(target_include, "cccl")]

tilelang/jit/adapter/nvrtc/libgen.py

@@ -18,8 +18,6 @@
 import importlib
 import logging
 import os.path as osp
-import platform
-import sys
 import tempfile
 from types import ModuleType
 
@@ -29,6 +27,7 @@
 from tilelang.jit.adapter.libgen import LibraryGenerator
 from tilelang.jit.adapter.utils import is_cuda_target
 from tilelang.jit.adapter.nvrtc import is_nvrtc_available, NVRTC_UNAVAILABLE_MESSAGE
+from tilelang.jit.adapter.nvrtc.include_paths import discover_cuda_include_paths
 
 logger = logging.getLogger(__name__)
 
@@ -185,29 +184,13 @@ def compile_lib(self, timeout: float | None = None):
 
             cuda_home = CUDA_HOME if CUDA_HOME else "/usr/local/cuda"
 
-            # CUDA Toolkit include layout differs by platform:
-            # * Linux pip wheel ``nvidia-cuXX`` and system CUDA both expose
-            #   per-target trees under ``targets/{arch}-linux/include``.
-            # * Windows pip wheel ``nvidia-cuXX`` and the system CUDA Toolkit
-            #   ship a single flat ``include/`` directory — no ``targets/``
-            #   subtree exists. Hard-coding ``x86_64-linux`` there sends nvrtc
-            #   to a non-existent path so headers like ``nvrtc_std.h`` fail to
-            #   resolve.
-            cuda_include = osp.join(cuda_home, "include")
-            if sys.platform.startswith("win32"):
-                arch_include = cuda_include
-            else:
-                machine = platform.machine()
-                target_arch = "sbsa-linux" if machine in ("aarch64", "arm64") else "x86_64-linux"
-                arch_include = osp.join(cuda_home, "targets", target_arch, "include")
+            cuda_include_paths = discover_cuda_include_paths(cuda_home)
 
             __CUDACC_VER_MAJOR__ = get_nvrtc_version()[0]
             options = [
                 f"-I{tl_template_path}",
                 f"-I{cutlass_path}",
-                f"-I{cuda_include}",
-                f"-I{arch_include}",
-                f"-I{arch_include}/cccl",
+                *(f"-I{include_path}" for include_path in cuda_include_paths),
                 f"-D__CUDACC_VER_MAJOR__={__CUDACC_VER_MAJOR__}",
             ]
 
@@ -220,7 +203,11 @@ def compile_lib(self, timeout: float | None = None):
             # forward declarations in ``cute/container/tuple.hpp`` under NVRTC
             # (cute uses variadic packs, cccl uses a single ``_Tp``).
             if __CUDACC_VER_MAJOR__ < 13:
-                options += [f"-I{arch_include}/cuda/std"]
+                options += [
+                    f"-I{include_path}/cuda/std"
+                    for include_path in cuda_include_paths
+                    if not include_path.endswith(osp.join("include", "cccl"))
+                ]
 
             if self.compile_flags:
                 options += [item for flag in self.compile_flags for item in flag.split() if item not in options]