| # Find the nccl libraries |
| # |
| # The following variables are optionally searched for defaults |
| # NCCL_ROOT_DIR: Base directory where all NCCL components are found |
| # NCCL_INCLUDE_DIR: Directory where NCCL header is found |
| # NCCL_LIB_DIR: Directory where NCCL library is found |
| # |
| # The following are set after configuration is done: |
| # NCCL_FOUND |
| # NCCL_INCLUDE_DIRS |
| # NCCL_LIBRARIES |
| # |
| # The path hints include CUDA_TOOLKIT_ROOT_DIR seeing as some folks |
| # install NCCL in the same location as the CUDA toolkit. |
| # See https://github.com/caffe2/caffe2/issues/1601 |
| |
| set(NCCL_ROOT_DIR "" CACHE PATH "Folder contains NVIDIA NCCL") |
| |
| find_path(NCCL_INCLUDE_DIRS |
| NAMES nccl.h |
| HINTS |
| ${NCCL_INCLUDE_DIR} |
| ${NCCL_ROOT_DIR} |
| ${NCCL_ROOT_DIR}/include |
| ${CUDA_TOOLKIT_ROOT_DIR}/include) |
| |
| IF ($ENV{USE_STATIC_NCCL}) |
| MESSAGE(STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library") |
| SET(NCCL_LIBNAME "libnccl_static.a") |
| ELSE() |
| SET(NCCL_LIBNAME "nccl") |
| ENDIF() |
| |
| find_library(NCCL_LIBRARIES |
| NAMES ${NCCL_LIBNAME} |
| HINTS |
| ${NCCL_LIB_DIR} |
| ${NCCL_ROOT_DIR} |
| ${NCCL_ROOT_DIR}/lib |
| ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu |
| ${NCCL_ROOT_DIR}/lib64 |
| ${CUDA_TOOLKIT_ROOT_DIR}/lib64) |
| |
| include(FindPackageHandleStandardArgs) |
| find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS NCCL_LIBRARIES) |
| |
| if(NCCL_FOUND) |
| message(STATUS "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})") |
| mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES) |
| endif() |
