Merge pull request #344 from shivak/master

Various fixes for obtaining and finding the CUDNN library on Windows.

2 files changed, 15 insertions, 13 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b4deeed..4cf64b8 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -14,7 +14,7 @@ FIND_PACKAGE(CUDA 7.0 REQUIRED)
 
 FIND_PACKAGE(CUDNN 5 EXACT QUIET)
 IF(NOT CUDNN_FOUND)
-  CUDNN_INSTALL(5.1 "${Torch_INSTALL_LIB}" "${Torch_INSTALL_INCLUDE}" "")
+  CUDNN_INSTALL(5.1 "${Torch_INSTALL_LIB}" "${Torch_INSTALL_INCLUDE}" "${Torch_INSTALL_BIN}")
   FIND_PACKAGE(CUDNN 5 EXACT REQUIRED)
 ENDIF()
 
diff --git a/cmake/FindCUDNN.cmake b/cmake/FindCUDNN.cmake
index b1f7349..5aa6158 100644
--- a/cmake/FindCUDNN.cmake
+++ b/cmake/FindCUDNN.cmake
@@ -41,6 +41,7 @@ function(CUDNN_INSTALL version dest_libdir dest_incdir dest_bindir)
   set(tar_incdir cuda/include)
 
   if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
+    set(url_extension tgz)
     if("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64")
       set(url_arch_name linux-x64 )
     elseif("${CMAKE_SYSTEM_PROCESSOR}" MATCHES "ppc")
@@ -48,15 +49,17 @@ function(CUDNN_INSTALL version dest_libdir dest_incdir dest_bindir)
       #  TX1 has to be installed via JetPack
     endif()
   elseif  (APPLE)
+    set(url_extension tgz)
     set(tar_libdir cuda/lib)
     set(url_arch_name osx-x64)
   elseif(WIN32)
+    set(url_extension zip)
     set(tar_bindir cuda/bin)
     set(tar_libdir cuda/lib/x64)
     if(CMAKE_SYSTEM_VERSION MATCHES "10")
-      set(url_arch_name windows10)
+      set(url_arch_name windows10-x64)
     else()
-      set(url_arch_name windows7)
+      set(url_arch_name windows7-x64)
     endif()
   endif()
   
@@ -64,17 +67,17 @@ function(CUDNN_INSTALL version dest_libdir dest_incdir dest_bindir)
   if(url_arch_name) 
     set(download_dir ${CMAKE_CURRENT_BINARY_DIR}/downloads/cudnn${version})
     file(MAKE_DIRECTORY ${download_dir})
-    set(cudnn_filename cudnn-${CUDA_VERSION}-${url_arch_name}-v${version}.tgz)
+    set(cudnn_filename cudnn-${CUDA_VERSION}-${url_arch_name}-v${version}.${url_extension})
     set(base_url http://developer.download.nvidia.com/compute/redist/cudnn)
     set(cudnn_url ${base_url}/v${version_base}/${cudnn_filename})
-    set(cudnn_tgz ${download_dir}/${cudnn_filename})
+    set(cudnn_file ${download_dir}/${cudnn_filename})
     
-    if(NOT EXISTS ${cudnn_tgz})
+    if(NOT EXISTS ${cudnn_file})
       message(STATUS "Downloading CUDNN library from NVIDIA...")
-      file(DOWNLOAD ${cudnn_url} ${cudnn_tgz}
+      file(DOWNLOAD ${cudnn_url} ${cudnn_file}
 	SHOW_PROGRESS STATUS cudnn_status
 	)
-      execute_process(COMMAND ${CMAKE_COMMAND} -E tar xzvf ${cudnn_tgz} WORKING_DIRECTORY ${download_dir} RESULT_VARIABLE cudnn_status)
+    execute_process(COMMAND ${CMAKE_COMMAND} -E tar xzvf --format=${url_extension} ${cudnn_file} WORKING_DIRECTORY ${download_dir} RESULT_VARIABLE cudnn_status)
 
       if(NOT "${cudnn_status}" MATCHES "0")
 	message(STATUS "Was not able to download CUDNN from ${cudnn_url}. Please install CuDNN manually from https://developer.nvidia.com/cuDNN")
@@ -91,7 +94,7 @@ function(CUDNN_INSTALL version dest_libdir dest_incdir dest_bindir)
 
     file(COPY ${download_dir}/${tar_libdir}/ DESTINATION  ${dest_libdir} )
 
-    get_filename_component(dest_dir ${dest_libdir} PATH)
+    get_filename_component(dest_dir ${dest_libdir} DIRECTORY)
 
     set(CUDNN_ROOT_DIR ${dest_dir} PARENT_SCOPE)
     unset(CUDNN_LIBRARY CACHE)
@@ -112,18 +115,17 @@ get_filename_component(__libpath_cudart "${CUDA_CUDART_LIBRARY}" PATH)
 # More strict enforcement of minor/patch version is done if/when the header file is examined.
 if(CUDNN_FIND_VERSION_EXACT)
   SET(__cudnn_ver_suffix ".${CUDNN_FIND_VERSION_MAJOR}")
-  SET(__cudnn_lib_win_name cudnn64_${CUDNN_FIND_VERSION_MAJOR}.dll)
+  SET(__cudnn_lib_win_name cudnn64_${CUDNN_FIND_VERSION_MAJOR})
 else()
-  SET(__cudnn_lib_win_name cudnn64.dll)
+  SET(__cudnn_lib_win_name cudnn64)
 endif()
 
 find_library(CUDNN_LIBRARY 
   NAMES libcudnn.so${__cudnn_ver_suffix} libcudnn${__cudnn_ver_suffix}.dylib ${__cudnn_lib_win_name}
   PATHS $ENV{LD_LIBRARY_PATH} ${__libpath_cudart} ${CUDNN_ROOT_DIR} ${PC_CUDNN_LIBRARY_DIRS} ${CMAKE_INSTALL_PREFIX}
-  PATH_SUFFIXES lib lib64
+  PATH_SUFFIXES lib lib64 bin
   DOC "CUDNN library." )
 
-
 if(CUDNN_LIBRARY)
   SET(CUDNN_MAJOR_VERSION ${CUDNN_FIND_VERSION_MAJOR})
   set(CUDNN_VERSION ${CUDNN_MAJOR_VERSION})