MegEngine/cmake/cudnn.cmake

find_package(PkgConfig)
if(${PkgConfig_FOUND})
    pkg_check_modules(PC_CUDNN QUIET CUDNN)
endif()

if("${CUDNN_ROOT_DIR}" STREQUAL "" AND NOT "$ENV{CUDNN_ROOT_DIR}"  STREQUAL "")
    set(CUDNN_ROOT_DIR $ENV{CUDNN_ROOT_DIR})
endif()

if(MGE_CUDA_USE_STATIC AND NOT MGE_WITH_CUDNN_SHARED)
    find_library(CUDNN_LIBRARY 
        NAMES libcudnn_static.a cudnn.lib
        PATHS ${ALTER_LD_LIBRARY_PATHS} ${CUDNN_ROOT_DIR} ${PC_CUDNN_LIBRARY_DIRS} ${CMAKE_INSTALL_PREFIX}
        HINTS ${ALTER_LIBRARY_PATHS}
        PATH_SUFFIXES lib lib64
        DOC "CUDNN library." )
else()
    find_library(CUDNN_LIBRARY 
        NAMES libcudnn.so libcudnn.dylib cudnn64.dll
        PATHS ${ALTER_LD_LIBRARY_PATHS} ${CUDNN_ROOT_DIR} ${PC_CUDNN_LIBRARY_DIRS} ${CMAKE_INSTALL_PREFIX}
        HINTS ${ALTER_LIBRARY_PATHS}
        PATH_SUFFIXES lib lib64
        DOC "CUDNN library." )
endif()

if(CUDNN_LIBRARY STREQUAL "CUDNN_LIBRARY-NOTFOUND")
    message(FATAL_ERROR "Can not find CuDNN Library, please refer to scripts/cmake-build/BUILD_README.md to init CUDNN env")
endif()

get_filename_component(__found_cudnn_root ${CUDNN_LIBRARY}/../.. REALPATH)
find_path(CUDNN_INCLUDE_DIR 
    NAMES cudnn.h
    HINTS $ENV{PC_CUDNN_INCLUDE_DIRS} ${CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_INCLUDE} ${__found_cudnn_root}
    PATH_SUFFIXES include 
    DOC "Path to CUDNN include directory." )

if(CUDNN_INCLUDE_DIR STREQUAL "CUDNN_INCLUDE_DIR-NOTFOUND")
    message(FATAL_ERROR "Can not find CuDNN INCLUDE, please refer to scripts/cmake-build/BUILD_README.md to init CUDNN env")
endif()

if(EXISTS ${CUDNN_INCLUDE_DIR}/cudnn_version.h)
    file(READ ${CUDNN_INCLUDE_DIR}/cudnn_version.h CUDNN_VERSION_FILE_CONTENTS)
else()
    file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_VERSION_FILE_CONTENTS)
endif()

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_PATCH_VERSION "${CUDNN_VERSION_FILE_CONTENTS}")
string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1"
    CUDNN_PATCH_VERSION "${CUDNN_PATCH_VERSION}")  
set(CUDNN_VERSION ${CUDNN_MAJOR_VERSION}.${CUDNN_MINOR_VERSION}.${CUDNN_PATCH_VERSION})



if(MGE_CUDA_USE_STATIC)
    add_library(libcudnn STATIC IMPORTED)
else()
    add_library(libcudnn SHARED IMPORTED)
endif()

set_target_properties(libcudnn PROPERTIES
    IMPORTED_LOCATION ${CUDNN_LIBRARY}
    INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR})

message(STATUS "Found CuDNN: ${__found_cudnn_root} (found version: ${CUDNN_VERSION})")