Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
cudnn.cmake 3.43 KB
Newer Older
yuguo960516yuguo's avatar
2.4.1-dtk-23.04  
yuguo960516yuguo committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31 
if(NOT WITH_GPU)
  return()
endif()

if(WIN32)
  set(CUDNN_ROOT ${CUDA_TOOLKIT_ROOT_DIR})
else()
  set(CUDNN_ROOT
      "/usr"
      CACHE PATH "CUDNN ROOT")
endif()

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)

set(TARGET_ARCH "x86_64")
if(NOT ${CMAKE_SYSTEM_PROCESSOR})
  set(TARGET_ARCH ${CMAKE_SYSTEM_PROCESSOR})
endif()

list(
  APPEND
  CUDNN_CHECK_LIBRARY_DIRS
  ${CUDNN_ROOT}
  ${CUDNN_ROOT}/lib64
  ${CUDNN_ROOT}/lib
“yuguo”'s avatar
2.5  
“yuguo” committed
32 
  ${CUDNN_ROOT}/lib/x64
yuguo960516yuguo's avatar
2.4.1-dtk-23.04  
yuguo960516yuguo committed
33
34
35
36
37 
  ${CUDNN_ROOT}/lib/${TARGET_ARCH}-linux-gnu
  ${CUDNN_ROOT}/local/cuda-${CUDA_VERSION}/targets/${TARGET_ARCH}-linux/lib/
  $ENV{CUDNN_ROOT}
  $ENV{CUDNN_ROOT}/lib64
  $ENV{CUDNN_ROOT}/lib
“yuguo”'s avatar
2.5  
“yuguo” committed
38 
  $ENV{CUDNN_ROOT}/lib/x64
yuguo960516yuguo's avatar
2.4.1-dtk-23.04  
yuguo960516yuguo committed
39
40
41
42
43
44
45
46
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115 
  /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()

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

if(APPLE)
  set(CUDNN_LIB_NAME "libcudnn.dylib" "libcudnn.so")
endif()

find_library(
  CUDNN_LIBRARY
  NAMES ${CUDNN_LIB_NAME} # libcudnn_static.a
  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()

macro(find_cudnn_version cudnn_header_file)
  file(READ ${cudnn_header_file} 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 "???")
    else()
      add_definitions("-DCUDNN_MAJOR_VERSION=\"${CUDNN_MAJOR_VERSION}\"")
      math(EXPR CUDNN_VERSION "${CUDNN_MAJOR_VERSION} * 1000 +
                 ${CUDNN_MINOR_VERSION} * 100 + ${CUDNN_PATCHLEVEL_VERSION}")
      message(
        STATUS
          "Current cuDNN header is ${cudnn_header_file} "
          "Current cuDNN version is v${CUDNN_MAJOR_VERSION}.${CUDNN_MINOR_VERSION}.${CUDNN_PATCHLEVEL_VERSION}. "
      )
    endif()
  endif()
endmacro()

if(CUDNN_FOUND)
  find_cudnn_version(${CUDNN_INCLUDE_DIR}/cudnn.h)
  if(NOT CUDNN_MAJOR_VERSION)
    find_cudnn_version(${CUDNN_INCLUDE_DIR}/cudnn_version.h)
  endif()
endif()