[Inference] save_optimized_model_pass support tensorrt (#55893)

* fix cudnn 8.7+ bug on cudnnConvolutionBiasActivationForward

* save_optimized_model_pass support tensorrt

* update

* update

* fix compile

* update

* fix ut timeout

Author: yuanlehome

paddle/fluid/framework/ir/auto_mixed_precision_pass.cc

@@ -249,7 +249,7 @@ void AutoMixedPrecisionPass::Init(Graph* graph) const {
     subgraphes_[i] = graph->GetSubGraph(i);
     all_op_nodes_[i] = TopologySortOperations(*subgraphes_[i]);
     VLOG(4) << "subgraph " << i << " has " << all_op_nodes_[i].size()
-            << "op nodes";
+            << " op nodes";
     for (auto* var_node : subgraphes_[i]->Nodes()) {
       if (!var_node->IsVar()) continue;
 

paddle/fluid/framework/naive_executor.cc

@@ -64,10 +64,6 @@ void NaiveExecutor::Run() {
     VLOG(4) << std::this_thread::get_id() << " run "
             << op->DebugStringEx(scope_) << " on scope " << scope_;
     op->SetIsCalledByExecutor(false);
-#ifdef PADDLE_WITH_NVTX
-    platform::CudaNvtxRangePush(op->Type() + "|" + op->OutputVars(true).front(),
-                                platform::NvtxRangeColor::Green);
-#endif
 
     for (auto &func : input_hookfuncs_) {
       func(op.get(), scope_);
@@ -77,7 +73,14 @@ void NaiveExecutor::Run() {
       op->SetOutputHooks(output_hookfuncs_);
     }
 
+#ifdef PADDLE_WITH_NVTX
+    platform::CudaNvtxRangePush(op->Type() + "|" + op->OutputVars(true).front(),
+                                platform::NvtxRangeColor::Green);
+#endif
     op->Run(*scope_, place_);
+#ifdef PADDLE_WITH_NVTX
+    platform::CudaNvtxRangePop();
+#endif
 
     // Update the shared_holder so that only records the max one.
     if (reuse_cache_.count(op.get())) {
@@ -105,9 +108,6 @@ void NaiveExecutor::Run() {
       }
     }
 
-#ifdef PADDLE_WITH_NVTX
-    platform::CudaNvtxRangePop();
-#endif
     for (auto &func : output_hookfuncs_) {
       func(op.get(), scope_);
     }

paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/inference/analysis/passes/ir_graph_build_pass.h"
 
+#include <memory>
 #include <string>
 
 #include "paddle/fluid/framework/executor.h"
@@ -63,8 +64,7 @@ void IrGraphBuildPass::RunImpl(Argument *argument) {
         "set."));
   }
 
-  auto graph = std::unique_ptr<framework::ir::Graph>(
-      new framework::ir::Graph(argument->main_program()));
+  auto graph = std::make_unique<framework::ir::Graph>(argument->main_program());
   argument->SetMainGraph(graph.release());
   auto *scope_ptr = argument->scope_ptr();
   PADDLE_ENFORCE_NOT_NULL(scope_ptr,

paddle/fluid/inference/analysis/passes/ir_graph_build_pass.cc

@@ -24,16 +24,6 @@ namespace inference {
 namespace analysis {
 
 void SaveOptimizedModelPass::SaveOptimizedModel(Argument* argument) {
-  if (!argument->save_optimized_model()) {
-    LOG(WARNING) << "save_optim_cache_model is turned off, skip "
-                    "save_optimized_model_pass";
-    return;
-  }
-  if (!argument->enable_ir_optim()) {
-    LOG(WARNING) << "ir_optim is turned off, skip save_optimized_model_pass";
-    return;
-  }
-
   std::string model_opt_cache_dir = argument->optim_cache_dir();
   if (!model_opt_cache_dir.empty()) {
     if (!PathExists(model_opt_cache_dir)) {
@@ -55,9 +45,11 @@ void SaveOptimizedModelPass::SaveOptimizedModel(Argument* argument) {
   auto* graph = argument->main_graph_ptr();
 
   framework::ProgramDesc optimized_program_desc;
+
   // NOTE(liuyuanle): If the following line of code is not added, an error
   // [SegmentFault] may occur!
   optimized_program_desc.CopyFrom(*argument->main_program().Proto());
+
   framework::ir::GraphToProgram(*graph, &optimized_program_desc);
 
   auto IsPersistable = [](const framework::VarDesc* var) {
@@ -133,11 +125,10 @@ void SaveOptimizedModelPass::SaveOptimizedModel(Argument* argument) {
 }
 
 void SaveOptimizedModelPass::RunImpl(Argument* argument) {
-  // TODO(inference): Support trt.
-  if (argument->use_xpu() ||
-      (argument->use_gpu() && !argument->use_tensorrt())) {
-    SaveOptimizedModel(argument);
+  if (!argument->save_optimized_model() || !argument->enable_ir_optim()) {
+    return;
   }
+  SaveOptimizedModel(argument);
 }
 
 std::string SaveOptimizedModelPass::repr() const {

paddle/fluid/inference/analysis/passes/save_optimized_model_pass.cc

@@ -641,7 +641,7 @@ bool AnalysisPredictor::PrepareProgram(
 }
 
 bool AnalysisPredictor::CreateExecutor() {
-  executor_.reset(new paddle::framework::NaiveExecutor(place_));
+  executor_ = std::make_unique<paddle::framework::NaiveExecutor>(place_);
   return true;
 }
 
@@ -1341,7 +1341,7 @@ bool AnalysisPredictor::GetFetch(std::vector<paddle::Tensor> *outputs,
 void AnalysisPredictor::PrepareArgument() {
   VLOG(3) << "AnalysisPredictor::PrepareArgument";
   // Init std::unique_ptr argument_.
-  argument_.reset(new Argument);
+  argument_ = std::make_unique<Argument>();
   argument_->SetUseGPU(config_.use_gpu());
   argument_->SetUseCutlass(config_.use_cutlass_);
   argument_->SetUseFcPadding(config_.use_fc_padding());
@@ -1570,7 +1570,8 @@ void AnalysisPredictor::PrepareArgument() {
 
   if (!config_.ir_optim()) {
     argument_->SetEnableIrOptim(false);
-    if (config_.enable_gpu_mixed_) {
+    if (config_.enable_gpu_mixed_ &&
+        model_precision_ == phi::DataType::FLOAT32) {
       argument_->SetEnableIrOptim(true);
       pass_builder->ClearPasses();
       pass_builder->AppendPass("auto_mixed_precision_pass");
@@ -1886,6 +1887,10 @@ AnalysisPredictor::GetInputTypes() {
       input_type[name] = paddle_infer::DataType::UINT8;
     } else if (dtype == paddle::framework::proto::VarType::INT8) {
       input_type[name] = paddle_infer::DataType::INT8;
+    } else if (dtype == paddle::framework::proto::VarType::FP64) {
+      input_type[name] = paddle_infer::DataType::FLOAT64;
+    } else if (dtype == paddle::framework::proto::VarType::BOOL) {
+      input_type[name] = paddle_infer::DataType::BOOL;
     } else {
       PADDLE_THROW(paddle::platform::errors::Unimplemented(
           "Unsupported data type `%s` when get input dtype ", dtype));
@@ -2609,7 +2614,7 @@ AnalysisPredictor::~AnalysisPredictor() {
 #ifdef PADDLE_WITH_TENSORRT
   if (config_.trt_engine_memory_sharing()) {
     inference::Singleton<inference::tensorrt::TRTEngineManager>::Global()
-        .releaseContextMemory(predictor_id_);
+        .ReleaseContextMemory(predictor_id_);
   }
 #endif
 }

paddle/fluid/inference/api/analysis_predictor.cc

@@ -167,7 +167,7 @@ class OpConverter {
                                         op_desc.Type()));
 
     it->SetEngine(engine);
-    engine->SetScope(scope);
+    engine->SetScope(&scope);
     it->SetBlockDesc(block);
     (*it)(op, scope, test_mode);
 
@@ -301,7 +301,7 @@ class OpConverter {
       nvinfer1::DataType in_dtype = FluidDataType2TRT(var->GetDataType());
       if (engine->precision() == phi::DataType::FLOAT16 &&
           in_dtype == nvinfer1::DataType::kFLOAT &&
-          engine->EnableLowPrecisionIO()) {
+          engine->LowPrecisionIOEnabled()) {
         in_dtype = nvinfer1::DataType::kHALF;
       }
 
@@ -360,7 +360,7 @@ class OpConverter {
       nvinfer1::DataType out_dtype = FluidDataType2TRT(var->GetDataType());
       if (engine->precision() == phi::DataType::FLOAT16 &&
           out_dtype == nvinfer1::DataType::kFLOAT &&
-          engine->EnableLowPrecisionIO()) {
+          engine->LowPrecisionIOEnabled()) {
         out_dtype = nvinfer1::DataType::kHALF;
       }
       engine->DeclareOutput(output, out_dtype);
@@ -470,7 +470,7 @@ class OpConverter {
       auto shape = newShape->getDimensions();
       shuffle->setReshapeDimensions(shape);
     }
-    if (name != "") {
+    if (!name.empty()) {
       shuffle->setName(name.c_str());
     }
     return shuffle->getOutput(0);
@@ -481,7 +481,7 @@ class OpConverter {
                              const std::string& name = "") {
     auto* shuffle = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
     shuffle->setReshapeDimensions(shape);
-    if (name != "") {
+    if (!name.empty()) {
       shuffle->setName(name.c_str());
     }
     return shuffle->getOutput(0);
@@ -774,11 +774,6 @@ class OpConverter {
   bool test_mode_;
 
  private:
-  // registered op converter map, whose key is the fluid op type, and value is
-  // the pointer position of corresponding OpConverter class.
-  std::unordered_map<std::string, OpConverter*> converters_;
-  // fluid inference scope
-  framework::Scope* scope_{nullptr};
   std::mutex mut_;
 };
 

paddle/fluid/inference/tensorrt/convert/op_converter.h

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include <gtest/gtest.h>  // NOLINT
+#include <memory>
 
 #include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
@@ -95,7 +96,11 @@ TEST(CustomPluginCreater, StaticShapePlugin) {
 
   // init trt engine
   std::unique_ptr<TensorRTEngine> engine_;
-  engine_.reset(new TensorRTEngine(5, 1 << 15));
+
+  TensorRTEngine::ConstructionParams params;
+  params.max_batch_size = 5;
+  params.max_workspace_size = 1 << 15;
+  engine_ = std::make_unique<TensorRTEngine>(params);
   engine_->InitNetwork();
 
   engine_->DeclareInput(
@@ -173,15 +178,10 @@ TEST(CustomPluginCreater, DynamicShapePlugin) {
   std::map<std::string, std::vector<int>> optim_input_shape = {
       {"x", {1, 2, 5, 5}}};
 
-  engine_.reset(new TensorRTEngine(5,
-                                   1 << 15,
-                                   phi::DataType::FLOAT32,
-                                   nullptr,
-                                   0,
-                                   true,
-                                   min_input_shape,
-                                   max_input_shape,
-                                   optim_input_shape));
+  TensorRTEngine::ConstructionParams params;
+  params.max_batch_size = 5;
+  params.max_workspace_size = 1 << 15;
+  engine_ = std::make_unique<TensorRTEngine>(params);
   engine_->InitNetwork();
 
   LOG(INFO) << "with_dynamic_shape " << engine_->with_dynamic_shape();

paddle/fluid/inference/tensorrt/convert/test_custom_plugin_creater.cc

@@ -14,6 +14,8 @@ limitations under the License. */
 
 #include <gtest/gtest.h>  // NOLINT
 
+#include <memory>
+
 #include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
 
@@ -28,7 +30,10 @@ TEST(OpConverter, ConvertBlock) {
 
   // init trt engine
   std::unique_ptr<TensorRTEngine> engine_;
-  engine_.reset(new TensorRTEngine(5, 1 << 15));
+  TensorRTEngine::ConstructionParams params;
+  params.max_batch_size = 5;
+  params.max_workspace_size = 1 << 15;
+  engine_ = std::make_unique<TensorRTEngine>(params);
   engine_->InitNetwork();
 
   engine_->DeclareInput(

paddle/fluid/inference/tensorrt/convert/test_op_converter.cc

@@ -88,7 +88,10 @@ class TRTConvertValidation {
     PADDLE_ENFORCE_EQ(cudaStreamCreate(&stream_),
                       0,
                       platform::errors::External("cudaStreamCreate error."));
-    engine_.reset(new TensorRTEngine(max_batch_size, workspace_size));
+    TensorRTEngine::ConstructionParams params;
+    params.max_batch_size = max_batch_size;
+    params.max_workspace_size = workspace_size;
+    engine_ = std::make_unique<TensorRTEngine>(params);
     engine_->InitNetwork();
   }
 
@@ -155,7 +158,7 @@ class TRTConvertValidation {
     engine_->FreezeNetwork();
 
     // Declare outputs.
-    op_desc_.reset(new framework::OpDesc(desc, nullptr));
+    op_desc_ = std::make_unique<framework::OpDesc>(desc, nullptr);
   }
 
   // We use the set 'neglected_output' here, because some Ops like batch norm,