From aa5927d8a168b44f32a13dfa0551d78594e9fe33 Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Wed, 6 Nov 2019 14:04:23 +0100
Subject: [PATCH] findCUDNN

---
 CMakeLists.txt          | 12 +++---------
 cmake/FindCUDNN.cmake   | 33 +++++++++++++++++++++++++++++++++
 cmake/tkDNNConfig.cmake | 16 +++++-----------
 3 files changed, 41 insertions(+), 20 deletions(-)
 create mode 100644 cmake/FindCUDNN.cmake

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 467e823..b4b3c38 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -18,17 +18,11 @@ find_package(CUDA 9.0 REQUIRED)
 SET(CUDA_SEPARABLE_COMPILATION ON)
 #set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -arch=sm_30 --compiler-options '-fPIC'")
 
-# compile Discovery only if TensorRT is installed
-find_library(NVINFER NAMES nvinfer)
-if(NVINFER STREQUAL "NVINFER-NOTFOUND")
-    set(NVINFER_INCLUDES "/usr/local/nvidia/tensorrt/include/")
-    link_directories(/usr/local/nvidia/tensorrt/targets/x86_64-linux-gnu/lib/ 
-		     /usr/local/cuda/targets/x86_64-linux/lib/)
-endif()
+find_package(CUDNN REQUIRED)
 
 # compile
 file(GLOB tkdnn_CUSRC "src/kernels/*.cu")
-cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${NVINFER_INCLUDES})
+cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIRS})
 cuda_add_library(kernels SHARED ${tkdnn_CUSRC})
 
 
@@ -43,7 +37,7 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPENCV")
 # Build Libraries
 #-------------------------------------------------------------------------------
 file(GLOB tkdnn_SRC "src/*.cpp")
-set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn -lnvinfer ${OpenCV_LIBS})
+set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRARIES} ${OpenCV_LIBS})
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -std=c++11")
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${OPENCV_INCLUDE_DIRS} ${NVINFER_INCLUDES})
diff --git a/cmake/FindCUDNN.cmake b/cmake/FindCUDNN.cmake
new file mode 100644
index 0000000..0e99052
--- /dev/null
+++ b/cmake/FindCUDNN.cmake
@@ -0,0 +1,33 @@
+# Find the header files
+
+find_path(CUDNN_INCLUDE_DIR
+        ${CMAKE_SYSROOT}/usr/local/include
+        ${CMAKE_SYSROOT}/usr/include
+        /usr/local/nvidia/tensorrt/include/
+        NO_DEFAULT_PATH
+)
+
+set(OLD_ROOT ${CMAKE_FIND_ROOT_PATH})
+list(APPEND CMAKE_FIND_ROOT_PATH /)
+list(APPEND CMAKE_FIND_LIBRARY_SUFFIXES .so.7)
+list(APPEND CMAKE_FIND_LIBRARY_SUFFIXES .so.5)
+find_library(CUDNN_LIB
+    NAMES cudnn
+    PATHS
+    /usr/local/driveworks-2.0/targets/${CMAKE_SYSTEM_PROCESSOR}-Linux/lib
+    /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/
+    NO_DEFAULT_PATH
+)
+find_library(CUDNN_NVLIB
+    NAMES "nvinfer"
+    PATHS
+    /usr/local/driveworks-2.0/targets/${CMAKE_SYSTEM_PROCESSOR}-Linux/lib
+    /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/
+    NO_DEFAULT_PATH
+)
+set(CMAKE_FIND_ROOT_PATH ${OLD_ROOT})
+
+set(CUDNN_LIBRARIES ${CUDNN_LIB} ${CUDNN_NVLIB})
+message("-- Found CUDNN: "  ${CUDNN_LIB})
+message("-- Found NVINFER: "  ${CUDNN_NVLIB})
+set(CUDNN_FOUND true)
diff --git a/cmake/tkDNNConfig.cmake b/cmake/tkDNNConfig.cmake
index fd06286..4dbaf21 100644
--- a/cmake/tkDNNConfig.cmake
+++ b/cmake/tkDNNConfig.cmake
@@ -4,27 +4,21 @@ set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} --std=c++11 -fPIC")
 
 find_package(CUDA REQUIRED)
 find_package(OpenCV REQUIRED)
-find_library(NVINFER NAMES nvinfer)
-if(NVINFER STREQUAL "NVINFER-NOTFOUND")
-    set(NVINFER_INCLUDES "/usr/local/nvidia/tensorrt/include/")
-    link_directories(/usr/local/nvidia/tensorrt/targets/x86_64-linux-gnu/lib/ 
-		     /usr/local/cuda/targets/x86_64-linux/lib/)
-endif()
+find_package(CUDNN REQUIRED)
 
 set(tkDNN_INCLUDE_DIRS 
 	${CUDA_INCLUDE_DIRS} 
 	${OPENCV_INCLUDE_DIRS} 
-    ${NVINFER_INCLUDES}
+    ${CUDNN_INCLUDE_DIRS}
 )
 
 set(tkDNN_LIBRARIES 
     tkDNN 
     kernels 
     ${CUDA_LIBRARIES} 
-    ${CUDA_CUBLAS_LIBRARIES} 
-    -lcudnn 
-    -lnvinfer 
-    ${OpenCV_LIBS}
+    ${CUDA_CUBLAS_LIBRARIES}
+	${CUDNN_LIBRARIES}
+	${OpenCV_LIBS}
 )
 
 set(tkDNN_FOUND true)