From 0a4002f5dcd901cf799952153fa7c0eb418a3a56 Mon Sep 17 00:00:00 2001
From: Zhaolong Xing <nhzlx.dragon@gmail.com>
Date: Fri, 6 Dec 2019 11:05:54 +0800
Subject: [PATCH] CHERRY_PICK: Better TensorRT support (#20858) (#21578)

* Fix TensorRT detection bug

1. Add new search path for TensorRT at tensorrt.cmake
2. Add better debug message
3. Fix the bug of detection of TensorRT version

In NVIDIA official docker image, TensorRT headers are located at
`/usr/include/x86_64-linux-gnu` and TensorRT libraries are located
at `/usr/lib/x86_64-linux-gnu`, so using `-DTENSORRT_ROOT` will
fail to detect TensorRT.

There is no debug/warning message to tell developer that TensorRT
is failed to be detected.

In later version of TensorRT (e.g. v6), `NV_TENSORRT_MAJOR` is
defined at `NvInferVersion.h` instead of `NvInfer.h`, so add
compatibility fix.

* Fix TensorRT variables in CMake

1. Replace `${TENSORRT_ROOT}/include` with `${TENSORRT_INCLUDE_DIR}`
2. Replace `${TENSORRT_ROOT}/lib` with `${TENSORRT_LIBRARY}`

Manually type path may locate incorrect path of TensorRT. Use the
paths detected by system instead.

* Fix TensorRT library path

1. Add new variable - `${TENSORRT_LIBRARY_DIR}`
2. Fix TensorRT library path

inference_lib.cmake and setup.py.in need the path of TensorRT library
instead of the file of TensorRT library, so add new variable to fix it.

* Add more general search rule for TensoRT

Let system detect architecture instead of manually assign it, so
replace `x86_64-linux-gnu` with `${CMAKE_LIBRARY_ARCHITECTURE}`.

* Add more general search rule for TensorRT

Remove duplicate search rules for TensorRT libraries. Use
`${TENSORRT_LIBRARY_DIR}` to get full path of libnvinfer.so

test=release/1.6
---
 cmake/inference_lib.cmake |  2 +-
 cmake/tensorrt.cmake      | 28 +++++++++++++++++++++++++++-
 python/setup.py.in        |  4 ++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/cmake/inference_lib.cmake b/cmake/inference_lib.cmake
index b3afae79ca3..20410361f48 100644
--- a/cmake/inference_lib.cmake
+++ b/cmake/inference_lib.cmake
@@ -161,7 +161,7 @@ endif ()
 if (TENSORRT_FOUND)
     set(dst_dir "${FLUID_INFERENCE_INSTALL_DIR}/third_party/install/tensorrt")
     copy(inference_lib_dist
-            SRCS ${TENSORRT_ROOT}/include/Nv*.h ${TENSORRT_ROOT}/lib/*nvinfer*
+            SRCS ${TENSORRT_INCLUDE_DIR}/Nv*.h ${TENSORRT_LIBRARY_DIR}/*nvinfer*
             DSTS ${dst_dir}/include ${dst_dir}/lib)
 endif ()
 
diff --git a/cmake/tensorrt.cmake b/cmake/tensorrt.cmake
index fc97fcbf20a..3db8a231ccb 100644
--- a/cmake/tensorrt.cmake
+++ b/cmake/tensorrt.cmake
@@ -19,13 +19,23 @@ endif()
 
 find_path(TENSORRT_INCLUDE_DIR NvInfer.h
     PATHS ${TENSORRT_ROOT} ${TENSORRT_ROOT}/include
+    ${TENSORRT_ROOT}/include/${CMAKE_LIBRARY_ARCHITECTURE}
     $ENV{TENSORRT_ROOT} $ENV{TENSORRT_ROOT}/include
+    $ENV{TENSORRT_ROOT}/include/${CMAKE_LIBRARY_ARCHITECTURE}
     NO_DEFAULT_PATH
 )
 
-find_library(TENSORRT_LIBRARY NAMES ${TR_INFER_LIB} ${TR_INFER_RT}
+find_path(TENSORRT_LIBRARY_DIR NAMES ${TR_INFER_LIB} ${TR_INFER_RT}
     PATHS ${TENSORRT_ROOT} ${TENSORRT_ROOT}/lib
+    ${TENSORRT_ROOT}/lib/${CMAKE_LIBRARY_ARCHITECTURE}
     $ENV{TENSORRT_ROOT} $ENV{TENSORRT_ROOT}/lib
+    $ENV{TENSORRT_ROOT}/lib/${CMAKE_LIBRARY_ARCHITECTURE}
+    NO_DEFAULT_PATH
+    DOC "Path to TensorRT library."
+)
+
+find_library(TENSORRT_LIBRARY NAMES ${TR_INFER_LIB} ${TR_INFER_RT}
+    PATHS ${TENSORRT_LIBRARY_DIR}
     NO_DEFAULT_PATH
     DOC "Path to TensorRT library.")
 
@@ -35,12 +45,28 @@ if(TENSORRT_INCLUDE_DIR AND TENSORRT_LIBRARY)
   endif(WITH_DSO)
 else()
     set(TENSORRT_FOUND OFF)
+    if(WITH_DSO)
+        message(WARNING "TensorRT is NOT found.")
+    else(WITH_DSO)
+        message(WARNING "TensorRT is disabled because WITH_DSO is OFF.")
+    endif(WITH_DSO)
 endif()
 
 if(TENSORRT_FOUND)
     file(READ ${TENSORRT_INCLUDE_DIR}/NvInfer.h TENSORRT_VERSION_FILE_CONTENTS)
     string(REGEX MATCH "define NV_TENSORRT_MAJOR +([0-9]+)" TENSORRT_MAJOR_VERSION
         "${TENSORRT_VERSION_FILE_CONTENTS}")
+
+    if("${TENSORRT_MAJOR_VERSION}" STREQUAL "")
+        file(READ ${TENSORRT_INCLUDE_DIR}/NvInferVersion.h TENSORRT_VERSION_FILE_CONTENTS)
+        string(REGEX MATCH "define NV_TENSORRT_MAJOR +([0-9]+)" TENSORRT_MAJOR_VERSION
+        "${TENSORRT_VERSION_FILE_CONTENTS}")
+    endif()
+
+    if("${TENSORRT_MAJOR_VERSION}" STREQUAL "")
+        message(SEND_ERROR "Failed to detect TensorRT version.")
+    endif()
+
     string(REGEX REPLACE "define NV_TENSORRT_MAJOR +([0-9]+)" "\\1"
         TENSORRT_MAJOR_VERSION "${TENSORRT_MAJOR_VERSION}")
 
diff --git a/python/setup.py.in b/python/setup.py.in
index 35033683a80..79edac467f5 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -173,8 +173,8 @@ package_data['paddle.libs']=[('libwarpctc' if os.name != 'nt' else 'warpctc') +
 shutil.copy('${WARPCTC_LIBRARIES}', libs_path)
 
 if '${TENSORRT_FOUND}' == 'ON' and os.name == 'nt':
-    shutil.copy(os.path.join('${TENSORRT_ROOT}', 'lib', '${TR_INFER_RT}'), libs_path)
-    shutil.copy(os.path.join('${TENSORRT_ROOT}', 'lib', '${TR_INFER_PLUGIN_RT}'), libs_path)
+    shutil.copy(os.path.join('${TENSORRT_LIBRARY_DIR}', '${TR_INFER_RT}'), libs_path)
+    shutil.copy(os.path.join('${TENSORRT_LIBRARY_DIR}', '${TR_INFER_PLUGIN_RT}'), libs_path)
     package_data['paddle.libs'] += ['${TR_INFER_RT}', '${TR_INFER_PLUGIN_RT}']
 
 if '${WITH_MKL}' == 'ON':
-- 
GitLab