From ce897a3dac45a9d78a857a8944cd15d90d318877 Mon Sep 17 00:00:00 2001
From: WolframRhodium <WolframRhodium@users.noreply.github.com>
Date: Mon, 23 Oct 2023 07:00:28 +0800
Subject: [PATCH] vstrt/{trtexec}/CMakeLists.txt: linking cuDNN

---
 vstrt/CMakeLists.txt         | 8 +++-----
 vstrt/trtexec/CMakeLists.txt | 4 ++++
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/vstrt/CMakeLists.txt b/vstrt/CMakeLists.txt
index fb3fe33..1185c64 100644
--- 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
index 36cadd4..d04208e 100644
--- 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,6 +29,8 @@ 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
@@ -35,6 +38,7 @@ target_link_libraries(trtexec PRIVATE
     nvparsers
     nvonnxparser
     CUDA::cublas
+    ${CUDNN_LIB}
 )
 
 install(TARGETS trtexec RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})