Add build flags and config settings to enable CUDA support and select the CUDA device compiler.

Trigger the flags in the .bazelrc when '--config=cuda' or '--config=cuda_clang' is selected.

The pre-existing //tensorflow:using_cuda_clang/nvcc alias to the new config settings, instead of triggering themselves on '--define=...'.

PiperOrigin-RevId: 360523367
Change-Id: I7a56156655220e9d91fbc355689d27b95a0c22ab

.bazelrc

@@ -50,7 +50,6 @@
 # Feature and Third party library support options:
 #     xla:          Build TF with XLA
 #     tpu:          Build TF with TPU support
-#     using_cuda:   CUDA is available to build system.
 #     cuda:         Build with full cuda support.
 #     rocm:         Build with AMD GPU support (rocm).
 #     mkl:          Enable full mkl support.
@@ -191,23 +190,22 @@ build:mkl_aarch64 -c opt
 
 # This config refers to building with CUDA available. It does not necessarily
 # mean that we build CUDA op kernels.
-build:using_cuda --define=using_cuda=true
-build:using_cuda --action_env TF_NEED_CUDA=1
-build:using_cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
+build:cuda_base --@local_config_cuda//:enable_cuda
+build:cuda_base --action_env TF_NEED_CUDA=1
+build:cuda_base --crosstool_top=@local_config_cuda//crosstool:toolchain
 
 # Enable the mlir generated GPU kernels only for cuda builds.
 build --define=tensorflow_enable_mlir_generated_gpu_kernels=0
 # This is a more specific option, so it takes precedence over the line above for cuda builds.
-build:using_cuda --define=tensorflow_enable_mlir_generated_gpu_kernels=1
+build:cuda_base --define=tensorflow_enable_mlir_generated_gpu_kernels=1
 
 # This config refers to building CUDA op kernels with nvcc.
-build:cuda --config=using_cuda
-build:cuda --define=using_cuda_nvcc=true
+build:cuda --config=cuda_base
+build:cuda --@local_config_cuda//:cuda_compiler=nvcc
 
 # This config refers to building CUDA op kernels with clang.
-build:cuda_clang --config=using_cuda
-build:cuda_clang --define=using_cuda_clang=true
-build:cuda_clang --define=using_clang=true
+build:cuda_clang --config=cuda_base
+build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
 build:cuda_clang --action_env TF_CUDA_CLANG=1
 
 # dbg config, as a shorthand for '--config=opt -c dbg'
@@ -431,6 +429,9 @@ build:rbe_cpu_linux --host_platform="@ubuntu16.04-manylinux2010-py3_config_platf
 build:rbe_cpu_linux --platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform"
 
 build:rbe_linux_cuda_base --config=rbe_linux
+build:rbe_linux_cuda_base --@local_config_cuda//:enable_cuda
+build:rbe_linux_cuda_base --@local_config_cuda//:cuda_compiler=nvcc
+# TODO(csigg): those probably don't do anything because cuda_config is remote.
 build:rbe_linux_cuda_base --repo_env=TF_NEED_TENSORRT=1
 build:rbe_linux_cuda_base --repo_env=TF_CUDA_VERSION=10
 build:rbe_linux_cuda_base --repo_env=TF_CUDNN_VERSION=7
@@ -439,7 +440,6 @@ build:rbe_linux_cuda_base --repo_env=TF_NEED_CUDA=1
 test:rbe_linux_cuda_base --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
 
 build:rbe_linux_cuda10.1_nvcc_base --config=rbe_linux_cuda_base
-build:rbe_linux_cuda10.1_nvcc_base --define=using_cuda_nvcc=true
 build:rbe_linux_cuda10.1_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0_config_cuda//crosstool:toolchain"
 build:rbe_linux_cuda10.1_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0_config_cuda//crosstool:toolchain"
 build:rbe_linux_cuda10.1_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0_config_cuda//crosstool:toolchain-linux-x86_64"
@@ -456,7 +456,6 @@ build:rbe_linux_cuda10.1_nvcc_py3.7 --config=rbe_linux_cuda10.1_nvcc_base --repo
 build:rbe_linux_cuda10.1_nvcc_py3.8 --config=rbe_linux_cuda10.1_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0_config_python3.8"
 
 build:rbe_linux_cuda11.0_nvcc_base --config=rbe_linux_cuda_base
-build:rbe_linux_cuda11.0_nvcc_base --define=using_cuda_nvcc=true
 build:rbe_linux_cuda11.0_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.0-cudnn8-tensorrt7.1_config_cuda//crosstool:toolchain"
 build:rbe_linux_cuda11.0_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.0-cudnn8-tensorrt7.1_config_cuda//crosstool:toolchain"
 build:rbe_linux_cuda11.0_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.0-cudnn8-tensorrt7.1_config_cuda//crosstool:toolchain-linux-x86_64"

tensorflow/BUILD

@@ -197,7 +197,10 @@ config_setting(
 
 config_setting(
     name = "windows",
-    values = {"cpu": "x64_windows"},
+    # Internal builds query the target OS.
+    # copybara:uncomment flag_values = {"//tools/cpp:cc_target_os": "windows"},
+    # OSS builds query the CPU type.
+    values = {"cpu": "x64_windows"},  # copybara:comment
     visibility = ["//visibility:public"],
 )
 
@@ -411,11 +414,12 @@ config_setting(
 
 # Crosses between platforms and file system libraries not supported on those
 # platforms due to limitations in nested select() statements.
-config_setting(
+selects.config_setting_group(
     name = "with_cuda_support_windows_override",
-    define_values = {"using_cuda_nvcc": "true"},
-    values = {"cpu": "x64_windows"},
-    visibility = ["//visibility:public"],
+    match_all = [
+        ":using_cuda_nvcc",
+        ":windows",
+    ],
 )
 
 config_setting(
@@ -479,9 +483,46 @@ selects.config_setting_group(
     ],
 )
 
-config_setting(
+# Config setting that is satisfied when TensorFlow is being built with CUDA
+# support through e.g. `--config=cuda` (or `--config=cuda_clang` in OSS).
+alias(
+    name = "is_cuda_enabled",
+    actual = if_oss(
+        "@local_config_cuda//:is_cuda_enabled",
+        "@local_config_cuda//cuda:using_clang",
+    ),
+)
+
+# Config setting that is satisfied when CUDA device code should be compiled
+# with clang. It does not imply that CUDA support has been enabled.
+alias(
+    name = "is_cuda_compiler_clang",
+    actual = if_oss(
+        "@local_config_cuda//:is_cuda_compiler_clang",
+        "@local_config_cuda//cuda:TRUE",
+    ),
+)
+
+# Config setting that is satisfied when CUDA device code should be compiled
+# with nvcc. It does not imply that CUDA support has been enabled.
+alias(
+    name = "is_cuda_compiler_nvcc",
+    actual = if_oss(
+        "@local_config_cuda//:is_cuda_compiler_nvcc",
+        "@local_config_cuda//cuda:FALSE",
+    ),
+)
+
+# Config setting whether TensorFlow is built with CUDA support using clang.
+alias(
     name = "using_cuda_clang",
-    define_values = {"using_cuda_clang": "true"},
+    actual = "@local_config_cuda//cuda:using_clang",
+)
+
+# Config setting whether TensorFlow is built with CUDA support using nvcc.
+alias(
+    name = "using_cuda_nvcc",
+    actual = "@local_config_cuda//cuda:using_nvcc",
 )
 
 # Config setting to use in select()s to distinguish open source build from
@@ -502,12 +543,12 @@ bool_setting(
     visibility = ["//visibility:private"],
 )
 
-config_setting(
+selects.config_setting_group(
     name = "using_cuda_clang_with_dynamic_build",
-    define_values = {
-        "using_cuda_clang": "true",
-        "framework_shared_object": "true",
-    },
+    match_all = [
+        ":using_cuda_clang",
+        ":framework_shared_object",
+    ],
 )
 
 selects.config_setting_group(
@@ -528,19 +569,6 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
-config_setting(
-    name = "using_cuda_nvcc",
-    define_values = {"using_cuda_nvcc": "true"},
-)
-
-config_setting(
-    name = "using_cuda_nvcc_with_dynamic_build",
-    define_values = {
-        "using_cuda_nvcc": "true",
-        "framework_shared_object": "true",
-    },
-)
-
 selects.config_setting_group(
     name = "build_oss_using_cuda_nvcc",
     match_all = [
@@ -568,15 +596,6 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
-# This flag is defined for select statements that match both
-# on 'windows' and 'api_version_2'. In this case, bazel requires
-# having a flag which is a superset of these two.
-config_setting(
-    name = "windows_and_api_version_2",
-    define_values = {"tf_api_version": "2"},
-    values = {"cpu": "x64_windows"},
-)
-
 # This flag enables experimental MLIR support.
 config_setting(
     name = "with_mlir_support",

tensorflow/core/platform/default/build_config/BUILD

@@ -27,6 +27,20 @@ cc_library(
     deps = [],
 )
 
+cc_library(
+    name = "_empty_lib",
+    visibility = ["//visibility:private"],
+)
+
+alias(
+    name = "_cuda_runtime",
+    actual = select({
+        "//tensorflow:framework_shared_object": ":_empty_lib",
+        "//conditions:default": "//tensorflow/stream_executor/cuda:all_runtime",
+    }),
+    visibility = ["//visibility:private"],
+)
+
 tf_cuda_library(
     name = "stream_executor",
     cuda_deps = [
@@ -47,11 +61,11 @@ tf_cuda_library(
         "//tensorflow:macos": ["IOKit"],
         "//conditions:default": [],
     }) + select({
-        "//tensorflow:using_cuda_clang": ["//tensorflow/stream_executor/cuda:all_runtime"],
-        "//tensorflow:using_cuda_nvcc": ["//tensorflow/stream_executor/cuda:all_runtime"],
-        "//tensorflow:using_cuda_clang_with_dynamic_build": [],
-        "//tensorflow:using_cuda_nvcc_with_dynamic_build": [],
-        "//tensorflow:using_rocm_hipcc": ["//tensorflow/stream_executor/rocm:all_runtime"],
+        "//tensorflow:using_cuda_clang": [":_cuda_runtime"],
+        "//tensorflow:using_cuda_nvcc": [":_cuda_runtime"],
+        "//tensorflow:using_rocm_hipcc": [
+            "//tensorflow/stream_executor/rocm:all_runtime",
+        ],
         "//conditions:default": [],
     }),
 )

third_party/gpus/cuda/BUILD.tpl

@@ -1,31 +1,48 @@
 load(":build_defs.bzl", "cuda_header_library")
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("@bazel_skylib//lib:selects.bzl", "selects")
 
 licenses(["restricted"])  # MPL2, portions GPL v3, LGPL v3, BSD-like
 
 package(default_visibility = ["//visibility:public"])
 
-config_setting(
-    name = "using_nvcc",
-    values = {
-        "define": "using_cuda_nvcc=true",
-    },
+# Config setting whether TensorFlow is built with CUDA support using clang.
+#
+# TODO(b/174244321), DEPRECATED: this target will be removed when all users
+# have been converted to :is_cuda_enabled (most) or :is_cuda_compiler_clang.
+selects.config_setting_group(
+    name = "using_clang",
+    match_all = [
+        "@local_config_cuda//:is_cuda_enabled",
+        "@local_config_cuda//:is_cuda_compiler_clang",
+    ],
 )
 
-config_setting(
-    name = "using_clang",
-    values = {
-        "define": "using_cuda_clang=true",
-    },
+# Config setting whether TensorFlow is built with CUDA support using nvcc.
+#
+# TODO(b/174244321), DEPRECATED: this target will be removed when all users
+# have been converted to :is_cuda_enabled (most) or :is_cuda_compiler_nvcc.
+selects.config_setting_group(
+    name = "using_nvcc",
+    match_all = [
+        "@local_config_cuda//:is_cuda_enabled",
+        "@local_config_cuda//:is_cuda_compiler_nvcc",
+    ],
 )
 
 # Equivalent to using_clang && -c opt.
-config_setting(
+selects.config_setting_group(
     name = "using_clang_opt",
-    values = {
-        "define": "using_cuda_clang=true",
-        "compilation_mode": "opt",
-    },
+    match_all = [
+        ":using_clang",
+        ":_opt",
+    ],
+)
+
+config_setting(
+    name = "_opt",
+    values = {"compilation_mode": "opt"},
+    visibility = ["//visibility:private"],
 )
 
 # Provides CUDA headers for '#include "third_party/gpus/cuda/include/cuda.h"'

third_party/gpus/cuda/BUILD.windows.tpl

@@ -1,31 +1,48 @@
 load(":build_defs.bzl", "cuda_header_library")
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("@bazel_skylib//lib:selects.bzl", "selects")
 
 licenses(["restricted"])  # MPL2, portions GPL v3, LGPL v3, BSD-like
 
 package(default_visibility = ["//visibility:public"])
 
-config_setting(
-    name = "using_nvcc",
-    values = {
-        "define": "using_cuda_nvcc=true",
-    },
+# Config setting whether TensorFlow is built with CUDA support using clang.
+#
+# TODO(b/174244321), DEPRECATED: this target will be removed when all users
+# have been converted to :is_cuda_enabled (most) or :is_cuda_compiler_clang.
+selects.config_setting_group(
+    name = "using_clang",
+    match_all = [
+        "@local_config_cuda//:is_cuda_enabled",
+        "@local_config_cuda//:is_cuda_compiler_clang",
+    ],
 )
 
-config_setting(
-    name = "using_clang",
-    values = {
-        "define": "using_cuda_clang=true",
-    },
+# Config setting whether TensorFlow is built with CUDA support using nvcc.
+#
+# TODO(b/174244321), DEPRECATED: this target will be removed when all users
+# have been converted to :is_cuda_enabled (most) or :is_cuda_compiler_nvcc.
+selects.config_setting_group(
+    name = "using_nvcc",
+    match_all = [
+        "@local_config_cuda//:is_cuda_enabled",
+        "@local_config_cuda//:is_cuda_compiler_nvcc",
+    ],
 )
 
 # Equivalent to using_clang && -c opt.
-config_setting(
+selects.config_setting_group(
     name = "using_clang_opt",
-    values = {
-        "define": "using_cuda_clang=true",
-        "compilation_mode": "opt",
-    },
+    match_all = [
+        ":using_clang",
+        ":_opt",
+    ],
+)
+
+config_setting(
+    name = "_opt",
+    values = {"compilation_mode": "opt"},
+    visibility = ["//visibility:private"],
 )
 
 # Provides CUDA headers for '#include "third_party/gpus/cuda/include/cuda.h"'

third_party/gpus/cuda_configure.bzl

@@ -1378,6 +1378,8 @@ def _create_remote_cuda_repository(repository_ctx, remote_config_repo):
 
 def _cuda_autoconf_impl(repository_ctx):
     """Implementation of the cuda_autoconf repository rule."""
+    build_file = Label("//third_party/gpus:local_config_cuda.BUILD")
+
     if not enable_cuda(repository_ctx):
         _create_dummy_repository(repository_ctx)
     elif get_host_environ(repository_ctx, _TF_CUDA_CONFIG_REPO) != None:
@@ -1393,6 +1395,8 @@ def _cuda_autoconf_impl(repository_ctx):
     else:
         _create_local_cuda_repository(repository_ctx)
 
+    repository_ctx.symlink(build_file, "BUILD")
+
 # For @bazel_tools//tools/cpp:windows_cc_configure.bzl
 _MSVC_ENVVARS = [
     "BAZEL_VC",

third_party/gpus/local_config_cuda.BUILD

@@ -0,0 +1,50 @@
+load(
+    "@bazel_skylib//rules:common_settings.bzl",
+    "bool_flag",
+    "string_flag",
+)
+
+package(default_visibility = ["//visibility:public"])
+
+# Build flag to enable CUDA support.
+#
+# Enable with '--@local_config_cuda//:enable_cuda', or indirectly with
+# ./configure or '--config=cuda'.
+bool_flag(
+    name = "enable_cuda",
+    build_setting_default = False,
+)
+
+# Config setting whether CUDA support has been requested.
+#
+# Enable path: ./configure > --config=cuda (.tf_configure.bazelrc)
+#     > --//tensorflow:enable_cuda (.bazelrc) > :is_cuda_enabled
+config_setting(
+    name = "is_cuda_enabled",
+    flag_values = {":enable_cuda": "True"},
+)
+
+# Build flag to select CUDA compiler.
+#
+# Set with '--@local_config_cuda//:cuda_compiler=...', or indirectly with
+# ./configure, '--config=cuda' or '--config=cuda_clang'.
+string_flag(
+    name = "cuda_compiler",
+    build_setting_default = "nvcc",
+    values = [
+        "clang",
+        "nvcc",
+    ],
+)
+
+# Config setting whether CUDA device code should be compiled with clang.
+config_setting(
+    name = "is_cuda_compiler_clang",
+    flag_values = {":cuda_compiler": "clang"},
+)
+
+# Config setting whether CUDA device code should be compiled with nvcc.
+config_setting(
+    name = "is_cuda_compiler_nvcc",
+    flag_values = {":cuda_compiler": "nvcc"},
+)

third_party/mkl_dnn/mkldnn.BUILD

@@ -5,14 +5,6 @@ load(
     "template_rule",
 )
 
-config_setting(
-    name = "clang_linux_x86_64",
-    values = {
-        "cpu": "k8",
-        "define": "using_clang=true",
-    },
-)
-
 template_rule(
     name = "mkldnn_config_h",
     src = "include/mkldnn_config.h.in",