diff --git a/cmake/Modules/FindNCCL.cmake b/cmake/Modules/FindNCCL.cmake index a16c9aca67..51eb54f14c 100644 --- a/cmake/Modules/FindNCCL.cmake +++ b/cmake/Modules/FindNCCL.cmake @@ -55,9 +55,10 @@ if(NCCL_FOUND) # obtaining NCCL version and some sanity checks set (OLD_CMAKE_REQUIRED_INCLUDES ${CMAKE_REQUIRED_INCLUDES}) list (APPEND CMAKE_REQUIRED_INCLUDES ${NCCL_INCLUDE_DIRS}) include(CheckCXXSymbolExists) - check_cxx_symbol_exists(NCCL_VERSION_CODE nccl.h NCCL_VERSION_DEFINED) + set(NCCL_VERSION_CODE $ENV{NCCL_VER_CODE}) + set(NCCL_VERSION_DEFINED $ENV{NCCL_VER_CODE}) - if (NCCL_VERSION_DEFINED) + if (DEFINED NCCL_VERSION_DEFINED) set(file "${PROJECT_BINARY_DIR}/detect_nccl_version.cc") file(WRITE ${file} " #include @@ -72,6 +73,7 @@ if(NCCL_FOUND) # obtaining NCCL version and some sanity checks } ") try_run(NCCL_VERSION_MATCHED compile_result ${PROJECT_BINARY_DIR} ${file} + CMAKE_FLAGS -DINCLUDE_DIRECTORIES=/opt/cuda/include RUN_OUTPUT_VARIABLE NCCL_VERSION_FROM_HEADER CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${NCCL_INCLUDE_DIRS}" LINK_LIBRARIES ${NCCL_LIBRARIES})