def _find_cuda_config()

in build_deps/toolchains/gpu/find_cuda_config.py [0:0]

• 62 lines of code
• 7 McCabe index (conditional complexity)

def _find_cuda_config(base_paths, required_version):
    def get_header_version(path):
        version = int(_get_header_version(path, "CUDA_VERSION"))
        if not version:
            return None
        return "%d.%d" % (version // 1000, version % 1000 // 10)

    cuda_header_path, header_version = _find_header(
        base_paths, "cuda.h", required_version, get_header_version
    )
    cuda_version = header_version  # x.y, see above.

    cuda_library_path = _find_library(base_paths, "cudart", cuda_version)

    def get_nvcc_version(path):
        pattern = r"Cuda compilation tools, release \d+\.\d+, V(\d+\.\d+\.\d+)"
        for line in subprocess.check_output([path, "--version"]).splitlines():
            match = re.match(pattern, line.decode("ascii"))
            if match:
                return match.group(1)
        return None

    nvcc_name = "nvcc.exe" if _is_windows() else "nvcc"
    nvcc_path, nvcc_version = _find_versioned_file(
        base_paths,
        [
            "",
            "bin",
            "local/cuda/bin",
        ],
        nvcc_name,
        cuda_version,
        get_nvcc_version,
    )

    nvvm_path = _find_file(
        base_paths,
        [
            "nvvm/libdevice",
            "share/cuda",
            "lib/nvidia-cuda-toolkit/libdevice",
            "local/cuda/nvvm/libdevice",
        ],
        "libdevice*.10.bc",
    )

    cupti_header_path = _find_file(base_paths, _header_paths(), "cupti.h")
    cupti_library_path = _find_library(base_paths, "cupti", required_version)

    cuda_binary_dir = os.path.dirname(nvcc_path)
    nvvm_library_dir = os.path.dirname(nvvm_path)

    # XLA requires the toolkit path to find ptxas and libdevice.
    # TODO(csigg): pass in both directories instead.
    cuda_toolkit_paths = (
        os.path.normpath(os.path.join(cuda_binary_dir, "..")),
        os.path.normpath(os.path.join(nvvm_library_dir, "../..")),
    )
    if cuda_toolkit_paths[0] != cuda_toolkit_paths[1]:
        raise ConfigError(
            "Inconsistent CUDA toolkit path: %s vs %s" % cuda_toolkit_paths
        )

    return {
        "cuda_version": cuda_version,
        "cuda_include_dir": os.path.dirname(cuda_header_path),
        "cuda_library_dir": os.path.dirname(cuda_library_path),
        "cuda_binary_dir": cuda_binary_dir,
        "nvvm_library_dir": nvvm_library_dir,
        "cupti_include_dir": os.path.dirname(cupti_header_path),
        "cupti_library_dir": os.path.dirname(cupti_library_path),
        "cuda_toolkit_path": cuda_toolkit_paths[0],
    }