Merge commit for internal changes

5 files changed, 97 insertions, 0 deletions

diff --git a/third_party/gpus/cuda/BUILD.tpl b/third_party/gpus/cuda/BUILD.tpl
index 20920f7fca..8754105619 100644
--- a/third_party/gpus/cuda/BUILD.tpl
+++ b/third_party/gpus/cuda/BUILD.tpl
@@ -58,6 +58,13 @@ cc_library(
 )
 
 cc_library(
+    name = "cuda_driver",
+    srcs = ["lib/%{cuda_driver_lib}"],
+    includes = ["include/"],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
     name = "cudart",
     srcs = ["lib/%{cudart_lib}"],
     data = ["lib/%{cudart_lib}"],
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 316850ef84..d58a32fb9d 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -410,6 +410,9 @@ def _find_cuda_lib(lib, repository_ctx, cpu_value, basedir, version="",
     path = repository_ctx.path("%s/lib64/%s" % (basedir, file_name))
     if path.exists:
       return struct(file_name=file_name, path=str(path.realpath))
+    path = repository_ctx.path("%s/lib64/stubs/%s" % (basedir, file_name))
+    if path.exists:
+      return struct(file_name=file_name, path=str(path.realpath))
     path = repository_ctx.path(
         "%s/lib/x86_64-linux-gnu/%s" % (basedir, file_name))
     if path.exists:
@@ -493,6 +496,7 @@ def _find_libs(repository_ctx, cuda_config):
   cudnn_ext = ".%s" % cudnn_version if cudnn_version else ""
   cpu_value = cuda_config.cpu_value
   return {
+      "cuda": _find_cuda_lib("cuda", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path),
       "cudart": _find_cuda_lib(
           "cudart", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
           cuda_config.cuda_version),
@@ -649,6 +653,7 @@ def _create_dummy_repository(repository_ctx):
        })
   _tpl(repository_ctx, "cuda:BUILD",
        {
+           "%{cuda_driver_lib}": _lib_name("cuda", cpu_value),
            "%{cudart_static_lib}": _lib_name("cudart_static", cpu_value,
                                              static=True),
            "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value),
@@ -661,6 +666,7 @@ def _create_dummy_repository(repository_ctx):
        })
   _tpl(repository_ctx, "cuda:BUILD",
        {
+           "%{cuda_driver_lib}": _lib_name("cuda", cpu_value),
            "%{cudart_static_lib}": _lib_name("cudart_static", cpu_value,
                                              static=True),
            "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value),
@@ -684,6 +690,7 @@ def _create_dummy_repository(repository_ctx):
   repository_ctx.file("cuda/include/cublas.h", "")
   repository_ctx.file("cuda/include/cudnn.h", "")
   repository_ctx.file("cuda/extras/CUPTI/include/cupti.h", "")
+  repository_ctx.file("cuda/lib/%s" % _lib_name("cuda", cpu_value))
   repository_ctx.file("cuda/lib/%s" % _lib_name("cudart", cpu_value))
   repository_ctx.file("cuda/lib/%s" % _lib_name("cudart_static", cpu_value))
   repository_ctx.file("cuda/lib/%s" % _lib_name("cublas", cpu_value))
@@ -757,6 +764,7 @@ def _create_cuda_repository(repository_ctx):
        })
   _tpl(repository_ctx, "cuda:BUILD",
        {
+           "%{cuda_driver_lib}": cuda_libs["cuda"].file_name,
            "%{cudart_static_lib}": cuda_libs["cudart_static"].file_name,
            "%{cudart_static_linkopt}": _cudart_static_linkopt(
                cuda_config.cpu_value),
diff --git a/third_party/highwayhash.BUILD b/third_party/highwayhash.BUILD
new file mode 100644
index 0000000000..1b8e40765e
--- /dev/null
+++ b/third_party/highwayhash.BUILD
@@ -0,0 +1,32 @@
+# Description:
+#   SipHash and HighwayHash: cryptographically-strong pseudorandom functions
+
+licenses(["notice"])  # Apache 2.0
+
+exports_files(["LICENSE"])
+
+cc_library(
+    name = "sip_hash",
+    srcs = ["highwayhash/sip_hash.cc"],
+    hdrs = [
+        "highwayhash/sip_hash.h",
+        "highwayhash/state_helpers.h",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":arch_specific",
+        ":compiler_specific",
+    ],
+)
+
+cc_library(
+    name = "arch_specific",
+    srcs = ["highwayhash/arch_specific.cc"],
+    hdrs = ["highwayhash/arch_specific.h"],
+    deps = [":compiler_specific"],
+)
+
+cc_library(
+    name = "compiler_specific",
+    hdrs = ["highwayhash/compiler_specific.h"],
+)
diff --git a/third_party/llvm/llvm.BUILD b/third_party/llvm/llvm.BUILD
index 330d8b79ce..02f13939a9 100644
--- a/third_party/llvm/llvm.BUILD
+++ b/third_party/llvm/llvm.BUILD
@@ -329,6 +329,7 @@ llvm_target_list = [
         "lower_name": "aarch64",
         "short_name": "AArch64",
         "tbl_outs": [
+            ("-gen-register-bank", "lib/Target/AArch64/AArch64GenRegisterBank.inc"),
             ("-gen-register-info", "lib/Target/AArch64/AArch64GenRegisterInfo.inc"),
             ("-gen-instr-info", "lib/Target/AArch64/AArch64GenInstrInfo.inc"),
             ("-gen-emitter", "lib/Target/AArch64/AArch64GenMCCodeEmitter.inc"),
@@ -1266,6 +1267,7 @@ cc_library(
         ":code_gen",
         ":config",
         ":core",
+        ":ipo",
         ":mc",
         ":nvptx_asm_printer",
         ":nvptx_desc",
diff --git a/third_party/nccl.BUILD b/third_party/nccl.BUILD
new file mode 100644
index 0000000000..bb460a05e0
--- /dev/null
+++ b/third_party/nccl.BUILD
@@ -0,0 +1,48 @@
+# NVIDIA nccl
+# A package of optimized primitives for collective multi-GPU communication.
+
+licenses(["notice"])  # BSD
+
+exports_files(["LICENSE.txt"])
+
+load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts", "if_cuda")
+
+SRCS = [
+    "src/all_gather.cu",
+    "src/all_reduce.cu",
+    "src/broadcast.cu",
+    "src/core.cu",
+    "src/libwrap.cu",
+    "src/reduce.cu",
+    "src/reduce_scatter.cu",
+]
+
+# Copy .cu to .cu.cc so they can be in srcs of cc_library.
+[
+    genrule(
+        name = "gen_" + src,
+        srcs = [src],
+        outs = [src + ".cc"],
+        cmd = "cp $(location " + src + ") $(location " + src + ".cc)",
+    )
+    for src in SRCS
+]
+
+SRCS_CU_CC = [src + ".cc" for src in SRCS]
+
+cc_library(
+    name = "nccl",
+    srcs = if_cuda(SRCS_CU_CC + glob(["src/*.h"])),
+    hdrs = if_cuda(["src/nccl.h"]),
+    copts = [
+        "-DCUDA_MAJOR=0",
+        "-DCUDA_MINOR=0",
+        "-DNCCL_MAJOR=0",
+        "-DNCCL_MINOR=0",
+        "-DNCCL_PATCH=0",
+        "-Iexternal/nccl_archive/src",
+        "-O3",
+    ] + cuda_default_copts(),
+    visibility = ["//visibility:public"],
+    deps = ["@local_config_cuda//cuda:cuda_headers"],
+)