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