Cuda.cmakeをcuDNN 8に対応させた

lltcggie · Sep 5, 2020 · ed6cff4 · ed6cff4
1 parent 98f9aea
commit ed6cff4
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/cmake/Cuda.cmake b/cmake/Cuda.cmake
@@ -190,7 +190,7 @@ function(detect_cuDNN)
             PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT} ${CUDA_TOOLKIT_INCLUDE}
             PATH_SUFFIXES include
             DOC "Path to cuDNN include directory." )
-           
+
   unset(_path_suffixes)
   if(MSVC AND ${CMAKE_SIZEOF_VOID_P} EQUAL 8)
     set(_path_suffixes PATH_SUFFIXES lib/x64)
@@ -217,7 +217,11 @@ function(detect_cuDNN)
     set(HAVE_CUDNN  TRUE PARENT_SCOPE)
     set(CUDNN_FOUND TRUE PARENT_SCOPE)
 
-    file(READ ${CUDNN_INCLUDE}/cudnn.h CUDNN_VERSION_FILE_CONTENTS)
+    if (EXISTS ${CUDNN_INCLUDE}/cudnn_version.h)
+      file(READ ${CUDNN_INCLUDE}/cudnn_version.h CUDNN_VERSION_FILE_CONTENTS)
+    else()
+      file(READ ${CUDNN_INCLUDE}/cudnn.h CUDNN_VERSION_FILE_CONTENTS)
+    endif()
 
     # cuDNN v3 and beyond
     string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"