From fcd6d6757990f12a0c37f0f543080ffaca07a544 Mon Sep 17 00:00:00 2001
From: xiaoxiaohehe001 <49090790+xiaoxiaohehe001@users.noreply.github.com>
Date: Thu, 5 Jan 2023 20:33:20 +0800
Subject: [PATCH] [Paddle Inference] Add ci flags for a persistent IBuilder.
 (#49538)

---
 paddle/fluid/inference/tensorrt/engine.h | 31 +++++++++---------------
 paddle/fluid/inference/tensorrt/helper.h | 11 +++++++++
 paddle/phi/core/flags.cc                 | 13 ++++++++++
 paddle/scripts/paddle_build.sh           |  2 ++
 4 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/engine.h b/paddle/fluid/inference/tensorrt/engine.h
index 8a4afa3255..a3f3921db4 100644
--- a/paddle/fluid/inference/tensorrt/engine.h
+++ b/paddle/fluid/inference/tensorrt/engine.h
@@ -42,6 +42,8 @@ limitations under the License. */
 #include "paddle/phi/core/stream.h"
 #include "paddle/utils/any.h"
 
+DECLARE_bool(trt_ibuilder_cache);
+
 namespace paddle {
 namespace inference {
 namespace tensorrt {
@@ -679,16 +681,6 @@ class TensorRTEngine {
   std::vector<std::unique_ptr<nvinfer1::IPluginV2IOExt>> owned_plugin_v2ioext_;
 
   // TensorRT related internal members
-  template <typename T>
-  struct Destroyer {
-    void operator()(T* x) {
-      if (x) {
-        x->destroy();
-      }
-    }
-  };
-  template <typename T>
-  using infer_ptr = std::unique_ptr<T, Destroyer<T>>;
   infer_ptr<nvinfer1::IBuilder> infer_builder_;
   infer_ptr<nvinfer1::INetworkDefinition> infer_network_;
   infer_ptr<nvinfer1::ICudaEngine> infer_engine_;
@@ -733,6 +725,15 @@ class TRTEngineManager {
   using AllocationPtr = phi::Allocator::AllocationPtr;
 
  public:
+  TRTEngineManager() {
+    // createInferBuilder loads trt kernels and take a few second
+    // But as long as one IBuilder lives, trt kernel will not be unloaded
+    // Hence, a persistent IBuilder to avoid TensorRT unload/reload kernels
+    if (FLAGS_trt_ibuilder_cache) {
+      holder_.reset(createInferBuilder(&NaiveLogger::Global()));
+    }
+  }
+
   bool Empty() const {
     std::lock_guard<std::mutex> lock(mutex_);
     return engines_.size() == 0;
@@ -855,15 +856,7 @@ class TRTEngineManager {
   size_t max_ctx_mem_size_{0};
   std::unordered_map<PredictorID, AllocationPtr> context_memorys_;
   std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_;
-  // createInferBuilder loads trt kernels and take a few second
-  // But as long as one IBuilder lives, trt kernel will not be unloaded
-  // Hence, a persistent IBuilder to avoid TensorRT unload/reload kernels
-  std::unique_ptr<nvinfer1::IBuilder, std::function<void(nvinfer1::IBuilder*)>>
-      holder{createInferBuilder(&NaiveLogger::Global()), [](auto* ptr) {
-               if (ptr) {
-                 ptr->destroy();
-               }
-             }};
+  infer_ptr<nvinfer1::IBuilder> holder_;
 };
 
 }  // namespace tensorrt
diff --git a/paddle/fluid/inference/tensorrt/helper.h b/paddle/fluid/inference/tensorrt/helper.h
index 3d6bb6da3d..1effa5c671 100644
--- a/paddle/fluid/inference/tensorrt/helper.h
+++ b/paddle/fluid/inference/tensorrt/helper.h
@@ -95,6 +95,17 @@ static std::tuple<int, int, int> GetTrtCompileVersion() {
       NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH};
 }
 
+template <typename T>
+struct Destroyer {
+  void operator()(T* x) {
+    if (x) {
+      x->destroy();
+    }
+  }
+};
+template <typename T>
+using infer_ptr = std::unique_ptr<T, Destroyer<T>>;
+
 // A logger for create TensorRT infer builder.
 class NaiveLogger : public nvinfer1::ILogger {
  public:
diff --git a/paddle/phi/core/flags.cc b/paddle/phi/core/flags.cc
index 6850d91b8d..5f7058fdc1 100644
--- a/paddle/phi/core/flags.cc
+++ b/paddle/phi/core/flags.cc
@@ -1165,3 +1165,16 @@ PADDLE_DEFINE_EXPORTED_bool(enable_cudnn_frontend, false, "");
  */
 PADDLE_DEFINE_EXPORTED_int32(cudnn_cache_saturation_count, 1, "");
 #endif  // PADDLE_WITH_CUDNN_FRONTEND
+
+/**
+ * CI related FLAG
+ * Name: trt_ibuilder_cache
+ * Since Version: 2.5.0
+ * Value Range: bool, default=false
+ * Example:
+ * Note: This FLAG is only enabled when CI is running. If True, a persistent
+ * IBuilder is added to avoid TensorRT unload/reload kernels.
+ */
+PADDLE_DEFINE_EXPORTED_bool(trt_ibuilder_cache,
+                            false,
+                            "Add a persistent ibuilder.");
diff --git a/paddle/scripts/paddle_build.sh b/paddle/scripts/paddle_build.sh
index 1587ff83fb..6b3a3c4487 100644
--- a/paddle/scripts/paddle_build.sh
+++ b/paddle/scripts/paddle_build.sh
@@ -1403,6 +1403,7 @@ EOF
 set -x
         # set trt_convert ut to run 15% cases.
         export TEST_NUM_PERCENT_CASES=0.15
+        export FLAGS_trt_ibuilder_cache=1
         precison_cases=""
         bash $PADDLE_ROOT/tools/check_added_ut.sh
         if [ ${PRECISION_TEST:-OFF} == "ON" ]; then
@@ -2547,6 +2548,7 @@ EOF
 set -x
         # set trt_convert ut to run 15% cases.
         export TEST_NUM_PERCENT_CASES=0.15
+        export FLAGS_trt_ibuilder_cache=1
         precison_cases=""
         bash $PADDLE_ROOT/tools/check_added_ut.sh
         #check change of pr_unnitests and dev_unnitests
-- 
GitLab