From a077330c3e6164c29b174d0e406c77b6c4c0835c Mon Sep 17 00:00:00 2001
From: Chi Lo <54722500+chilo-ms@users.noreply.github.com>
Date: Wed, 24 Apr 2024 21:46:54 -0700
Subject: [PATCH] [TensorRT] adapt for TRT lib name change after TRT 10 GA
 (#20445)

For TensorRT 10 GA onwards, the TensorRT libraries will have major
version appended to the end on Windows, for example, nvinfer_10.dll,
nvinfer_plugin_10.dll, nvonnxparser_10.dll ...

Change cmake file accordingly.
---
 cmake/onnxruntime_providers_tensorrt.cmake | 74 ++++++++++++++++++++--
 1 file changed, 67 insertions(+), 7 deletions(-)

diff --git a/cmake/onnxruntime_providers_tensorrt.cmake b/cmake/onnxruntime_providers_tensorrt.cmake
index 15ffc29e79ff4..42275873ceeb0 100644
--- a/cmake/onnxruntime_providers_tensorrt.cmake
+++ b/cmake/onnxruntime_providers_tensorrt.cmake
@@ -34,21 +34,80 @@
     MESSAGE(STATUS "[Note] There is an issue when running \"Debug build\" TRT EP with \"Release build\" TRT built-in parser on Windows. This build will use tensorrt oss parser instead.")
   endif()
 
+  find_path(TENSORRT_INCLUDE_DIR NvInfer.h
+    HINTS ${TENSORRT_ROOT}
+    PATH_SUFFIXES include)
+
+  # TensorRT 10 GA onwards, the TensorRT libraries will have major version appended to the end on Windows,
+  # for example, nvinfer_10.dll, nvinfer_plugin_10.dll, nvonnxparser_10.dll ...
+  if (WIN32)
+    file(READ ${TENSORRT_INCLUDE_DIR}/NvInferVersion.h NVINFER_VER_CONTENT)
+    string(REGEX MATCH "define NV_TENSORRT_MAJOR * +([0-9]+)" NV_TENSORRT_MAJOR "${NVINFER_VER_CONTENT}")
+    string(REGEX REPLACE "define NV_TENSORRT_MAJOR * +([0-9]+)" "\\1" NV_TENSORRT_MAJOR "${NV_TENSORRT_MAJOR}")
+    string(REGEX MATCH "define NV_TENSORRT_MINOR * +([0-9]+)" NV_TENSORRT_MINOR "${NVINFER_VER_CONTENT}")
+    string(REGEX REPLACE "define NV_TENSORRT_MINOR * +([0-9]+)" "\\1" NV_TENSORRT_MINOR "${NV_TENSORRT_MINOR}")
+    string(REGEX MATCH "define NV_TENSORRT_PATCH * +([0-9]+)" NV_TENSORRT_PATCH "${NVINFER_VER_CONTENT}")
+    string(REGEX REPLACE "define NV_TENSORRT_PATCH * +([0-9]+)" "\\1" NV_TENSORRT_PATCH "${NV_TENSORRT_PATCH}")
+    math(EXPR NV_TENSORRT_MAJOR_INT "${NV_TENSORRT_MAJOR}")
+    math(EXPR NV_TENSORRT_MINOR_INT "${NV_TENSORRT_MINOR}")
+    math(EXPR NV_TENSORRT_PATCH_INT "${NV_TENSORRT_PATCH}")
+
+    if (NV_TENSORRT_MAJOR)
+      MESSAGE(STATUS "NV_TENSORRT_MAJOR is ${NV_TENSORRT_MAJOR}")
+    else()
+      MESSAGE(STATUS "Can't find NV_TENSORRT_MAJOR macro")
+    endif()
+
+    # Check TRT version >= 10.0.1.6 (Note: TRT 10 EA is 10.0.0.6 but with no major version appended to the end)
+    if ((NV_TENSORRT_MAJOR_INT GREATER 10) OR
+        (NV_TENSORRT_MAJOR_INT EQUAL 10 AND NV_TENSORRT_MINOR_INT GREATER 0) OR
+        (NV_TENSORRT_MAJOR_INT EQUAL 10 AND NV_TENSORRT_PATCH_INT GREATER 0))
+       set(NVINFER_LIB "nvinfer_${NV_TENSORRT_MAJOR}")
+       set(NVINFER_PLUGIN_LIB "nvinfer_plugin_${NV_TENSORRT_MAJOR}")
+       set(PARSER_LIB "nvonnxparser_${NV_TENSORRT_MAJOR}")
+    endif()
+  endif()
+
+  if (NOT NVINFER_LIB)
+     set(NVINFER_LIB "nvinfer")
+  endif()
+
+  if (NOT NVINFER_PLUGIN_LIB)
+     set(NVINFER_PLUGIN_LIB "nvinfer_plugin")
+  endif()
+
+  if (NOT PARSER_LIB)
+     set(PARSER_LIB "nvonnxparser")
+  endif()
+
   if (onnxruntime_USE_TENSORRT_BUILTIN_PARSER)
     # Add TensorRT library
-    find_path(TENSORRT_INCLUDE_DIR NvInfer.h
-      HINTS ${TENSORRT_ROOT}
-      PATH_SUFFIXES include)
-    MESSAGE(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
-    find_library(TENSORRT_LIBRARY_INFER nvinfer
+    MESSAGE(STATUS "Search for ${NVINFER_LIB}, ${NVINFER_PLUGIN_LIB} and ${PARSER_LIB}")
+
+    find_library(TENSORRT_LIBRARY_INFER ${NVINFER_LIB}
       HINTS ${TENSORRT_ROOT}
       PATH_SUFFIXES lib lib64 lib/x64)
-    find_library(TENSORRT_LIBRARY_INFER_PLUGIN nvinfer_plugin
+
+    if (NOT TENSORRT_LIBRARY_INFER)
+      MESSAGE(STATUS "Can't find ${NVINFER_LIB}")
+    endif()
+
+    find_library(TENSORRT_LIBRARY_INFER_PLUGIN ${NVINFER_PLUGIN_LIB}
       HINTS  ${TENSORRT_ROOT}
       PATH_SUFFIXES lib lib64 lib/x64)
-    find_library(TENSORRT_LIBRARY_NVONNXPARSER nvonnxparser
+
+    if (NOT TENSORRT_LIBRARY_INFER_PLUGIN)
+      MESSAGE(STATUS "Can't find ${NVINFER_PLUGIN_LIB}")
+    endif()
+
+    find_library(TENSORRT_LIBRARY_NVONNXPARSER ${PARSER_LIB}
       HINTS  ${TENSORRT_ROOT}
       PATH_SUFFIXES lib lib64 lib/x64)
+
+    if (NOT TENSORRT_LIBRARY_NVONNXPARSER)
+      MESSAGE(STATUS "Can't find ${PARSER_LIB}")
+    endif()
+
     set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_INFER_PLUGIN} ${TENSORRT_LIBRARY_NVONNXPARSER})
     MESSAGE(STATUS "Find TensorRT libs at ${TENSORRT_LIBRARY}")
   else()
@@ -77,6 +136,7 @@
       target_compile_options(nvonnxparser_static PRIVATE /FIio.h /wd4100)
       target_compile_options(nvonnxparser PRIVATE /FIio.h /wd4100)
     endif()
+    # Static libraries are just nvonnxparser_static on all platforms
     set(onnxparser_link_libs nvonnxparser_static)
   endif()