1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Prints CUDA library and header directories and versions found on the system. 16 17The script searches for CUDA library and header files on the system, inspects 18them to determine their version and prints the configuration to stdout. 19The paths to inspect and the required versions are specified through environment 20variables. If no valid configuration is found, the script prints to stderr and 21returns an error code. 22 23The list of libraries to find is specified as arguments. Supported libraries are 24CUDA (includes cuBLAS), cuDNN, NCCL, and TensorRT. 25 26The script takes a list of base directories specified by the TF_CUDA_PATHS 27environment variable as comma-separated glob list. The script looks for headers 28and library files in a hard-coded set of subdirectories from these base paths. 29If TF_CUDA_PATHS is not specified, a OS specific default is used: 30 31 Linux: /usr/local/cuda, /usr, and paths from 'ldconfig -p'. 32 Windows: CUDA_PATH environment variable, or 33 C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\* 34 35For backwards compatibility, some libraries also use alternative base 36directories from other environment variables if they are specified. List of 37library-specific environment variables: 38 39 Library Version env variable Additional base directories 40 ---------------------------------------------------------------- 41 CUDA TF_CUDA_VERSION CUDA_TOOLKIT_PATH 42 cuBLAS TF_CUBLAS_VERSION CUDA_TOOLKIT_PATH 43 cuDNN TF_CUDNN_VERSION CUDNN_INSTALL_PATH 44 NCCL TF_NCCL_VERSION NCCL_INSTALL_PATH, NCCL_HDR_PATH 45 TensorRT TF_TENSORRT_VERSION TENSORRT_INSTALL_PATH 46 47Versions environment variables can be of the form 'x' or 'x.y' to request a 48specific version, empty or unspecified to accept any version. 49 50The output of a found library is of the form: 51tf_<library>_version: x.y.z 52tf_<library>_header_dir: ... 53tf_<library>_library_dir: ... 54""" 55 56import io 57import os 58import glob 59import platform 60import re 61import subprocess 62import sys 63 64# pylint: disable=g-import-not-at-top 65try: 66 from shutil import which 67except ImportError: 68 from distutils.spawn import find_executable as which 69# pylint: enable=g-import-not-at-top 70 71 72class ConfigError(Exception): 73 pass 74 75 76def _is_linux(): 77 return platform.system() == "Linux" 78 79 80def _is_windows(): 81 return platform.system() == "Windows" 82 83 84def _is_macos(): 85 return platform.system() == "Darwin" 86 87 88def _matches_version(actual_version, required_version): 89 """Checks whether some version meets the requirements. 90 91 All elements of the required_version need to be present in the 92 actual_version. 93 94 required_version actual_version result 95 ----------------------------------------- 96 1 1.1 True 97 1.2 1 False 98 1.2 1.3 False 99 1 True 100 101 Args: 102 required_version: The version specified by the user. 103 actual_version: The version detected from the CUDA installation. 104 Returns: Whether the actual version matches the required one. 105 """ 106 if actual_version is None: 107 return False 108 109 # Strip spaces from the versions. 110 actual_version = actual_version.strip() 111 required_version = required_version.strip() 112 return actual_version.startswith(required_version) 113 114 115def _at_least_version(actual_version, required_version): 116 actual = [int(v) for v in actual_version.split(".")] 117 required = [int(v) for v in required_version.split(".")] 118 return actual >= required 119 120 121def _get_header_version(path, name): 122 """Returns preprocessor defines in C header file.""" 123 for line in io.open(path, "r", encoding="utf-8").readlines(): 124 match = re.match("#define %s +(\d+)" % name, line) 125 if match: 126 return match.group(1) 127 return "" 128 129 130def _cartesian_product(first, second): 131 """Returns all path combinations of first and second.""" 132 return [os.path.join(f, s) for f in first for s in second] 133 134 135def _get_ld_config_paths(): 136 """Returns all directories from 'ldconfig -p'.""" 137 if not _is_linux(): 138 return [] 139 ldconfig_path = which("ldconfig") or "/sbin/ldconfig" 140 output = subprocess.check_output([ldconfig_path, "-p"]) 141 pattern = re.compile(".* => (.*)") 142 result = set() 143 for line in output.splitlines(): 144 try: 145 match = pattern.match(line.decode("ascii")) 146 except UnicodeDecodeError: 147 match = False 148 if match: 149 result.add(os.path.dirname(match.group(1))) 150 return sorted(list(result)) 151 152 153def _get_default_cuda_paths(cuda_version): 154 if not cuda_version: 155 cuda_version = "*" 156 elif not "." in cuda_version: 157 cuda_version = cuda_version + ".*" 158 159 if _is_windows(): 160 return [ 161 os.environ.get( 162 "CUDA_PATH", 163 "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v%s\\" % 164 cuda_version) 165 ] 166 return ["/usr/local/cuda-%s" % cuda_version, "/usr/local/cuda", "/usr", 167 "/usr/local/cudnn"] + _get_ld_config_paths() 168 169 170def _header_paths(): 171 """Returns hard-coded set of relative paths to look for header files.""" 172 return [ 173 "", 174 "include", 175 "include/cuda", 176 "include/*-linux-gnu", 177 "extras/CUPTI/include", 178 "include/cuda/CUPTI", 179 "local/cuda/extras/CUPTI/include", 180 ] 181 182 183def _library_paths(): 184 """Returns hard-coded set of relative paths to look for library files.""" 185 return [ 186 "", 187 "lib64", 188 "lib", 189 "lib/*-linux-gnu", 190 "lib/x64", 191 "extras/CUPTI/*", 192 "local/cuda/lib64", 193 "local/cuda/extras/CUPTI/lib64", 194 ] 195 196 197def _not_found_error(base_paths, relative_paths, filepattern): 198 base_paths = "".join(["\n '%s'" % path for path in sorted(base_paths)]) 199 relative_paths = "".join(["\n '%s'" % path for path in relative_paths]) 200 return ConfigError( 201 "Could not find any %s in any subdirectory:%s\nof:%s\n" % 202 (filepattern, relative_paths, base_paths)) 203 204 205def _find_file(base_paths, relative_paths, filepattern): 206 for path in _cartesian_product(base_paths, relative_paths): 207 for file in glob.glob(os.path.join(path, filepattern)): 208 return file 209 raise _not_found_error(base_paths, relative_paths, filepattern) 210 211 212def _find_library(base_paths, library_name, required_version): 213 """Returns first valid path to the requested library.""" 214 if _is_windows(): 215 filepattern = library_name + ".lib" 216 elif _is_macos(): 217 filepattern = "%s*.dylib" % (".".join(["lib" + library_name] + 218 required_version.split(".")[:1])) 219 else: 220 filepattern = ".".join(["lib" + library_name, "so"] + 221 required_version.split(".")[:1]) + "*" 222 return _find_file(base_paths, _library_paths(), filepattern) 223 224 225def _find_versioned_file(base_paths, relative_paths, filepatterns, 226 required_version, get_version): 227 """Returns first valid path to a file that matches the requested version.""" 228 if type(filepatterns) not in [list, tuple]: 229 filepatterns = [filepatterns] 230 for path in _cartesian_product(base_paths, relative_paths): 231 for filepattern in filepatterns: 232 for file in glob.glob(os.path.join(path, filepattern)): 233 actual_version = get_version(file) 234 if _matches_version(actual_version, required_version): 235 return file, actual_version 236 raise _not_found_error( 237 base_paths, relative_paths, 238 ", ".join(filepatterns) + " matching version '%s'" % required_version) 239 240 241def _find_header(base_paths, header_name, required_version, get_version): 242 """Returns first valid path to a header that matches the requested version.""" 243 return _find_versioned_file(base_paths, _header_paths(), header_name, 244 required_version, get_version) 245 246 247def _find_cuda_config(base_paths, required_version): 248 249 def get_header_version(path): 250 version = int(_get_header_version(path, "CUDA_VERSION")) 251 if not version: 252 return None 253 return "%d.%d" % (version // 1000, version % 1000 // 10) 254 255 cuda_header_path, header_version = _find_header(base_paths, "cuda.h", 256 required_version, 257 get_header_version) 258 cuda_version = header_version # x.y, see above. 259 260 cuda_library_path = _find_library(base_paths, "cudart", cuda_version) 261 262 def get_nvcc_version(path): 263 pattern = "Cuda compilation tools, release \d+\.\d+, V(\d+\.\d+\.\d+)" 264 for line in subprocess.check_output([path, "--version"]).splitlines(): 265 match = re.match(pattern, line.decode("ascii")) 266 if match: 267 return match.group(1) 268 return None 269 270 nvcc_name = "nvcc.exe" if _is_windows() else "nvcc" 271 nvcc_path, nvcc_version = _find_versioned_file(base_paths, [ 272 "", 273 "bin", 274 "local/cuda/bin", 275 ], nvcc_name, cuda_version, get_nvcc_version) 276 277 nvvm_path = _find_file(base_paths, [ 278 "nvvm/libdevice", 279 "share/cuda", 280 "lib/nvidia-cuda-toolkit/libdevice", 281 "local/cuda/nvvm/libdevice", 282 ], "libdevice*.10.bc") 283 284 cupti_header_path = _find_file(base_paths, _header_paths(), "cupti.h") 285 cupti_library_path = _find_library(base_paths, "cupti", required_version) 286 287 cuda_binary_dir = os.path.dirname(nvcc_path) 288 nvvm_library_dir = os.path.dirname(nvvm_path) 289 290 # XLA requires the toolkit path to find ptxas and libdevice. 291 # TODO(csigg): pass in both directories instead. 292 cuda_toolkit_paths = ( 293 os.path.normpath(os.path.join(cuda_binary_dir, "..")), 294 os.path.normpath(os.path.join(nvvm_library_dir, "../..")), 295 ) 296 if cuda_toolkit_paths[0] != cuda_toolkit_paths[1]: 297 raise ConfigError("Inconsistent CUDA toolkit path: %s vs %s" % 298 cuda_toolkit_paths) 299 300 return { 301 "cuda_version": cuda_version, 302 "cuda_include_dir": os.path.dirname(cuda_header_path), 303 "cuda_library_dir": os.path.dirname(cuda_library_path), 304 "cuda_binary_dir": cuda_binary_dir, 305 "nvvm_library_dir": nvvm_library_dir, 306 "cupti_include_dir": os.path.dirname(cupti_header_path), 307 "cupti_library_dir": os.path.dirname(cupti_library_path), 308 "cuda_toolkit_path": cuda_toolkit_paths[0], 309 } 310 311 312def _find_cublas_config(base_paths, required_version, cuda_version): 313 314 if _at_least_version(cuda_version, "10.1"): 315 316 def get_header_version(path): 317 version = ( 318 _get_header_version(path, name) 319 for name in ("CUBLAS_VER_MAJOR", "CUBLAS_VER_MINOR", 320 "CUBLAS_VER_PATCH")) 321 return ".".join(version) 322 323 header_path, header_version = _find_header(base_paths, "cublas_api.h", 324 required_version, 325 get_header_version) 326 # cuBLAS uses the major version only. 327 cublas_version = header_version.split(".")[0] 328 329 else: 330 # There is no version info available before CUDA 10.1, just find the file. 331 header_version = cuda_version 332 header_path = _find_file(base_paths, _header_paths(), "cublas_api.h") 333 # cuBLAS version is the same as CUDA version (x.y). 334 cublas_version = required_version 335 336 library_path = _find_library(base_paths, "cublas", cublas_version) 337 338 return { 339 "cublas_version": header_version, 340 "cublas_include_dir": os.path.dirname(header_path), 341 "cublas_library_dir": os.path.dirname(library_path), 342 } 343 344 345def _find_cusolver_config(base_paths, required_version, cuda_version): 346 347 if _at_least_version(cuda_version, "11.0"): 348 349 def get_header_version(path): 350 version = ( 351 _get_header_version(path, name) 352 for name in ("CUSOLVER_VER_MAJOR", "CUSOLVER_VER_MINOR", 353 "CUSOLVER_VER_PATCH")) 354 return ".".join(version) 355 356 header_path, header_version = _find_header(base_paths, "cusolver_common.h", 357 required_version, 358 get_header_version) 359 cusolver_version = header_version.split(".")[0] 360 361 else: 362 header_version = cuda_version 363 header_path = _find_file(base_paths, _header_paths(), "cusolver_common.h") 364 cusolver_version = required_version 365 366 library_path = _find_library(base_paths, "cusolver", cusolver_version) 367 368 return { 369 "cusolver_version": header_version, 370 "cusolver_include_dir": os.path.dirname(header_path), 371 "cusolver_library_dir": os.path.dirname(library_path), 372 } 373 374 375def _find_curand_config(base_paths, required_version, cuda_version): 376 377 if _at_least_version(cuda_version, "11.0"): 378 379 def get_header_version(path): 380 version = ( 381 _get_header_version(path, name) 382 for name in ("CURAND_VER_MAJOR", "CURAND_VER_MINOR", 383 "CURAND_VER_PATCH")) 384 return ".".join(version) 385 386 header_path, header_version = _find_header(base_paths, "curand.h", 387 required_version, 388 get_header_version) 389 curand_version = header_version.split(".")[0] 390 391 else: 392 header_version = cuda_version 393 header_path = _find_file(base_paths, _header_paths(), "curand.h") 394 curand_version = required_version 395 396 library_path = _find_library(base_paths, "curand", curand_version) 397 398 return { 399 "curand_version": header_version, 400 "curand_include_dir": os.path.dirname(header_path), 401 "curand_library_dir": os.path.dirname(library_path), 402 } 403 404 405def _find_cufft_config(base_paths, required_version, cuda_version): 406 407 if _at_least_version(cuda_version, "11.0"): 408 409 def get_header_version(path): 410 version = ( 411 _get_header_version(path, name) 412 for name in ("CUFFT_VER_MAJOR", "CUFFT_VER_MINOR", "CUFFT_VER_PATCH")) 413 return ".".join(version) 414 415 header_path, header_version = _find_header(base_paths, "cufft.h", 416 required_version, 417 get_header_version) 418 cufft_version = header_version.split(".")[0] 419 420 else: 421 header_version = cuda_version 422 header_path = _find_file(base_paths, _header_paths(), "cufft.h") 423 cufft_version = required_version 424 425 library_path = _find_library(base_paths, "cufft", cufft_version) 426 427 return { 428 "cufft_version": header_version, 429 "cufft_include_dir": os.path.dirname(header_path), 430 "cufft_library_dir": os.path.dirname(library_path), 431 } 432 433 434def _find_cudnn_config(base_paths, required_version): 435 436 def get_header_version(path): 437 version = [ 438 _get_header_version(path, name) 439 for name in ("CUDNN_MAJOR", "CUDNN_MINOR", "CUDNN_PATCHLEVEL")] 440 return ".".join(version) if version[0] else None 441 442 header_path, header_version = _find_header(base_paths, 443 ("cudnn.h", "cudnn_version.h"), 444 required_version, 445 get_header_version) 446 cudnn_version = header_version.split(".")[0] 447 448 library_path = _find_library(base_paths, "cudnn", cudnn_version) 449 450 return { 451 "cudnn_version": cudnn_version, 452 "cudnn_include_dir": os.path.dirname(header_path), 453 "cudnn_library_dir": os.path.dirname(library_path), 454 } 455 456 457def _find_cusparse_config(base_paths, required_version, cuda_version): 458 459 if _at_least_version(cuda_version, "11.0"): 460 461 def get_header_version(path): 462 version = ( 463 _get_header_version(path, name) 464 for name in ("CUSPARSE_VER_MAJOR", "CUSPARSE_VER_MINOR", 465 "CUSPARSE_VER_PATCH")) 466 return ".".join(version) 467 468 header_path, header_version = _find_header(base_paths, "cusparse.h", 469 required_version, 470 get_header_version) 471 cusparse_version = header_version.split(".")[0] 472 473 else: 474 header_version = cuda_version 475 header_path = _find_file(base_paths, _header_paths(), "cusparse.h") 476 cusparse_version = required_version 477 478 library_path = _find_library(base_paths, "cusparse", cusparse_version) 479 480 return { 481 "cusparse_version": header_version, 482 "cusparse_include_dir": os.path.dirname(header_path), 483 "cusparse_library_dir": os.path.dirname(library_path), 484 } 485 486 487def _find_nccl_config(base_paths, required_version): 488 489 def get_header_version(path): 490 version = ( 491 _get_header_version(path, name) 492 for name in ("NCCL_MAJOR", "NCCL_MINOR", "NCCL_PATCH")) 493 return ".".join(version) 494 495 header_path, header_version = _find_header(base_paths, "nccl.h", 496 required_version, 497 get_header_version) 498 nccl_version = header_version.split(".")[0] 499 500 library_path = _find_library(base_paths, "nccl", nccl_version) 501 502 return { 503 "nccl_version": nccl_version, 504 "nccl_include_dir": os.path.dirname(header_path), 505 "nccl_library_dir": os.path.dirname(library_path), 506 } 507 508 509def _find_tensorrt_config(base_paths, required_version): 510 511 def get_header_version(path): 512 version = ( 513 _get_header_version(path, name) 514 for name in ("NV_TENSORRT_MAJOR", "NV_TENSORRT_MINOR", 515 "NV_TENSORRT_PATCH")) 516 # `version` is a generator object, so we convert it to a list before using 517 # it (muitiple times below). 518 version = list(version) 519 if not all(version): 520 return None # Versions not found, make _matches_version returns False. 521 return ".".join(version) 522 523 try: 524 header_path, header_version = _find_header(base_paths, "NvInfer.h", 525 required_version, 526 get_header_version) 527 except ConfigError: 528 # TensorRT 6 moved the version information to NvInferVersion.h. 529 header_path, header_version = _find_header(base_paths, "NvInferVersion.h", 530 required_version, 531 get_header_version) 532 533 tensorrt_version = header_version.split(".")[0] 534 library_path = _find_library(base_paths, "nvinfer", tensorrt_version) 535 536 return { 537 "tensorrt_version": tensorrt_version, 538 "tensorrt_include_dir": os.path.dirname(header_path), 539 "tensorrt_library_dir": os.path.dirname(library_path), 540 } 541 542 543def _list_from_env(env_name, default=[]): 544 """Returns comma-separated list from environment variable.""" 545 if env_name in os.environ: 546 return os.environ[env_name].split(",") 547 return default 548 549 550def _get_legacy_path(env_name, default=[]): 551 """Returns a path specified by a legacy environment variable. 552 553 CUDNN_INSTALL_PATH, NCCL_INSTALL_PATH, TENSORRT_INSTALL_PATH set to 554 '/usr/lib/x86_64-linux-gnu' would previously find both library and header 555 paths. Detect those and return '/usr', otherwise forward to _list_from_env(). 556 """ 557 if env_name in os.environ: 558 match = re.match("^(/[^/ ]*)+/lib/\w+-linux-gnu/?$", os.environ[env_name]) 559 if match: 560 return [match.group(1)] 561 return _list_from_env(env_name, default) 562 563 564def _normalize_path(path): 565 """Returns normalized path, with forward slashes on Windows.""" 566 path = os.path.realpath(path) 567 if _is_windows(): 568 path = path.replace("\\", "/") 569 return path 570 571 572def find_cuda_config(): 573 """Returns a dictionary of CUDA library and header file paths.""" 574 libraries = [argv.lower() for argv in sys.argv[1:]] 575 cuda_version = os.environ.get("TF_CUDA_VERSION", "") 576 base_paths = _list_from_env("TF_CUDA_PATHS", 577 _get_default_cuda_paths(cuda_version)) 578 base_paths = [path for path in base_paths if os.path.exists(path)] 579 580 result = {} 581 if "cuda" in libraries: 582 cuda_paths = _list_from_env("CUDA_TOOLKIT_PATH", base_paths) 583 result.update(_find_cuda_config(cuda_paths, cuda_version)) 584 585 cuda_version = result["cuda_version"] 586 cublas_paths = base_paths 587 if tuple(int(v) for v in cuda_version.split(".")) < (10, 1): 588 # Before CUDA 10.1, cuBLAS was in the same directory as the toolkit. 589 cublas_paths = cuda_paths 590 cublas_version = os.environ.get("TF_CUBLAS_VERSION", "") 591 result.update( 592 _find_cublas_config(cublas_paths, cublas_version, cuda_version)) 593 594 cusolver_paths = base_paths 595 if tuple(int(v) for v in cuda_version.split(".")) < (11, 0): 596 cusolver_paths = cuda_paths 597 cusolver_version = os.environ.get("TF_CUSOLVER_VERSION", "") 598 result.update( 599 _find_cusolver_config(cusolver_paths, cusolver_version, cuda_version)) 600 601 curand_paths = base_paths 602 if tuple(int(v) for v in cuda_version.split(".")) < (11, 0): 603 curand_paths = cuda_paths 604 curand_version = os.environ.get("TF_CURAND_VERSION", "") 605 result.update( 606 _find_curand_config(curand_paths, curand_version, cuda_version)) 607 608 cufft_paths = base_paths 609 if tuple(int(v) for v in cuda_version.split(".")) < (11, 0): 610 cufft_paths = cuda_paths 611 cufft_version = os.environ.get("TF_CUFFT_VERSION", "") 612 result.update(_find_cufft_config(cufft_paths, cufft_version, cuda_version)) 613 614 cusparse_paths = base_paths 615 if tuple(int(v) for v in cuda_version.split(".")) < (11, 0): 616 cusparse_paths = cuda_paths 617 cusparse_version = os.environ.get("TF_CUSPARSE_VERSION", "") 618 result.update( 619 _find_cusparse_config(cusparse_paths, cusparse_version, cuda_version)) 620 621 if "cudnn" in libraries: 622 cudnn_paths = _get_legacy_path("CUDNN_INSTALL_PATH", base_paths) 623 cudnn_version = os.environ.get("TF_CUDNN_VERSION", "") 624 result.update(_find_cudnn_config(cudnn_paths, cudnn_version)) 625 626 if "nccl" in libraries: 627 nccl_paths = _get_legacy_path("NCCL_INSTALL_PATH", base_paths) 628 nccl_version = os.environ.get("TF_NCCL_VERSION", "") 629 result.update(_find_nccl_config(nccl_paths, nccl_version)) 630 631 if "tensorrt" in libraries: 632 tensorrt_paths = _get_legacy_path("TENSORRT_INSTALL_PATH", base_paths) 633 tensorrt_version = os.environ.get("TF_TENSORRT_VERSION", "") 634 result.update(_find_tensorrt_config(tensorrt_paths, tensorrt_version)) 635 636 for k, v in result.items(): 637 if k.endswith("_dir") or k.endswith("_path"): 638 result[k] = _normalize_path(v) 639 640 return result 641 642 643def main(): 644 try: 645 for key, value in sorted(find_cuda_config().items()): 646 print("%s: %s" % (key, value)) 647 except ConfigError as e: 648 sys.stderr.write(str(e) + '\n') 649 sys.exit(1) 650 651 652if __name__ == "__main__": 653 main() 654