diff --git a/third_party/nccl/build_defs.bzl.tpl b/third_party/nccl/build_defs.bzl.tpl index bb6518753be..bd65cb98f73 100644 --- a/third_party/nccl/build_defs.bzl.tpl +++ b/third_party/nccl/build_defs.bzl.tpl @@ -182,8 +182,9 @@ def device_link(name, srcs): # Generate fatbin header from all cubins. fatbin_hdr = name + ".fatbin.h" fatbinary = "@local_config_nccl//:cuda/bin/fatbinary" - cmd = ("PATH=$$CUDA_TOOLKIT_PATH/bin:$$PATH " + # for bin2c - "$(location %s) -64 --cmdline=--compile-only --link " % fatbinary + + bin2c = "@local_config_nccl//:cuda/bin/bin2c" + cmd = ("$(location %s) -64 --cmdline=--compile-only " % fatbinary + + "--link --bin2c-path $$(dirname $(location %s)) " % bin2c + "--compress-all %s --create=%%{name}.fatbin " % " ".join(images) + "--embedded-fatbin=$@") native.genrule( @@ -191,7 +192,7 @@ def device_link(name, srcs): outs = [fatbin_hdr], srcs = cubins, cmd = cmd, - tools = [fatbinary], + tools = [fatbinary, bin2c], ) # Generate the source file #including the headers generated above. diff --git a/third_party/nccl/nccl_configure.bzl b/third_party/nccl/nccl_configure.bzl index 7f00df09620..1e6422b49ef 100644 --- a/third_party/nccl/nccl_configure.bzl +++ b/third_party/nccl/nccl_configure.bzl @@ -3,7 +3,7 @@ `nccl_configure` depends on the following environment variables: - * `TF_NCCL_VERSION`: The NCCL version. + * `TF_NCCL_VERSION`: Installed NCCL version or empty to build from source. * `NCCL_INSTALL_PATH`: The installation path of the NCCL library. * `NCCL_HDR_PATH`: The installation path of the NCCL header files. """ @@ -44,6 +44,7 @@ _NCCL_ARCHIVE_BUILD_CONTENT = """ exports_files([ "cuda/bin/crt/link.stub", "cuda/bin/fatbinary", + "cuda/bin/bin2c", "nvlink", ])