cudnn.cmake 3.0 KB
Newer Older
1 2 3 4
if(NOT WITH_GPU)
    return()
endif()

5
set(CUDNN_ROOT "/usr" CACHE PATH "CUDNN ROOT")
Z
zhangjinchao01 已提交
6 7 8 9 10 11 12 13
find_path(CUDNN_INCLUDE_DIR cudnn.h
    PATHS ${CUDNN_ROOT} ${CUDNN_ROOT}/include
    $ENV{CUDNN_ROOT} $ENV{CUDNN_ROOT}/include ${CUDA_TOOLKIT_INCLUDE}
    NO_DEFAULT_PATH
)

get_filename_component(__libpath_hist ${CUDA_CUDART_LIBRARY} PATH)

14 15 16 17
set(TARGET_ARCH "x86_64")
if(NOT ${CMAKE_SYSTEM_PROCESSOR})
    set(TARGET_ARCH ${CMAKE_SYSTEM_PROCESSOR})
endif()
18

Z
zhangjinchao01 已提交
19 20 21 22
list(APPEND CUDNN_CHECK_LIBRARY_DIRS
    ${CUDNN_ROOT}
    ${CUDNN_ROOT}/lib64
    ${CUDNN_ROOT}/lib
23
    ${CUDNN_ROOT}/lib/${TARGET_ARCH}-linux-gnu
24
    ${CUDNN_ROOT}/local/cuda-${CUDA_VERSION}/targets/${TARGET_ARCH}-linux/lib/
Z
zhangjinchao01 已提交
25 26 27
    $ENV{CUDNN_ROOT}
    $ENV{CUDNN_ROOT}/lib64
    $ENV{CUDNN_ROOT}/lib
D
dzhwinter 已提交
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44
    /usr/lib
	${CUDA_TOOLKIT_ROOT_DIR}
	${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
	)
set(CUDNN_LIB_NAME "")
if (LINUX)
set(CUDNN_LIB_NAME "libcudnn.so")
endif(LINUX)

if(WIN32)
# only support cudnn7
set(CUDNN_LIB_NAME "cudnn.lib" "cudnn64_7.dll")
endif(WIN32)

if(Apple)
set(CUDNN_LIB_NAME "libcudnn.dylib" "libcudnn.so")
endif(Apple)
D
dzhwinter 已提交
45

D
dzhwinter 已提交
46
find_library(CUDNN_LIBRARY NAMES ${CUDNN_LIB_NAME} # libcudnn_static.a
Z
zhangjinchao01 已提交
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
    PATHS ${CUDNN_CHECK_LIBRARY_DIRS} ${CUDNN_INCLUDE_DIR} ${__libpath_hist}
          NO_DEFAULT_PATH
    DOC "Path to cuDNN library.")

if(CUDNN_INCLUDE_DIR AND CUDNN_LIBRARY)
    set(CUDNN_FOUND ON)
else()
    set(CUDNN_FOUND OFF)
endif()

if(CUDNN_FOUND)
    file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_VERSION_FILE_CONTENTS)

    get_filename_component(CUDNN_LIB_PATH ${CUDNN_LIBRARY} DIRECTORY)

    string(REGEX MATCH "define CUDNN_VERSION +([0-9]+)"
        CUDNN_VERSION "${CUDNN_VERSION_FILE_CONTENTS}")
    string(REGEX REPLACE "define CUDNN_VERSION +([0-9]+)" "\\1"
        CUDNN_VERSION "${CUDNN_VERSION}")

    if("${CUDNN_VERSION}" STREQUAL "2000")
        message(STATUS "Current cuDNN version is v2. ")
    else()
        string(REGEX MATCH "define CUDNN_MAJOR +([0-9]+)" CUDNN_MAJOR_VERSION
            "${CUDNN_VERSION_FILE_CONTENTS}")
        string(REGEX REPLACE "define CUDNN_MAJOR +([0-9]+)" "\\1"
            CUDNN_MAJOR_VERSION "${CUDNN_MAJOR_VERSION}")
        string(REGEX MATCH "define CUDNN_MINOR +([0-9]+)" CUDNN_MINOR_VERSION
            "${CUDNN_VERSION_FILE_CONTENTS}")
        string(REGEX REPLACE "define CUDNN_MINOR +([0-9]+)" "\\1"
            CUDNN_MINOR_VERSION "${CUDNN_MINOR_VERSION}")
        string(REGEX MATCH "define CUDNN_PATCHLEVEL +([0-9]+)"
            CUDNN_PATCHLEVEL_VERSION "${CUDNN_VERSION_FILE_CONTENTS}")
        string(REGEX REPLACE "define CUDNN_PATCHLEVEL +([0-9]+)" "\\1"
            CUDNN_PATCHLEVEL_VERSION "${CUDNN_PATCHLEVEL_VERSION}")

        if(NOT CUDNN_MAJOR_VERSION)
            set(CUDNN_VERSION "???")
D
dzhwinter 已提交
85
        else() 
Z
zhangjinchao01 已提交
86 87 88 89 90 91 92 93 94 95
            math(EXPR CUDNN_VERSION
                "${CUDNN_MAJOR_VERSION} * 1000 +
                 ${CUDNN_MINOR_VERSION} * 100 + ${CUDNN_PATCHLEVEL_VERSION}")
        endif()

        message(STATUS "Current cuDNN header is ${CUDNN_INCLUDE_DIR}/cudnn.h. "
            "Current cuDNN version is v${CUDNN_MAJOR_VERSION}. ")

    endif()
endif()