From 43bc7a79f0d2f576d0f3d719f3a55a10f8fba61c Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Tue, 19 May 2020 10:04:20 -0700
Subject: [PATCH] Simplify nccl_configure.

It can make use of the newly added cuda_gpu_architectures() macro.

PiperOrigin-RevId: 312301158
Change-Id: Ifa5229831f8d17093a0649b64457ae9d97ba6737
---
 third_party/nccl/build_defs.bzl.tpl |  4 ++--
 third_party/nccl/nccl_configure.bzl | 10 ----------
 2 files changed, 2 insertions(+), 12 deletions(-)

diff --git a/third_party/nccl/build_defs.bzl.tpl b/third_party/nccl/build_defs.bzl.tpl
index e734e49f9dc..7585949ea92 100644
--- a/third_party/nccl/build_defs.bzl.tpl
+++ b/third_party/nccl/build_defs.bzl.tpl
@@ -1,6 +1,6 @@
 """Repository rule for NCCL."""
 
-load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts")
+load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts", "cuda_gpu_architectures")
 load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")
 
 def _gen_device_srcs_impl(ctx):
@@ -285,7 +285,7 @@ def cuda_rdc_library(name, hdrs = None, copts = None, linkstatic = True, **kwarg
         name = dlink_hdrs,
         deps = [lib],
         out = dlink_cc,
-        gpu_archs = %{gpu_architectures},
+        gpu_archs = cuda_gpu_architectures(),
         nvlink_args = select({
             "@org_tensorflow//tensorflow:linux_x86_64": ["--cpu-arch=X86_64"],
             "@org_tensorflow//tensorflow:linux_ppc64le": ["--cpu-arch=PPC64LE"],
diff --git a/third_party/nccl/nccl_configure.bzl b/third_party/nccl/nccl_configure.bzl
index 92acb204097..d59e861d70b 100644
--- a/third_party/nccl/nccl_configure.bzl
+++ b/third_party/nccl/nccl_configure.bzl
@@ -13,7 +13,6 @@
 
 load(
     "//third_party/gpus:cuda_configure.bzl",
-    "compute_capabilities",
     "enable_cuda",
     "find_cuda_config",
 )
@@ -84,16 +83,7 @@ def _create_local_nccl_repository(repository_ctx):
         # Alias to open source build from @nccl_archive.
         repository_ctx.file("BUILD", _NCCL_ARCHIVE_BUILD_CONTENT)
 
-        # TODO(csigg): implement and reuse in cuda_configure.bzl.
-        gpu_architectures = [
-            "sm_" + capability.replace(".", "")
-            for capability in compute_capabilities(repository_ctx)
-        ]
-
-        # Round-about way to make the list unique.
-        gpu_architectures = dict(zip(gpu_architectures, gpu_architectures)).keys()
         config_wrap = {
-            "%{gpu_architectures}": str(gpu_architectures),
             "%{use_bin2c_path}": "False",
         }
         if (int(cuda_major), int(cuda_minor)) <= (10, 1):