add IPluginV2Layer: AddPluginV2Ext (#36493)

paddle/fluid/inference/tensorrt/engine.cc

@@ -135,12 +135,6 @@ void TensorRTEngine::FreezeNetwork() {
         }
         for (int j = 0; j < layer->getNbOutputs(); j++) {
           auto *temp_out = layer->getOutput(j);
-          if (temp_out->isNetworkOutput()) {
-            VLOG(1) << "Layer(Name: " << layer->getName()
-                    << ") is set to float32 because its output("
-                    << temp_out->getName() << ") is the output of the network.";
-            return false;
-          }
           if (!temp_out->dynamicRangeIsSet()) {
             VLOG(1) << "Layer(Name: " << layer->getName()
                     << ") is set to float32 because its output("
@@ -357,6 +351,13 @@ nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2Ext(
   return network()->addPluginV2(inputs, num_inputs, *plugin);
 }
 
+nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2IOExt(
+    nvinfer1::ITensor *const *inputs, int num_inputs,
+    nvinfer1::IPluginV2IOExt *plugin) {
+  owned_plugin_v2ioext_.emplace_back(plugin);
+  return network()->addPluginV2(inputs, num_inputs, *plugin);
+}
+
 void TensorRTEngine::freshDeviceId() {
   int count;
   cudaGetDeviceCount(&count);

paddle/fluid/inference/tensorrt/engine.h

@@ -323,6 +323,10 @@ class TensorRTEngine {
                                            int num_inputs,
                                            plugin::PluginTensorRTV2Ext* plugin);
 
+  nvinfer1::IPluginV2Layer* AddPluginV2IOExt(nvinfer1::ITensor* const* inputs,
+                                             int num_inputs,
+                                             nvinfer1::IPluginV2IOExt* plugin);
+
   void SetTensorDynamicRange(nvinfer1::ITensor* tensor, float range) {
     quant_dynamic_range_[tensor] = range;
   }
@@ -429,6 +433,7 @@ class TensorRTEngine {
   bool with_ernie() { return with_ernie_; }
   bool disable_trt_plugin_fp16() { return disable_trt_plugin_fp16_; }
   bool with_dynamic_shape() { return with_dynamic_shape_; }
+  AnalysisConfig::Precision precision() { return precision_; }
 
 #if IS_TRT_VERSION_GE(6000)
   nvinfer1::IPluginV2Layer* AddDynamicPlugin(
@@ -550,6 +555,7 @@ class TensorRTEngine {
 
   std::vector<std::unique_ptr<plugin::PluginTensorRT>> owned_plugin_;
   std::vector<std::unique_ptr<plugin::PluginTensorRTV2Ext>> owned_plugin_v2ext_;
+  std::vector<std::unique_ptr<nvinfer1::IPluginV2IOExt>> owned_plugin_v2ioext_;
 
   // TensorRT related internal members
   template <typename T>