[Inference]Fix PIR-TRT bugs in PaddleX (#71183)

YuanRisheng · web-flow · commit 2622dcacb24b · 2025-02-21T11:35:03.000+08:00
* fix trt model bugs

* fix bugs

* fix inplace bugs
diff --git a/paddle/fluid/framework/new_executor/collect_shape_manager.cc b/paddle/fluid/framework/new_executor/collect_shape_manager.cc
@@ -29,12 +29,23 @@ void CollectShapeManager::CollectShapeInfo(
     framework::Scope *scope) {
   std::lock_guard<std::mutex> lock(info_mutex_);
   is_shape_range_info_ready_ = false;
-  for (auto &input : instr->Inputs()) {
-    auto var_name = value_exe_info->GetVarName(input.first);
+  for (auto &output : instr->Outputs()) {
+    auto var_name = value_exe_info->GetVarName(output.first);
     auto *var = scope->FindVar(var_name);
     if (!var || !var->IsType<phi::DenseTensor>()) continue;
+
+    // inplace var in kernel program has same value between input and output,
+    // we should use ‘continue’ to prevent duplicate data entries.
+    bool is_inplace_var = false;
+    for (const auto &inplace_var_pair : instr->InplaceInfo()) {
+      if (inplace_var_pair.first == var) {
+        is_inplace_var = true;
+      }
+    }
+    if (is_inplace_var) continue;
+
     auto tensor = var->Get<phi::DenseTensor>();
-    if (!tensor.initialized() && !instr->NoNeedBuffer().count(input.first)) {
+    if (!tensor.initialized() && !instr->NoNeedBuffer().count(output.first)) {
       continue;
     }
     paddle::platform::DeviceContextPool &pool =
@@ -54,7 +65,7 @@ void CollectShapeManager::CollectShapeInfo(
     for (int i = 0; i < static_cast<int>(shape.size()); ++i)
       shape[i] = static_cast<int32_t>(dim[i]);
     if (!shape.empty()) {
-      shape_info_[input.first].emplace_back(shape);
+      shape_info_[output.first].emplace_back(shape);
     } else if (tensor.numel() > 0) {
       // This must be a zero dimension tensor.
       PADDLE_ENFORCE_EQ(tensor.numel(),
@@ -63,7 +74,7 @@ void CollectShapeManager::CollectShapeInfo(
                             "This tensor must have one element, but got %ld.",
                             tensor.numel()));
       std::vector<int32_t> zero_shape(1, 1);
-      shape_info_[input.first].emplace_back(zero_shape);
+      shape_info_[output.first].emplace_back(zero_shape);
     }
 
     // We need collect value range for shape tensor for Paddle-TRT's use.
@@ -109,7 +120,7 @@ void CollectShapeManager::CollectShapeInfo(
                              nullptr);
 #endif
       }
-      shape_tensor_info_[input.first].emplace_back(int32_host);
+      shape_tensor_info_[output.first].emplace_back(int32_host);
     }
   }
 }
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -1958,11 +1958,6 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {
             << "Before: " << cur_place << " "
             << instr_node->DebugStringEx(scope_, value_exe_info_.get());
 
-    if (FLAGS_enable_collect_shape) {
-      CollectShapeManager::Instance().CollectShapeInfo(
-          instr_node, value_exe_info_.get(), scope_);
-    }
-
     if (execution_config_.used_for_inference) {
       for (auto& hook : pir_input_hookfuncs_) {
         hook(instr_node, value_exe_info_.get(), scope_);
@@ -1984,6 +1979,10 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {
                 << "): context wait and get last error";
 #endif
       }
+      if (FLAGS_enable_collect_shape) {
+        CollectShapeManager::Instance().CollectShapeInfo(
+            instr_node, value_exe_info_.get(), scope_);
+      }
       if (FLAGS_check_nan_inf) {
         CheckTensorHasNanOrInf(instr_node, scope_, value_exe_info_.get());
       }
@@ -2009,6 +2008,7 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {
     }
 
     VLOG(5) << "after run kernel";
+
     instr_node->RecordEvent(cur_place);
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     if (enable_job_schedule_profiler_) {
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -66,7 +66,6 @@ DEFINE_GENERAL_PATTERN(Layer_norm, paddle::dialect::LayerNormOp)
 DEFINE_GENERAL_PATTERN(Add, paddle::dialect::AddOp)
 DEFINE_GENERAL_PATTERN(Full, paddle::dialect::FullOp)
 DEFINE_GENERAL_PATTERN(Silu, paddle::dialect::SiluOp)
-DEFINE_GENERAL_PATTERN(Conv2d, paddle::dialect::Conv2dOp)
 DEFINE_GENERAL_PATTERN(FusedConv2dAddAct, paddle::dialect::FusedConv2dAddActOp)
 DEFINE_GENERAL_PATTERN(DepthwiseConv2d, paddle::dialect::DepthwiseConv2dOp)
 DEFINE_GENERAL_PATTERN(Shape, paddle::dialect::ShapeOp)
@@ -460,6 +459,22 @@ class Pool2dOpPattern
   }
 };
 
+class Conv2dOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::Conv2dOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::Conv2dOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::Conv2dOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    auto filter_define_op = pir::GetDefiningOpForInput(op, 1);
+    if (filter_define_op->name() != "builtin.parameter" &&
+        filter_define_op->name() != "builtin.constant") {
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class Conv2dTransposeOpPattern
     : public pir::OpRewritePattern<paddle::dialect::Conv2dTransposeOp> {
  public:
@@ -2822,6 +2837,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
 #endif
 #undef ADD_PATTERN
     ps.Add(std::make_unique<Pool2dOpPattern>(context));
+    ps.Add(std::make_unique<Conv2dOpPattern>(context));
     ps.Add(std::make_unique<Conv2dTransposeOpPattern>(context));
     ps.Add(std::make_unique<DepthwiseConv2dTransposeOpPattern>(context));
     ps.Add(std::make_unique<DeformableConvOpPattern>(context));
