Enabling cuda configuration to use a remotely generated configuration. This will prepare TF for use of Bazel's new remote execution api.

PiperOrigin-RevId: 159098871

3 files changed, 150 insertions, 7 deletions

diff --git a/third_party/gpus/crosstool/remote.BUILD.tpl b/third_party/gpus/crosstool/remote.BUILD.tpl
new file mode 100644
index 0000000000..b2316331db
--- /dev/null
+++ b/third_party/gpus/crosstool/remote.BUILD.tpl
@@ -0,0 +1,10 @@
+# Description:
+#   Template for crosstool Build file to use a pre-generated config.
+licenses(["restricted"])
+
+package(default_visibility = ["//visibility:public"])
+
+alias(
+    name = "toolchain",
+    actual = "%{remote_cuda_repo}:toolchain",
+)
diff --git a/third_party/gpus/cuda/remote.BUILD.tpl b/third_party/gpus/cuda/remote.BUILD.tpl
new file mode 100644
index 0000000000..d88d512b90
--- /dev/null
+++ b/third_party/gpus/cuda/remote.BUILD.tpl
@@ -0,0 +1,105 @@
+# Description:
+#   Template for cuda Build file to use a pre-generated config.
+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(
+    name = "using_clang",
+    values = {
+        "define": "using_cuda_clang=true",
+    },
+)
+
+# Equivalent to using_clang && -c opt.
+config_setting(
+    name = "using_clang_opt",
+    values = {
+        "define": "using_cuda_clang=true",
+        "compilation_mode": "opt",
+    },
+)
+
+config_setting(
+    name = "darwin",
+    values = {"cpu": "darwin"},
+    visibility = ["//visibility:public"],
+)
+
+config_setting(
+    name = "freebsd",
+    values = {"cpu": "freebsd"},
+    visibility = ["//visibility:public"],
+)
+
+alias(
+    name = "cuda_headers",
+    actual = "%{remote_cuda_repo}cuda:cuda_headers",
+)
+
+alias(
+    name = "cudart_static",
+    actual = "%{remote_cuda_repo}cuda:cudart_static",
+)
+
+alias(
+    name = "cuda_driver",
+    actual = "%{remote_cuda_repo}cuda:cuda_driver",
+)
+
+alias(
+    name = "cudart",
+    actual = "%{remote_cuda_repo}cuda:cudart",
+)
+
+alias(
+    name = "cublas",
+    actual = "%{remote_cuda_repo}cuda:cublas",
+)
+
+alias(
+    name = "cusolver",
+    actual = "%{remote_cuda_repo}cuda:cusolver",
+)
+
+alias(
+    name = "cudnn",
+    actual = "%{remote_cuda_repo}cuda:cudnn",
+)
+
+alias(
+    name = "cufft",
+    actual = "%{remote_cuda_repo}cuda:cufft",
+)
+
+alias(
+    name = "curand",
+    actual = "%{remote_cuda_repo}cuda:curand",
+)
+
+alias(
+    name = "cuda",
+    actual = "%{remote_cuda_repo}cuda:cuda",
+)
+
+alias(
+    name = "cupti_headers",
+    actual = "%{remote_cuda_repo}cuda:cupti_headers",
+)
+
+alias(
+    name = "cupti_dsos",
+    actual = "%{remote_cuda_repo}cuda:cupti_dsos",
+)
+
+alias(
+    name = "libdevice_root",
+    actual = "%{remote_cuda_repo}cuda:libdevice_root",
+)
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 61932a8e6d..77dc602fd9 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -883,15 +883,16 @@ def _use_cuda_clang(repository_ctx):
     return enable_cuda == "1"
   return False
 
-def _compute_cuda_extra_copts(repository_ctx, cuda_config):
+def _compute_cuda_extra_copts(repository_ctx, compute_capabilities):
   if _use_cuda_clang(repository_ctx):
-    capability_flags = ["--cuda-gpu-arch=sm_" + cap.replace(".", "") for cap in cuda_config.compute_capabilities]
+    capability_flags = ["--cuda-gpu-arch=sm_" +
+        cap.replace(".", "") for cap in compute_capabilities]
   else:
     # Capabilities are handled in the "crosstool_wrapper_driver_is_not_gcc" for nvcc
     capability_flags = []
   return str(capability_flags)
 
-def _create_cuda_repository(repository_ctx):
+def _create_local_cuda_repository(repository_ctx):
   """Creates the repository containing files set up to build with CUDA."""
   cuda_config = _get_cuda_config(repository_ctx)
 
@@ -939,7 +940,8 @@ def _create_cuda_repository(repository_ctx):
   _tpl(repository_ctx, "cuda:build_defs.bzl",
        {
            "%{cuda_is_configured}": "True",
-           "%{cuda_extra_copts}": _compute_cuda_extra_copts(repository_ctx, cuda_config),
+           "%{cuda_extra_copts}": _compute_cuda_extra_copts(
+               repository_ctx, cuda_config.compute_capabilities),
 
        })
   _tpl(repository_ctx, "cuda:BUILD",
@@ -1000,17 +1002,39 @@ def _create_cuda_repository(repository_ctx):
        })
 
 
+def _create_remote_cuda_repository(repository_ctx):
+  """Creates pointers to a remotely configured repo set up to build with CUDA."""
+  _tpl(repository_ctx, "cuda:build_defs.bzl",
+       {
+           "%{cuda_is_configured}": "True",
+           "%{cuda_extra_copts}": _compute_cuda_extra_copts(
+               repository_ctx, _compute_capabilities(repository_ctx)),
+
+       })
+  _tpl(repository_ctx, "cuda:remote.BUILD",
+       {
+           "%{remote_cuda_repo}": repository_ctx.attr.remote_config_repo,
+       }, "cuda/BUILD")
+  _tpl(repository_ctx, "crosstool:remote.BUILD", {
+           "%{remote_cuda_repo}": repository_ctx.attr.remote_config_repo,
+       }, "crosstool/BUILD")
+
 def _cuda_autoconf_impl(repository_ctx):
   """Implementation of the cuda_autoconf repository rule."""
   if not _enable_cuda(repository_ctx):
     _create_dummy_repository(repository_ctx)
   else:
-    _create_cuda_repository(repository_ctx)
-
+    if repository_ctx.attr.remote_config_repo != "":
+      _create_remote_cuda_repository(repository_ctx)
+    else:
+      _create_local_cuda_repository(repository_ctx)
 
 
 cuda_configure = repository_rule(
     implementation = _cuda_autoconf_impl,
+    attrs = {
+        "remote_config_repo": attr.string(mandatory = False, default =""),
+    },
     environ = [
         _GCC_HOST_COMPILER_PATH,
         "TF_NEED_CUDA",
@@ -1027,9 +1051,13 @@ cuda_configure = repository_rule(
 Add the following to your WORKSPACE FILE:
 
 ```python
-cuda_configure(name = "local_config_cuda")
+cuda_configure(
+  name = "local_config_cuda"
+  remote_config_repo = "@remote_cuda_config_tf//"
+)
 ```
 
 Args:
   name: A unique name for this workspace rule.
+  remote_config_repo: Location of a pre-generated config (optional).
 """