experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Simplify nccl_configure. 

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

PiperOrigin-RevId: 312301158
Change-Id: Ifa5229831f8d17093a0649b64457ae9d97ba6737
This commit is contained in:
Adrian Kuegel 2020-05-19 10:04:20 -07:00 committed by TensorFlower Gardener
parent 81065dbaba
commit 43bc7a79f0
2 changed files with 2 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
third_party/nccl/build_defs.bzl.tpl vendored
View File

@ -1,6 +1,6 @@
"""Repository rule for NCCL.""" """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") load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")
def _gen_device_srcs_impl(ctx): 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, name = dlink_hdrs,
deps = [lib], deps = [lib],
out = dlink_cc, out = dlink_cc,
gpu_archs = %{gpu_architectures}, gpu_archs = cuda_gpu_architectures(),
nvlink_args = select({ nvlink_args = select({
"@org_tensorflow//tensorflow:linux_x86_64": ["--cpu-arch=X86_64"], "@org_tensorflow//tensorflow:linux_x86_64": ["--cpu-arch=X86_64"],
"@org_tensorflow//tensorflow:linux_ppc64le": ["--cpu-arch=PPC64LE"], "@org_tensorflow//tensorflow:linux_ppc64le": ["--cpu-arch=PPC64LE"],

10
third_party/nccl/nccl_configure.bzl vendored
View File

@ -13,7 +13,6 @@
load( load(
"//third_party/gpus:cuda_configure.bzl", "//third_party/gpus:cuda_configure.bzl",
"compute_capabilities",
"enable_cuda", "enable_cuda",
"find_cuda_config", "find_cuda_config",
) )
@ -84,16 +83,7 @@ def _create_local_nccl_repository(repository_ctx):
# Alias to open source build from @nccl_archive. # Alias to open source build from @nccl_archive.
repository_ctx.file("BUILD", _NCCL_ARCHIVE_BUILD_CONTENT) 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 = { config_wrap = {
"%{gpu_architectures}": str(gpu_architectures),
"%{use_bin2c_path}": "False", "%{use_bin2c_path}": "False",
} }
if (int(cuda_major), int(cuda_minor)) <= (10, 1): if (int(cuda_major), int(cuda_minor)) <= (10, 1):