text: # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Prints CUDA library and header directories and versions found on the system.
  
  The script searches for CUDA library and header files on the system, inspects
  them to determine their version and prints the configuration to stdout.
  The paths to inspect and the required versions are specified through environment
  variables. If no valid configuration is found, the script prints to stderr and
  returns an error code.
  
  The list of libraries to find is specified as arguments. Supported libraries are
  CUDA (includes cuBLAS), cuDNN, NCCL, and TensorRT.
  
  The script takes a list of base directories specified by the TF_CUDA_PATHS
  environment variable as comma-separated glob list. The script looks for headers
  and library files in a hard-coded set of subdirectories from these base paths.
  If TF_CUDA_PATHS is not specified, a OS specific default is used:
  
    Linux:   /usr/local/cuda, /usr, and paths from 'ldconfig -p'.
    Windows: CUDA_PATH environment variable, or
             C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\*
  
  For backwards compatibility, some libraries also use alternative base
  directories from other environment variables if they are specified. List of
  library-specific environment variables:
  
    Library   Version env variable  Additional base directories
    ----------------------------------------------------------------
    CUDA      TF_CUDA_VERSION       CUDA_TOOLKIT_PATH
    cuBLAS    TF_CUBLAS_VERSION     CUDA_TOOLKIT_PATH
    cuDNN     TF_CUDNN_VERSION      CUDNN_INSTALL_PATH
    NCCL      TF_NCCL_VERSION       NCCL_INSTALL_PATH, NCCL_HDR_PATH
    TensorRT  TF_TENSORRT_VERSION   TENSORRT_INSTALL_PATH
  
  Versions environment variables can be of the form 'x' or 'x.y' to request a
  specific version, empty or unspecified to accept any version.
  
  The output of a found library is of the form:
  tf_<library>_version: x.y.z
  tf_<library>_header_dir: ...
  tf_<library>_library_dir: ...
  """
  
  import os
  import glob
  import platform
  import re
  import subprocess
  import sys
  
  # pylint: disable=g-import-not-at-top
  try:
    from shutil import which
  except ImportError:
    from distutils.spawn import find_executable as which
  # pylint: enable=g-import-not-at-top
  
  
  class ConfigError(Exception):
    pass
  
  
  def _is_linux():
    return platform.system() == "Linux"
  
  
  def _is_windows():
    return platform.system() == "Windows"
  
  
  def _is_macos():
    return platform.system() == "Darwin"
  
  
  def _matches_version(actual_version, required_version):
    """Checks whether some version meets the requirements.
  
        All elements of the required_version need to be present in the
        actual_version.
  
            required_version  actual_version  result
            -----------------------------------------
            1                 1.1             True
            1.2               1               False
            1.2               1.3             False
                              1               True
  
        Args:
          required_version: The version specified by the user.
          actual_version: The version detected from the CUDA installation.
        Returns: Whether the actual version matches the required one.
    """
    if actual_version is None:
      return False
  
    # Strip spaces from the versions.
    actual_version = actual_version.strip()
    required_version = required_version.strip()
    return actual_version.startswith(required_version)
  
  
  def _at_least_version(actual_version, required_version):
    actual = [int(v) for v in actual_version.split(".")]
    required = [int(v) for v in required_version.split(".")]
    return actual >= required
  
  
  def _get_header_version(path, name):
    """Returns preprocessor defines in C header file."""
    for line in open(path, "r").readlines():
      match = re.match("#define %s +(\d+)" % name, line)
      if match:
        return match.group(1)
    return ""
  
  
  def _cartesian_product(first, second):
    """Returns all path combinations of first and second."""
    return [os.path.join(f, s) for f in first for s in second]
  
  
  def _get_ld_config_paths():
    """Returns all directories from 'ldconfig -p'."""
    if not _is_linux():
      return []
    ldconfig_path = which("ldconfig") or "/sbin/ldconfig"
    output = subprocess.check_output([ldconfig_path, "-p"])
    pattern = re.compile(".* => (.*)")
    result = set()
    for line in output.splitlines():
      match = pattern.match(line.decode("ascii"))
      if match:
        result.add(os.path.dirname(match.group(1)))
    return sorted(list(result))
  
  
  def _get_default_cuda_paths(cuda_version):
    if not cuda_version:
      cuda_version = "*"
    elif not "." in cuda_version:
      cuda_version = cuda_version + ".*"
  
    if _is_windows():
      return [
          os.environ.get(
              "CUDA_PATH",
              "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v%s\\" %
              cuda_version)
      ]
    return ["/usr/local/cuda-%s" % cuda_version, "/usr/local/cuda", "/usr"
           ] + _get_ld_config_paths()
  
  
  def _header_paths():
    """Returns hard-coded set of relative paths to look for header files."""
    return [
        "",
        "include",
        "include/cuda",
        "include/*-linux-gnu",
        "extras/CUPTI/include",
        "include/cuda/CUPTI",
    ]
  
  
  def _library_paths():
    """Returns hard-coded set of relative paths to look for library files."""
    return [
        "",
        "lib64",
        "lib",
        "lib/*-linux-gnu",
        "lib/x64",
        "extras/CUPTI/*",
    ]
  
  
  def _not_found_error(base_paths, relative_paths, filepattern):
    base_paths = "".join(["\n        '%s'" % path for path in sorted(base_paths)])
    relative_paths = "".join(["\n        '%s'" % path for path in relative_paths])
    return ConfigError(
        "Could not find any %s in any subdirectory:%s\nof:%s\n" %
        (filepattern, relative_paths, base_paths))
  
  
  def _find_file(base_paths, relative_paths, filepattern):
    for path in _cartesian_product(base_paths, relative_paths):
      for file in glob.glob(os.path.join(path, filepattern)):
        return file
    raise _not_found_error(base_paths, relative_paths, filepattern)
  
  
  def _find_library(base_paths, library_name, required_version):
    """Returns first valid path to the requested library."""
    if _is_windows():
      filepattern = library_name + ".lib"
    elif _is_macos():
      filepattern = "%s*.dylib" % (".".join(["lib" + library_name] +
                                            required_version.split(".")[:1]))
    else:
      filepattern = ".".join(["lib" + library_name, "so"] +
                             required_version.split(".")[:1]) + "*"
    return _find_file(base_paths, _library_paths(), filepattern)
  
  
  def _find_versioned_file(base_paths, relative_paths, filepattern,
                           required_version, get_version):
    """Returns first valid path to a file that matches the requested version."""
    for path in _cartesian_product(base_paths, relative_paths):
      for file in glob.glob(os.path.join(path, filepattern)):
        actual_version = get_version(file)
        if _matches_version(actual_version, required_version):
          return file, actual_version
    raise _not_found_error(
        base_paths, relative_paths,
        filepattern + " matching version '%s'" % required_version)
  
  
  def _find_header(base_paths, header_name, required_version, get_version):
    """Returns first valid path to a header that matches the requested version."""
    return _find_versioned_file(base_paths, _header_paths(), header_name,
                                required_version, get_version)
  
  
  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 = "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",
    ], nvcc_name, cuda_version, get_nvcc_version)
  
    nvvm_path = _find_file(base_paths, [
        "nvvm/libdevice",
        "share/cuda",
        "lib/nvidia-cuda-toolkit/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],
    }
  
  
  def _find_cublas_config(base_paths, required_version, cuda_version):
  
    if _at_least_version(cuda_version, "10.1"):
  
      def get_header_version(path):
        version = (
            _get_header_version(path, name)
            for name in ("CUBLAS_VER_MAJOR", "CUBLAS_VER_MINOR",
                         "CUBLAS_VER_PATCH"))
        return ".".join(version)
  
      header_path, header_version = _find_header(base_paths, "cublas_api.h",
                                                 required_version,
                                                 get_header_version)
      # cuBLAS uses the major version only.
      cublas_version = header_version.split(".")[0]
  
      if not _matches_version(cuda_version, cublas_version):
        raise ConfigError("cuBLAS version %s does not match CUDA version %s" %
                          (cublas_version, cuda_version))
  
    else:
      # There is no version info available before CUDA 10.1, just find the file.
      header_path = _find_file(base_paths, _header_paths(), "cublas_api.h")
      # cuBLAS version is the same as CUDA version (x.y).
      cublas_version = required_version
  
    library_path = _find_library(base_paths, "cublas", cublas_version)
  
    return {
        "cublas_include_dir": os.path.dirname(header_path),
        "cublas_library_dir": os.path.dirname(library_path),
    }
  
  
  def _find_cudnn_config(base_paths, required_version):
  
    def get_header_version(path):
      version = (
          _get_header_version(path, name)
          for name in ("CUDNN_MAJOR", "CUDNN_MINOR", "CUDNN_PATCHLEVEL"))
      return ".".join(version)
  
    header_path, header_version = _find_header(base_paths, "cudnn.h",
                                               required_version,
                                               get_header_version)
    cudnn_version = header_version.split(".")[0]
  
    library_path = _find_library(base_paths, "cudnn", cudnn_version)
  
    return {
        "cudnn_version": cudnn_version,
        "cudnn_include_dir": os.path.dirname(header_path),
        "cudnn_library_dir": os.path.dirname(library_path),
    }
  
  
  def _find_nccl_config(base_paths, required_version):
  
    def get_header_version(path):
      version = (
          _get_header_version(path, name)
          for name in ("NCCL_MAJOR", "NCCL_MINOR", "NCCL_PATCH"))
      return ".".join(version)
  
    header_path, header_version = _find_header(base_paths, "nccl.h",
                                               required_version,
                                               get_header_version)
    nccl_version = header_version.split(".")[0]
  
    library_path = _find_library(base_paths, "nccl", nccl_version)
  
    return {
        "nccl_version": nccl_version,
        "nccl_include_dir": os.path.dirname(header_path),
        "nccl_library_dir": os.path.dirname(library_path),
    }
  
  
  def _find_tensorrt_config(base_paths, required_version):
  
    def get_header_version(path):
      version = (
          _get_header_version(path, name)
          for name in ("NV_TENSORRT_MAJOR", "NV_TENSORRT_MINOR",
                       "NV_TENSORRT_PATCH"))
      return ".".join(version)
  
    header_path, header_version = _find_header(base_paths, "NvInfer.h",
                                               required_version,
                                               get_header_version)
    tensorrt_version = header_version.split(".")[0]
  
    library_path = _find_library(base_paths, "nvinfer", tensorrt_version)
  
    return {
        "tensorrt_version": tensorrt_version,
        "tensorrt_include_dir": os.path.dirname(header_path),
        "tensorrt_library_dir": os.path.dirname(library_path),
    }
  
  
  def _list_from_env(env_name, default=[]):
    """Returns comma-separated list from environment variable."""
    if env_name in os.environ:
      return os.environ[env_name].split(",")
    return default
  
  
  def _get_legacy_path(env_name, default=[]):
    """Returns a path specified by a legacy environment variable.
  
    CUDNN_INSTALL_PATH, NCCL_INSTALL_PATH, TENSORRT_INSTALL_PATH set to
    '/usr/lib/x86_64-linux-gnu' would previously find both library and header
    paths. Detect those and return '/usr', otherwise forward to _list_from_env().
    """
    if env_name in os.environ:
      match = re.match("^(/[^/ ]*)+/lib/\w+-linux-gnu/?$", os.environ[env_name])
      if match:
        return [match.group(1)]
    return _list_from_env(env_name, default)
  
  
  def _normalize_path(path):
    """Returns normalized path, with forward slashes on Windows."""
    path = os.path.normpath(path)
    if _is_windows():
      path = path.replace("\\", "/")
    return path
  
  
  def find_cuda_config():
    """Returns a dictionary of CUDA library and header file paths."""
    libraries = [argv.lower() for argv in sys.argv[1:]]
    cuda_version = os.environ.get("TF_CUDA_VERSION", "")
    base_paths = _list_from_env("TF_CUDA_PATHS",
                                _get_default_cuda_paths(cuda_version))
    base_paths = [path for path in base_paths if os.path.exists(path)]
  
    result = {}
    if "cuda" in libraries:
      cuda_paths = _list_from_env("CUDA_TOOLKIT_PATH", base_paths)
      result.update(_find_cuda_config(cuda_paths, cuda_version))
  
      cublas_paths = _list_from_env("CUBLAS_INSTALL_PATH", base_paths)
      # Add cuda paths in case CuBLAS is installed under CUDA_TOOLKIT_PATH.
      cublas_paths += list(set(cuda_paths) - set(cublas_paths))
      cuda_version = result["cuda_version"]
      cublas_version = os.environ.get("TF_CUBLAS_VERSION", "")
      result.update(
          _find_cublas_config(cublas_paths, cublas_version, cuda_version))
  
    if "cudnn" in libraries:
      cudnn_paths = _get_legacy_path("CUDNN_INSTALL_PATH", base_paths)
      cudnn_version = os.environ.get("TF_CUDNN_VERSION", "")
      result.update(_find_cudnn_config(cudnn_paths, cudnn_version))
  
    if "nccl" in libraries:
      nccl_paths = _get_legacy_path("NCCL_INSTALL_PATH", base_paths)
      nccl_version = os.environ.get("TF_NCCL_VERSION", "")
      result.update(_find_nccl_config(nccl_paths, nccl_version))
  
    if "tensorrt" in libraries:
      tensorrt_paths = _get_legacy_path("TENSORRT_INSTALL_PATH", base_paths)
      tensorrt_version = os.environ.get("TF_TENSORRT_VERSION", "")
      result.update(_find_tensorrt_config(tensorrt_paths, tensorrt_version))
  
    for k, v in result.items():
      if k.endswith("_dir") or k.endswith("_path"):
        result[k] = _normalize_path(v)
  
    return result
  
  
  def main():
    try:
      for key, value in sorted(find_cuda_config().items()):
        print("%s: %s" % (key, value))
    except ConfigError as e:
      sys.stderr.write(str(e))
      sys.exit(1)
  
  
  if __name__ == "__main__":
    main()
  