Function get_cuda_version

python/tvm/support/nvcc.py:723–762

Utility function to get CUDA version Parameters ---------- cuda_path : Optional[str] Path to CUDA root. If None is passed, will use `find_cuda_path()` as default. Returns ------- version : float The CUDA version

(cuda_path=None)

def get_cuda_version(cuda_path=None):
    """Utility function to get CUDA version

    Parameters
    ----------
    cuda_path : Optional[str]

        Path to CUDA root.  If None is passed, will use
        `find_cuda_path()` as default.

    Returns
    -------
    version : float
        The CUDA version

    """
    if cuda_path is None:
        cuda_path = find_cuda_path()

    version_file_path = os.path.join(cuda_path, "version.txt")
    if not os.path.exists(version_file_path):
        # Debian/Ubuntu repackaged CUDA path
        version_file_path = os.path.join(cuda_path, "lib", "cuda", "version.txt")
    try:
        with open(version_file_path) as f:
            version_str = f.read().strip().split()[-1]
            return tuple(int(field) for field in version_str.split("."))
    except FileNotFoundError:
        pass

    cmd = [os.path.join(cuda_path, "bin", "nvcc"), "--version"]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    (out, _) = proc.communicate()
    out = out.decode("utf-8", errors="replace")
    if proc.returncode == 0:
        release_line = next(line for line in out.split("\n") if "release" in line)
        release_fields = [s.strip() for s in release_line.split(",")]
        version_str = next(f[1:] for f in release_fields if f.startswith("V"))
        return tuple(int(field) for field in version_str.split("."))
    raise RuntimeError("Cannot read CUDA version file")


def find_nvshmem_paths() -> tuple[str, str]: