findCUDNN

2019-11-06 14:04:23 +01:00
parent 92f3d1c548
commit aa5927d8a1
3 changed files with 41 additions and 20 deletions
@@ -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)