From 67edc16326d6328e7ef096e1b06f81dae1bfb816 Mon Sep 17 00:00:00 2001
From: Sami <skama@nvidia.com>
Date: Thu, 5 Dec 2019 16:49:20 -0800
Subject: [PATCH] Make nccl bindings compilable with cuda 10.2

---
 third_party/nccl/build_defs.bzl.tpl | 16 +++++++++-------
 third_party/nccl/nccl_configure.bzl | 18 ++++++++++++++++--
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/third_party/nccl/build_defs.bzl.tpl b/third_party/nccl/build_defs.bzl.tpl
index 57191398553..e734e49f9dc 100644
--- a/third_party/nccl/build_defs.bzl.tpl
+++ b/third_party/nccl/build_defs.bzl.tpl
@@ -104,19 +104,21 @@ def _device_link_impl(ctx):
     tmp_fatbin = ctx.actions.declare_file("%s.fatbin" % name)
     fatbin_h = ctx.actions.declare_file("%s_fatbin.h" % name)
     bin2c = ctx.file._bin2c
-    ctx.actions.run(
-        outputs = [tmp_fatbin, fatbin_h],
-        inputs = cubins,
-        executable = ctx.file._fatbinary,
-        arguments = [
+    arguments_list = [
             "-64",
             "--cmdline=--compile-only",
             "--link",
             "--compress-all",
-            "--bin2c-path=%s" % bin2c.dirname,
             "--create=%s" % tmp_fatbin.path,
             "--embedded-fatbin=%s" % fatbin_h.path,
-        ] + images,
+        ]
+    if %{use_bin2c_path}:
+           arguments_list.append("--bin2c-path=%s" % bin2c.dirname)
+    ctx.actions.run(
+        outputs = [tmp_fatbin, fatbin_h],
+        inputs = cubins,
+        executable = ctx.file._fatbinary,
+        arguments = arguments_list + images,
         tools = [bin2c],
         mnemonic = "fatbinary",
     )
diff --git a/third_party/nccl/nccl_configure.bzl b/third_party/nccl/nccl_configure.bzl
index 4b341d02d22..68e093c1e7b 100644
--- a/third_party/nccl/nccl_configure.bzl
+++ b/third_party/nccl/nccl_configure.bzl
@@ -72,6 +72,11 @@ def _nccl_configure_impl(repository_ctx):
         nccl_version = repository_ctx.os.environ[_TF_NCCL_VERSION].strip()
         nccl_version = nccl_version.split(".")[0]
 
+    cuda_config = find_cuda_config(repository_ctx, ["cuda"])
+    cuda_version = cuda_config["cuda_version"].split(".")
+    cuda_major = cuda_version[0]
+    cuda_minor = cuda_version[1]
+
     if nccl_version == "":
         # Alias to open source build from @nccl_archive.
         repository_ctx.file("BUILD", _NCCL_ARCHIVE_BUILD_CONTENT)
@@ -84,9 +89,18 @@ def _nccl_configure_impl(repository_ctx):
 
         # Round-about way to make the list unique.
         gpu_architectures = dict(zip(gpu_architectures, gpu_architectures)).keys()
-        repository_ctx.template("build_defs.bzl", _label("build_defs.bzl.tpl"), {
+        config_wrap = {
             "%{gpu_architectures}": str(gpu_architectures),
-        })
+            "%{use_bin2c_path}": "False",
+        }
+        if (int(cuda_major), int(cuda_minor)) <= (10, 1):
+            config_wrap["%{use_bin2c_path}"] = "True"
+
+        repository_ctx.template(
+            "build_defs.bzl",
+            _label("build_defs.bzl.tpl"),
+            config_wrap,
+        )
     else:
         # Create target for locally installed NCCL.
         config = find_cuda_config(repository_ctx, ["nccl"])