From 469e717ae7a9cfce96a4117fcf022361b6185f88 Mon Sep 17 00:00:00 2001
From: zlsh80826 <zlsh80826@gmail.com>
Date: Fri, 7 Aug 2020 09:57:10 +0800
Subject: [PATCH] add nvinfer_plugin dyloader

---
 paddle/fluid/platform/dynload/tensorrt.cc | 37 +++++++++++++++++------
 paddle/fluid/platform/dynload/tensorrt.h  | 31 +++++++++++++++++++
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/paddle/fluid/platform/dynload/tensorrt.cc b/paddle/fluid/platform/dynload/tensorrt.cc
index c9c3a9456b7..8ddc9e982ba 100644
--- a/paddle/fluid/platform/dynload/tensorrt.cc
+++ b/paddle/fluid/platform/dynload/tensorrt.cc
@@ -22,19 +22,15 @@ namespace dynload {
 std::once_flag tensorrt_dso_flag;
 void* tensorrt_dso_handle;
 
+std::once_flag tensorrt_plugin_dso_flag;
+void* tensorrt_plugin_dso_handle;
+
 #define DEFINE_WRAP(__name) DynLoad__##__name __name
 
 TENSORRT_RAND_ROUTINE_EACH(DEFINE_WRAP);
+TENSORRT_PLUGIN_RAND_ROUTINE_EACH(DEFINE_WRAP);
 
-void* GetTensorRtHandle() {
-#if defined(__APPLE__) || defined(__OSX__)
-  std::string dso_name = "libnvinfer.dylib";
-#elif defined(_WIN32)
-  std::string dso_name = "nvinfer.dll";
-#else
-  std::string dso_name = "libnvinfer.so";
-#endif
-
+void* GetDsoHandle(const std::string& dso_name) {
 #if !defined(_WIN32)
   int dynload_flags = RTLD_LAZY | RTLD_LOCAL;
 #else
@@ -65,10 +61,31 @@ void* GetTensorRtHandle() {
 #endif  // !_WIN32
     std::cerr << string::Sprintf(error_msg, dso_name, errorno);
   }
-
   return dso_handle;
 }
 
+void* GetTensorRtHandle() {
+#if defined(__APPLE__) || defined(__OSX__)
+  std::string dso_name = "libnvinfer.dylib";
+#elif defined(_WIN32)
+  std::string dso_name = "nvinfer.dll";
+#else
+  std::string dso_name = "libnvinfer.so";
+#endif
+  return GetDsoHandle(dso_name);
+}
+
+void* GetTensorRtPluginHandle() {
+#if defined(__APPLE__) || defined(__OSX__)
+  std::string dso_name = "libnvinfer_plugin.dylib";
+#elif defined(_WIN32)
+  std::string dso_name = "nvinfer_plugin.dll";
+#else
+  std::string dso_name = "libnvinfer_plugin.so";
+#endif
+  return GetDsoHandle(dso_name);
+}
+
 }  // namespace dynload
 }  // namespace platform
 }  // namespace paddle
diff --git a/paddle/fluid/platform/dynload/tensorrt.h b/paddle/fluid/platform/dynload/tensorrt.h
index 60e299385d6..c9982274f73 100644
--- a/paddle/fluid/platform/dynload/tensorrt.h
+++ b/paddle/fluid/platform/dynload/tensorrt.h
@@ -14,6 +14,7 @@ limitations under the License. */
 #pragma once
 
 #include <NvInfer.h>
+#include <NvInferPlugin.h>
 #if !defined(_WIN32)
 #include <dlfcn.h>
 #endif
@@ -28,10 +29,14 @@ namespace platform {
 namespace dynload {
 
 void* GetTensorRtHandle();
+void* GetTensorRtPluginHandle();
 
 extern std::once_flag tensorrt_dso_flag;
 extern void* tensorrt_dso_handle;
 
+extern std::once_flag tensorrt_plugin_dso_flag;
+extern void* tensorrt_plugin_dso_handle;
+
 #define DECLARE_DYNAMIC_LOAD_TENSORRT_WRAP(__name)                            \
   struct DynLoad__##__name {                                                  \
     template <typename... Args>                                               \
@@ -52,12 +57,38 @@ extern void* tensorrt_dso_handle;
   };                                                                          \
   extern DynLoad__##__name __name
 
+#define DECLARE_DYNAMIC_LOAD_TENSORRT_PLUGIN_WRAP(__name)                      \
+  struct DynLoad__##__name {                                                   \
+    template <typename... Args>                                                \
+    auto operator()(Args... args) -> DECLARE_TYPE(__name, args...) {           \
+      using tensorrt_plugin_func = decltype(&::__name);                        \
+      std::call_once(tensorrt_plugin_dso_flag, []() {                          \
+        tensorrt_plugin_dso_handle =                                           \
+            paddle::platform::dynload::GetTensorRtPluginHandle();              \
+        PADDLE_ENFORCE_NOT_NULL(                                               \
+            tensorrt_plugin_dso_handle,                                        \
+            platform::errors::Unavailable("Load tensorrt plugin %s failed",    \
+                                          #__name));                           \
+      });                                                                      \
+      static void* p_##__name = dlsym(tensorrt_plugin_dso_handle, #__name);    \
+      PADDLE_ENFORCE_NOT_NULL(p_##__name,                                      \
+                              platform::errors::Unavailable(                   \
+                                  "Load tensorrt plugin %s failed", #__name)); \
+      return reinterpret_cast<tensorrt_plugin_func>(p_##__name)(args...);      \
+    }                                                                          \
+  };                                                                           \
+  extern DynLoad__##__name __name
+
 #define TENSORRT_RAND_ROUTINE_EACH(__macro) \
   __macro(createInferBuilder_INTERNAL);     \
   __macro(createInferRuntime_INTERNAL);     \
   __macro(getPluginRegistry);
 
+#define TENSORRT_PLUGIN_RAND_ROUTINE_EACH(__macro) \
+  __macro(initLibNvInferPlugins);
+
 TENSORRT_RAND_ROUTINE_EACH(DECLARE_DYNAMIC_LOAD_TENSORRT_WRAP)
+TENSORRT_PLUGIN_RAND_ROUTINE_EACH(DECLARE_DYNAMIC_LOAD_TENSORRT_PLUGIN_WRAP)
 
 }  // namespace dynload
 }  // namespace platform
-- 
GitLab