[CI] Prune unused archs from libnccl (#8179)

* [CI] Prune unused archs from libnccl

* Put pruning logic in CI directory

* Don't use --color in grep

Jenkinsfile

@@ -114,10 +114,6 @@ def checkoutSrcs() {
   }
 }
 
-def GetCUDABuildContainerType(cuda_version) {
-  return (cuda_version == ref_cuda_ver) ? 'gpu_build_centos7' : 'gpu_build'
-}
-
 def ClangTidy() {
   node('linux && cpu_build') {
     unstash name: 'srcs'
@@ -208,7 +204,7 @@ def BuildCUDA(args) {
   node('linux && cpu_build') {
     unstash name: 'srcs'
     echo "Build with CUDA ${args.cuda_version}"
-    def container_type = GetCUDABuildContainerType(args.cuda_version)
+    def container_type = "gpu_build_centos7"
     def docker_binary = "docker"
     def docker_args = "--build-arg CUDA_VERSION_ARG=${args.cuda_version}"
     def arch_flag = ""
@@ -217,7 +213,8 @@ def BuildCUDA(args) {
     }
     def wheel_tag = "manylinux2014_x86_64"
     sh """
-    ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh -DUSE_CUDA=ON -DUSE_NCCL=ON -DOPEN_MP:BOOL=ON -DHIDE_CXX_SYMBOLS=ON ${arch_flag}
+    ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/prune_libnccl.sh
+    ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh -DUSE_CUDA=ON -DUSE_NCCL=ON -DUSE_OPENMP=ON -DHIDE_CXX_SYMBOLS=ON -DUSE_NCCL_LIB_PATH=ON -DNCCL_INCLUDE_DIR=/usr/include -DNCCL_LIBRARY=/workspace/libnccl_static.a ${arch_flag}
     ${dockerRun} ${container_type} ${docker_binary} ${docker_args} bash -c "cd python-package && rm -rf dist/* && python setup.py bdist_wheel --universal"
     ${dockerRun} ${container_type} ${docker_binary} ${docker_args} python tests/ci_build/rename_whl.py python-package/dist/*.whl ${commit_id} ${wheel_tag}
     """

tests/ci_build/Dockerfile.gpu_build_centos7

@@ -15,8 +15,8 @@ RUN \
     bash Miniconda3.sh -b -p /opt/python && \
     /opt/python/bin/python -m pip install awscli && \
     # CMake
-    wget -nv -nc https://cmake.org/files/v3.14/cmake-3.14.0-Linux-x86_64.sh --no-check-certificate && \
-    bash cmake-3.14.0-Linux-x86_64.sh --skip-license --prefix=/usr
+    wget -nv -nc https://cmake.org/files/v3.18/cmake-3.18.0-Linux-x86_64.sh --no-check-certificate && \
+    bash cmake-3.18.0-Linux-x86_64.sh --skip-license --prefix=/usr
 
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \

tests/ci_build/prune_libnccl.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+set -e
+
+rm -rf tmp_nccl
+
+mkdir tmp_nccl
+pushd tmp_nccl
+
+set -x
+
+cat << EOF > test.cu
+int main(void) { return 0; }
+EOF
+
+cat << EOF > CMakeLists.txt
+cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
+project(gencode_extractor CXX C)
+cmake_policy(SET CMP0104 NEW)
+set(CMAKE_CUDA_HOST_COMPILER \${CMAKE_CXX_COMPILER})
+enable_language(CUDA)
+include(../cmake/Utils.cmake)
+set(GEN_CODE "")
+format_gencode_flags("" GEN_CODE)
+add_library(test OBJECT test.cu)
+set_property(TARGET test PROPERTY CUDA_ARCHITECTURES \${CMAKE_CUDA_ARCHITECTURES})
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+EOF
+
+cmake . -GNinja
+gen_code=$(grep -o -- '--generate-code=\S*' compile_commands.json | paste -sd ' ')
+
+nvprune ${gen_code} /usr/lib64/libnccl_static.a -o ../libnccl_static.a
+
+popd
+rm -rf tmp_nccl
+
+set +x