rccl.cmake 899 字节
Newer Older
1
if(NOT WITH_ROCM)
2
  return()
3 4 5 6
endif()

# Now we don't support RCCL on windows
if(WIN32)
7
  return()
8 9 10
endif()

if(WITH_RCCL)
11 12 13 14 15 16 17 18
  set(RCCL_ROOT
      ${ROCM_PATH}/rccl
      CACHE PATH "RCCL ROOT")
  find_path(
    RCCL_INCLUDE_DIR rccl.h
    PATHS ${RCCL_ROOT} ${RCCL_ROOT}/include ${RCCL_ROOT}/local/include
          $ENV{RCCL_ROOT} $ENV{RCCL_ROOT}/include $ENV{RCCL_ROOT}/local/include
    NO_DEFAULT_PATH)
19

20
  file(READ ${RCCL_INCLUDE_DIR}/rccl.h RCCL_VERSION_FILE_CONTENTS)
21

22 23 24 25
  string(REGEX MATCH "define NCCL_VERSION_CODE +([0-9]+)" RCCL_VERSION
               "${RCCL_VERSION_FILE_CONTENTS}")
  string(REGEX REPLACE "define NCCL_VERSION_CODE +([0-9]+)" "\\1" RCCL_VERSION
                       "${RCCL_VERSION}")
26

27 28 29
  # 2604 for ROCM3.5 and 2708 for ROCM 3.9
  message(STATUS "Current RCCL header is ${RCCL_INCLUDE_DIR}/rccl.h. "
                 "Current RCCL version is v${RCCL_VERSION}. ")
30
endif()