diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu16.04 b/tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu16.04
index e0533d871ce..32100d63768 100644
--- a/tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu16.04
+++ b/tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu16.04
@@ -16,7 +16,7 @@ RUN sh -c "echo deb [arch=amd64] $DEB_ROCM_REPO xenial main > /etc/apt/sources.
RUN apt-get update --allow-insecure-repositories && \
DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated \
rocm-dev rocm-libs rocm-utils rocm-cmake \
- rocfft miopen-hip miopengemm rocblas hipblas rocrand \
+ rocfft miopen-hip miopengemm rocblas hipblas rocrand rccl \
rocm-profiler cxlactivitylogger && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
diff --git a/tensorflow/tools/ci_build/Dockerfile.rocm b/tensorflow/tools/ci_build/Dockerfile.rocm
index 0affd6f2b49..999b108e1ba 100644
--- a/tensorflow/tools/ci_build/Dockerfile.rocm
+++ b/tensorflow/tools/ci_build/Dockerfile.rocm
@@ -58,7 +58,7 @@ RUN apt-get update --allow-insecure-repositories && DEBIAN_FRONTEND=noninteracti
RUN apt-get update --allow-insecure-repositories && \
DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated \
rocm-dev rocm-libs rocm-utils rocm-cmake \
- rocfft miopen-hip miopengemm rocblas hipblas rocrand \
+ rocfft miopen-hip miopengemm rocblas hipblas rocrand rccl \
rocm-profiler cxlactivitylogger && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
diff --git a/third_party/toolchains/preconfig/generate/containers.bzl b/third_party/toolchains/preconfig/generate/containers.bzl
index e2e125b6970..411e0b25536 100644
--- a/third_party/toolchains/preconfig/generate/containers.bzl
+++ b/third_party/toolchains/preconfig/generate/containers.bzl
@@ -8,5 +8,5 @@ container_digests = {
"cuda10.0-cudnn7-centos6": "sha256:a1909ba09c703340ee0074ce63dd94fe8fea48035a25264677907a609e2375e0",
"cuda10.1-cudnn7-centos6": "sha256:454b899657e87893ee5e68dc0f87df59b6a0a7418ae09cafcc3dd65ac71feca9",
"cuda10.0-cudnn7-ubuntu16.04-manylinux2010": "sha256:76cdd3956ce714bedca4b0c5b34c08e77fda7e888b8814da973d95f45628761c",
- "rocm-ubuntu16.04": "sha256:2df35a0b7f7513b4ca820a12792e98ecafafabd1076300ef26f89386277c10cc",
+ "rocm-ubuntu16.04": "sha256:ed6c04a1ec56eec6ef6276a55d69042261d5b1792845af3be06412c917d0ed51",
}
diff --git a/third_party/toolchains/preconfig/ubuntu16.04/gcc5-rocm/BUILD b/third_party/toolchains/preconfig/ubuntu16.04/gcc5-rocm/BUILD
index 50b746ede37..28b9ef6dfe6 100755
--- a/third_party/toolchains/preconfig/ubuntu16.04/gcc5-rocm/BUILD
+++ b/third_party/toolchains/preconfig/ubuntu16.04/gcc5-rocm/BUILD
@@ -77,6 +77,7 @@ cc_toolchain_config(
"/opt/rocm/rocfft/include",
"/opt/rocm/rocblas/include",
"/opt/rocm/miopen/include",
+ "/opt/rocm/rccl/include",
"/opt/rocm/hcc/include",
"/opt/rocm/hcc/compiler/lib/clang/7.0.0/include/",
"/opt/rocm/hcc/lib/clang/7.0.0/include",
@@ -137,6 +138,7 @@ cc_toolchain_config(
"/opt/rocm/rocfft/include",
"/opt/rocm/rocblas/include",
"/opt/rocm/miopen/include",
+ "/opt/rocm/rccl/include",
"/opt/rocm/hcc/include",
"/opt/rocm/hcc/compiler/lib/clang/7.0.0/include/",
"/opt/rocm/hcc/lib/clang/7.0.0/include",
@@ -197,6 +199,7 @@ cc_toolchain_config(
"/opt/rocm/rocfft/include",
"/opt/rocm/rocblas/include",
"/opt/rocm/miopen/include",
+ "/opt/rocm/rccl/include",
"/opt/rocm/hcc/include",
"/opt/rocm/hcc/compiler/lib/clang/7.0.0/include/",
"/opt/rocm/hcc/lib/clang/7.0.0/include",
diff --git a/third_party/toolchains/preconfig/ubuntu16.04/rocm/rocm/BUILD b/third_party/toolchains/preconfig/ubuntu16.04/rocm/rocm/BUILD
index 5cd99b42b0f..ea3f7efc6b3 100755
--- a/third_party/toolchains/preconfig/ubuntu16.04/rocm/rocm/BUILD
+++ b/third_party/toolchains/preconfig/ubuntu16.04/rocm/rocm/BUILD
@@ -14,6 +14,7 @@ cc_library(
hdrs = [
"rocm/rocm_config.h",
":miopen-include",
+ ":rccl-include",
":rocblas-include",
":rocfft-include",
":rocm-include",
@@ -87,6 +88,18 @@ cc_library(
visibility = ["//visibility:public"],
)
+cc_library(
+ name = "rccl",
+ srcs = ["rocm/lib/librccl.so"],
+ data = ["rocm/lib/librccl.so"],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
cc_library(
name = "rocm",
visibility = ["//visibility:public"],
@@ -218,13 +231,13 @@ genrule(
"rocm/include/hip/hcc_detail/hip_fp16_math_fwd.h",
"rocm/include/hip/hcc_detail/hip_ldg.h",
"rocm/include/hip/hcc_detail/hip_memory.h",
- "rocm/include/hip/hcc_detail/hip_prof_api.h",
"rocm/include/hip/hcc_detail/hip_prof_str.h",
"rocm/include/hip/hcc_detail/hip_runtime.h",
"rocm/include/hip/hcc_detail/hip_runtime_api.h",
"rocm/include/hip/hcc_detail/hip_surface_types.h",
"rocm/include/hip/hcc_detail/hip_texture_types.h",
"rocm/include/hip/hcc_detail/hip_vector_types.h",
+ "rocm/include/hip/hcc_detail/hiprtc.h",
"rocm/include/hip/hcc_detail/host_defines.h",
"rocm/include/hip/hcc_detail/hsa_helpers.hpp",
"rocm/include/hip/hcc_detail/llvm_intrinsics.h",
@@ -244,6 +257,7 @@ genrule(
"rocm/include/hip/hip_runtime_api.h",
"rocm/include/hip/hip_texture_types.h",
"rocm/include/hip/hip_vector_types.h",
+ "rocm/include/hip/hiprtc.h",
"rocm/include/hip/math_functions.h",
"rocm/include/hip/nvcc_detail/channel_descriptor.h",
"rocm/include/hip/nvcc_detail/hip_complex.h",
@@ -288,6 +302,7 @@ genrule(
"rocm/include/miopen/export.h",
"rocm/include/miopen/miopen.h",
"rocm/include/miopen/version.h",
+ "rocm/include/miopen_kernel_includes.h",
"rocm/include/miopen_kernels.h",
"rocm/include/miopengemm/accuracytests.hpp",
"rocm/include/miopengemm/alphagenerator.hpp",
@@ -340,12 +355,14 @@ genrule(
"rocm/include/opencl1.2-c.pch",
"rocm/include/opencl2.0-c.pch",
"rocm/include/profiler/CXLActivityLogger/CXLActivityLogger.h",
+ "rocm/include/rccl.h",
"rocm/include/rocblas-auxiliary.h",
"rocm/include/rocblas-export.h",
"rocm/include/rocblas-functions.h",
"rocm/include/rocblas-types.h",
"rocm/include/rocblas-version.h",
"rocm/include/rocblas.h",
+ "rocm/include/rocblas_bfloat16.h",
"rocm/include/rocfft-export.h",
"rocm/include/rocfft-version.h",
"rocm/include/rocfft.h",
@@ -394,6 +411,7 @@ genrule(
"rocm/include/rocblas/rocblas-types.h",
"rocm/include/rocblas/rocblas-version.h",
"rocm/include/rocblas/rocblas.h",
+ "rocm/include/rocblas/rocblas_bfloat16.h",
],
cmd = """cp -rLf "/opt/rocm/rocblas/include/." "$(@D)/rocm/include/rocblas/" """,
)
@@ -405,11 +423,20 @@ genrule(
"rocm/include/miopen/miopen/export.h",
"rocm/include/miopen/miopen/miopen.h",
"rocm/include/miopen/miopen/version.h",
+ "rocm/include/miopen/miopen_kernel_includes.h",
"rocm/include/miopen/miopen_kernels.h",
],
cmd = """cp -rLf "/opt/rocm/miopen/include/." "$(@D)/rocm/include/miopen/" """,
)
+genrule(
+ name = "rccl-include",
+ outs = [
+ "rocm/include/rccl/rccl.h",
+ ],
+ cmd = """cp -rLf "/opt/rocm/rccl/include/." "$(@D)/" """,
+)
+
genrule(
name = "rocm-lib",
outs = [
@@ -418,10 +445,12 @@ genrule(
"rocm/lib/librocfft.so",
"rocm/lib/libhiprand.so",
"rocm/lib/libMIOpen.so",
+ "rocm/lib/librccl.so",
],
cmd = """cp -f "/opt/rocm/hip/lib/libhip_hcc.so" "$(location rocm/lib/libhip_hcc.so)" && \
-cp -f "/opt/rocm/rocblas/lib/librocblas.so.2.2.6.0" "$(location rocm/lib/librocblas.so)" && \
-cp -f "/opt/rocm/rocfft/lib/librocfft.so.0.9.3.0" "$(location rocm/lib/librocfft.so)" && \
-cp -f "/opt/rocm/hiprand/lib/libhiprand.so.1.8.2" "$(location rocm/lib/libhiprand.so)" && \
-cp -f "/opt/rocm/miopen/lib/libMIOpen.so.1" "$(location rocm/lib/libMIOpen.so)" """,
+cp -f "/opt/rocm/rocblas/lib/librocblas.so.2.2.11.0" "$(location rocm/lib/librocblas.so)" && \
+cp -f "/opt/rocm/rocfft/lib/librocfft.so.0.9.4.0" "$(location rocm/lib/librocfft.so)" && \
+cp -f "/opt/rocm/hiprand/lib/libhiprand.so.2.6.0" "$(location rocm/lib/libhiprand.so)" && \
+cp -f "/opt/rocm/miopen/lib/libMIOpen.so.1" "$(location rocm/lib/libMIOpen.so)" && \
+cp -f "/opt/rocm/rccl/lib/librccl.so" "$(location rocm/lib/librccl.so)" """,
)