You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
56 lines
2.1 KiB
56 lines
2.1 KiB
if(NOT WITH_GPU)
|
|
return()
|
|
endif()
|
|
|
|
# Now we don't support NCCL on windows
|
|
if(WIN32)
|
|
return()
|
|
endif()
|
|
|
|
if(WITH_NCCL)
|
|
set(NCCL_ROOT "/usr" CACHE PATH "NCCL ROOT")
|
|
find_path(NCCL_INCLUDE_DIR nccl.h
|
|
PATHS ${NCCL_ROOT} ${NCCL_ROOT}/include ${NCCL_ROOT}/local/include
|
|
$ENV{NCCL_ROOT} $ENV{NCCL_ROOT}/include $ENV{NCCL_ROOT}/local/include
|
|
NO_DEFAULT_PATH
|
|
)
|
|
|
|
file(READ ${NCCL_INCLUDE_DIR}/nccl.h NCCL_VERSION_FILE_CONTENTS)
|
|
|
|
string(REGEX MATCH "define NCCL_VERSION_CODE +([0-9]+)"
|
|
NCCL_VERSION "${NCCL_VERSION_FILE_CONTENTS}")
|
|
string(REGEX REPLACE "define NCCL_VERSION_CODE +([0-9]+)" "\\1"
|
|
NCCL_VERSION "${NCCL_VERSION}")
|
|
|
|
if("${NCCL_VERSION}" GREATER "2000")
|
|
message(STATUS "Current NCCL header is ${NCCL_INCLUDE_DIR}/nccl.h. "
|
|
"Current NCCL version is v${NCCL_VERSION}. ")
|
|
else()
|
|
# in old version nccl, it may not define NCCL_VERSION_CODE
|
|
string(REGEX MATCH "define NCCL_MAJOR +([0-9]+)" NCCL_MAJOR_VERSION
|
|
"${NCCL_VERSION_FILE_CONTENTS}")
|
|
string(REGEX REPLACE "define NCCL_MAJOR +([0-9]+)" "\\1"
|
|
NCCL_MAJOR_VERSION "${NCCL_MAJOR_VERSION}")
|
|
string(REGEX MATCH "define NCCL_MINOR +([0-9]+)" NCCL_MINOR_VERSION
|
|
"${NCCL_VERSION_FILE_CONTENTS}")
|
|
string(REGEX REPLACE "define NCCL_MINOR +([0-9]+)" "\\1"
|
|
NCCL_MINOR_VERSION "${NCCL_MINOR_VERSION}")
|
|
string(REGEX MATCH "define NCCL_PATCH +([0-9]+)"
|
|
NCCL_PATCH_VERSION "${NCCL_VERSION_FILE_CONTENTS}")
|
|
string(REGEX REPLACE "define NCCL_PATCH +([0-9]+)" "\\1"
|
|
NCCL_PATCH_VERSION "${NCCL_PATCH_VERSION}")
|
|
|
|
if(NOT NCCL_MAJOR_VERSION)
|
|
set(NCCL_VERSION "0")
|
|
else()
|
|
math(EXPR NCCL_VERSION
|
|
"${NCCL_MAJOR_VERSION} * 1000 +
|
|
${NCCL_MINOR_VERSION} * 100 + ${NCCL_PATCH_VERSION}")
|
|
endif()
|
|
add_definitions("-DNCCL_VERSION_CODE=$NCCL_VERSION")
|
|
|
|
message(STATUS "Current NCCL header is ${NCCL_INCLUDE_DIR}/nccl.h. "
|
|
"Current NCCL version is v${NCCL_MAJOR_VERSION}.${NCCL_MINOR_VERSION}.${NCCL_PATCH_VERSION} ")
|
|
endif()
|
|
endif()
|