Repository: tensorflow/addons

Issue:
Build Addons with CUDA 10.1
TF 2.1 will be compiled with CUDA 10.1 by default:
https://github.com/tensorflow/tensorflow/issues/26289

Our current docker container supports only CUDA 10 so we'll want to upgrade at some point to match the environment. I'm not sure if this is an immediate breakage though as it looks like the nightlies are still compatible:
https://colab.research.google.com/drive/10Be-m2cRDRH-_l5qjSxRRFbpZsgdFV-I

@yifeif how would you like to track updating the build env? Should we create an issue in custom-op?


Discussion:

Resolving patch:
diff --git a/build_deps/toolchains/gpu/find_cuda_config.py b/build_deps/toolchains/gpu/find_cuda_config.py
new file mode 100644
--- /dev/null
+++ b/build_deps/toolchains/gpu/find_cuda_config.py
@@ -0,0 +1,498 @@
+# 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: ...
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+import os
+import glob
+import platform
+import re
+import subprocess
+import sys
+
+# pylint: disable=E0012
+try:
+    from shutil import which
+except ImportError:
+    from distutils.spawn import find_executable as which
+# pylint: enable=E0012
+
+
+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 io.open(path, "r", encoding="utf-8").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():
+        try:
+            match = pattern.match(line.decode("ascii"))
+        except UnicodeDecodeError:
+            match = False
+        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",
+        "/usr/local/cudnn"
+    ] + _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_path in glob.glob(os.path.join(path, filepattern)):
+            return file_path
+    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_path in glob.glob(os.path.join(path, filepattern)):
+            actual_version = get_version(file_path)
+            if _matches_version(actual_version, required_version):
+                return file_path, 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, cublas_version = _find_header(
+            base_paths, "cublas_api.h", required_version, get_header_version)
+
+        cublas_major_verison = cublas_version.split(".")[0]
+        if not _matches_version(cuda_version, cublas_major_verison):
+            raise ConfigError(
+                "cuBLAS version %s does not match CUDA version %s" %
+                (cublas_major_verison, 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_major_verison)
+
+    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"))
+        # `version` is a generator object, so we convert it to a list before using
+        # it (muitiple times below).
+        version = list(version)
+        if not all(version):
+            return None  # Versions not found, make _matches_version returns False.
+        return ".".join(version)
+
+    try:
+        header_path, header_version = _find_header(
+            base_paths, "NvInfer.h", required_version, get_header_version)
+    except ConfigError:
+        # TensorRT 6 moved the version information to NvInferVersion.h.
+        header_path, header_version = _find_header(
+            base_paths, "NvInferVersion.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))
+
+        cuda_version = result["cuda_version"]
+        cublas_paths = base_paths
+        if tuple(int(v) for v in cuda_version.split(".")) < (10, 1):
+            # Before CUDA 10.1, cuBLAS was in the same directory as the toolkit.
+            cublas_paths = cuda_paths
+        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()


Test patch:
diff --git a/tools/ci_testing/addons_cpu.sh b/tools/ci_testing/addons_cpu.sh
--- a/tools/ci_testing/addons_cpu.sh
+++ b/tools/ci_testing/addons_cpu.sh
@@ -50,6 +50,7 @@ yes 'y' | ./configure.sh
 bazel test -c opt -k \
     --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
     --test_output=errors --local_test_jobs=8 \
+    --extra_toolchains=@bazel_tools//tools/python:autodetecting_toolchain_nonstrict \
     //tensorflow_addons/...
 
 exit $?
diff --git a/tools/ci_testing/addons_gpu.sh b/tools/ci_testing/addons_gpu.sh
--- a/tools/ci_testing/addons_gpu.sh
+++ b/tools/ci_testing/addons_gpu.sh
@@ -43,7 +43,8 @@ yes 'y' | ./configure.sh
 bazel test -c opt -k \
     --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
     --test_output=errors --local_test_jobs=8 \
-    --crosstool_top=//build_deps/toolchains/gcc7_manylinux2010-nvcc-cuda10.0:toolchain \
+    --crosstool_top=//build_deps/toolchains/gcc7_manylinux2010-nvcc-cuda10.1:toolchain \
+    --extra_toolchains=@bazel_tools//tools/python:autodetecting_toolchain_nonstrict \
     //tensorflow_addons/...
 
 exit $?