vstrt/{trtexec}/CMakeLists.txt: linking cuDNN

AmusementClub · Oct 22, 2023 · ce897a3 · ce897a3
1 parent 0fe13e7
commit ce897a3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/vstrt/CMakeLists.txt b/vstrt/CMakeLists.txt
@@ -39,17 +39,15 @@ endif()
 
 if (USE_NVINFER_PLUGIN)
     add_definitions(-DUSE_NVINFER_PLUGIN)
+    find_library(CUDNN_LIB cudnn HINTS ${CUDNN_ROOT_DIR} PATH_SUFFIXES lib)
     if (USE_NVINFER_PLUGIN_STATIC)
-        target_link_libraries(vstrt PRIVATE nvinfer_plugin_static)
-
-        find_library(CUDNN_LIB cudnn HINTS ${CUDNN_ROOT_DIR} PATH_SUFFIXES lib)
-        target_link_libraries(vstrt PRIVATE ${CUDNN_LIB} CUDA::cublas)
+        target_link_libraries(vstrt PRIVATE nvinfer_plugin_static ${CUDNN_LIB} CUDA::cublas)
 
         if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
             target_link_options(vstrt PRIVATE "/DELAYLOAD:cublas64_12.dll" "/DELAYLOAD:cudnn64_8.dll")
         endif()
     else()
-        target_link_libraries(vstrt PRIVATE nvinfer_plugin CUDA::cublas)
+        target_link_libraries(vstrt PRIVATE nvinfer_plugin ${CUDNN_LIB} CUDA::cublas)
 
         if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
             target_link_options(vstrt PRIVATE "/DELAYLOAD:nvinfer_plugin.dll" "/DELAYLOAD:cublas64_12.dll")

diff --git a/vstrt/trtexec/CMakeLists.txt b/vstrt/trtexec/CMakeLists.txt
@@ -3,6 +3,7 @@ cmake_minimum_required(VERSION 3.20)
 project(trtexec LANGUAGES CXX)
 
 set(TENSORRT_HOME "" CACHE PATH "Path to TensorRT")
+set(CUDNN_HOME "" CACHE PATH "Path to cuDNN")
 
 find_package(CUDAToolkit REQUIRED)
 
@@ -28,13 +29,16 @@ target_link_directories(trtexec PRIVATE
     ${TENSORRT_HOME}/lib
 )
 
+find_library(CUDNN_LIB cudnn HINTS ${CUDNN_ROOT_DIR} PATH_SUFFIXES lib)
+
 target_link_libraries(trtexec PRIVATE
     CUDA::cudart_static
     nvinfer
     nvinfer_plugin
     nvparsers
     nvonnxparser
     CUDA::cublas
+    ${CUDNN_LIB}
 )
 
 install(TARGETS trtexec RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})