[Inference]Fix PIR-TRT bugs in PaddleX Part-2 (#71233)

YuanRisheng · web-flow · commit 56c7f0632be2 · 2025-02-24T13:37:24.000+08:00
* fix pir-trt bugs

* fix coverage
diff --git a/paddle/fluid/framework/new_executor/collect_shape_manager.cc b/paddle/fluid/framework/new_executor/collect_shape_manager.cc
@@ -29,23 +29,22 @@ void CollectShapeManager::CollectShapeInfo(
     framework::Scope *scope) {
   std::lock_guard<std::mutex> lock(info_mutex_);
   is_shape_range_info_ready_ = false;
-  for (auto &output : instr->Outputs()) {
-    auto var_name = value_exe_info->GetVarName(output.first);
+  for (auto &input : instr->Inputs()) {
+    if (!op_value2instr_id_.count(input.first)) {
+      // Because the input value maybe same between different ops.
+      // To prevent duplicate shape collection, we only select one op for
+      // getting shape of value
+      op_value2instr_id_[input.first] = instr->Id();
+    }
+    if (op_value2instr_id_[input.first] != instr->Id()) {
+      continue;
+    }
+    auto var_name = value_exe_info->GetVarName(input.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(output.first)) {
+    if (!tensor.initialized() && !instr->NoNeedBuffer().count(input.first)) {
       continue;
     }
     paddle::platform::DeviceContextPool &pool =
@@ -65,7 +64,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_[output.first].emplace_back(shape);
+      shape_info_[input.first].emplace_back(shape);
     } else if (tensor.numel() > 0) {
       // This must be a zero dimension tensor.
       PADDLE_ENFORCE_EQ(tensor.numel(),
@@ -74,7 +73,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_[output.first].emplace_back(zero_shape);
+      shape_info_[input.first].emplace_back(zero_shape);
     }
 
     // We need collect value range for shape tensor for Paddle-TRT's use.
@@ -120,7 +119,7 @@ void CollectShapeManager::CollectShapeInfo(
                              nullptr);
 #endif
       }
-      shape_tensor_info_[output.first].emplace_back(int32_host);
+      shape_tensor_info_[input.first].emplace_back(int32_host);
     }
   }
 }
diff --git a/paddle/fluid/framework/new_executor/collect_shape_manager.h b/paddle/fluid/framework/new_executor/collect_shape_manager.h
@@ -65,12 +65,14 @@ class CollectShapeManager {
     max_values_.clear();
     opt_values_.clear();
     op_value2kernel_value_.clear();
+    op_value2instr_id_.clear();
     is_shape_range_info_ready_ = false;
   }
 
  private:
   CollectShapeManager() {}
   std::unordered_map<pir::Value, pir::Value> op_value2kernel_value_;
+  std::unordered_map<pir::Value, size_t> op_value2instr_id_;
   std::map<pir::Value, std::vector<std::vector<int32_t>>> shape_info_;
   std::map<pir::Value, std::vector<std::vector<int32_t>>> shape_tensor_info_;
   std::map<pir::Value, std::vector<int32_t>> min_shapes_;
diff --git a/paddle/fluid/framework/new_executor/garbage_collector/event_garbage_collector.cc b/paddle/fluid/framework/new_executor/garbage_collector/event_garbage_collector.cc
@@ -133,6 +133,7 @@ void InterpreterCoreEventGarbageCollector::Add(Variable* var,
     for (auto& t : *tensor_arr) {
       Add(t.MoveMemoryHolder(), event, ctx);
     }
+    tensor_arr->clear();
   } else if (var->IsType<std::vector<Scope*>>()) {
     // NOTE(@xiongkun03) conditional_op / while_op will create a STEP_SCOPE
     // refer to executor.cc to see what old garbage collector does.
diff --git a/paddle/fluid/framework/new_executor/garbage_collector/fast_garbage_collector.cc b/paddle/fluid/framework/new_executor/garbage_collector/fast_garbage_collector.cc
@@ -50,6 +50,7 @@ void InterpreterCoreFastGarbageCollector::Add(Variable* var) {
     for (auto& t : *tensor_arr) {
       Add(t.MoveMemoryHolder());
     }
+    tensor_arr->clear();
   } else if (var->IsType<phi::SparseCooTensor>()) {
     Add(var->GetMutable<phi::SparseCooTensor>()
             ->mutable_indices()
diff --git a/paddle/fluid/framework/new_executor/garbage_collector/no_event_garbage_collector.cc b/paddle/fluid/framework/new_executor/garbage_collector/no_event_garbage_collector.cc
@@ -92,6 +92,7 @@ void InterpreterCoreNoEventGarbageCollector::Add(
     for (auto& t : *tensor_arr) {
       Add(t.MoveMemoryHolder(), ctx);
     }
+    tensor_arr->clear();
   } else if (var->IsType<std::vector<Scope*>>()) {
     // NOTE(@xiongkun03) conditional_op / while_op will create a STEP_SCOPE
     // refer to executor.cc to see what old garbage collector does.
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -1964,6 +1964,11 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {
       }
     }
 
+    if (FLAGS_enable_collect_shape) {
+      CollectShapeManager::Instance().CollectShapeInfo(
+          instr_node, value_exe_info_.get(), scope_);
+    }
+
     if (!instr_node->IsArtificial()) {
       {
         phi::RecordEvent record(
@@ -1979,10 +1984,7 @@ 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());
       }
diff --git a/python/paddle/tensorrt/impls/ops.py b/python/paddle/tensorrt/impls/ops.py
@@ -91,7 +91,8 @@ def roi_align_converter(network, paddle_op, inputs):
         ),
         trt.PluginField(
             "aligned",
-            np.array(aligned, dtype=np.bool),
+            np.array(aligned, dtype=np.bool_),
+            trt.PluginFieldType.INT32,
         ),
     ]
     plugin_field_collection = trt.PluginFieldCollection(plugin_fields)

Original file line number	Diff line number	Diff line change
`@@ -133,6 +133,7 @@ void InterpreterCoreEventGarbageCollector::Add(Variable* var,`
`133`	`133`	`for (auto& t : *tensor_arr) {`
`134`	`134`	`Add(t.MoveMemoryHolder(), event, ctx);`
`135`	`135`	`}`
	`136`	`+ tensor_arr->clear();`
`136`	`137`	`} else if (var->IsType<std::vector<Scope*>>()) {`
`137`	`138`	`// NOTE(@xiongkun03) conditional_op / while_op will create a STEP_SCOPE`
`138`	`139`	`// refer to executor.cc to see what old garbage collector does.`
Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,7 @@ void InterpreterCoreFastGarbageCollector::Add(Variable* var) {`
`50`	`50`	`for (auto& t : *tensor_arr) {`
`51`	`51`	`Add(t.MoveMemoryHolder());`
`52`	`52`	`}`
	`53`	`+ tensor_arr->clear();`
`53`	`54`	`} else if (var->IsType<phi::SparseCooTensor>()) {`
`54`	`55`	`Add(var->GetMutable<phi::SparseCooTensor>()`
`55`	`56`	`->mutable_indices()`
Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,7 @@ void InterpreterCoreNoEventGarbageCollector::Add(`
`92`	`92`	`for (auto& t : *tensor_arr) {`
`93`	`93`	`Add(t.MoveMemoryHolder(), ctx);`
`94`	`94`	`}`
	`95`	`+ tensor_arr->clear();`
`95`	`96`	`} else if (var->IsType<std::vector<Scope*>>()) {`
`96`	`97`	`// NOTE(@xiongkun03) conditional_op / while_op will create a STEP_SCOPE`
`97`	`98`	`// refer to executor.cc to see what old garbage collector does.`
Original file line number	Diff line number	Diff line change
`@@ -1964,6 +1964,11 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {`
`1964`	`1964`	`}`
`1965`	`1965`	`}`
`1966`	`1966`
	`1967`	`+ if (FLAGS_enable_collect_shape) {`
	`1968`	`+ CollectShapeManager::Instance().CollectShapeInfo(`
	`1969`	`+ instr_node, value_exe_info_.get(), scope_);`
	`1970`	`+ }`
	`1971`	`+`
`1967`	`1972`	`if (!instr_node->IsArtificial()) {`
`1968`	`1973`	`{`
`1969`	`1974`	`phi::RecordEvent record(`
`@@ -1979,10 +1984,7 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {`
`1979`	`1984`	`<< "): context wait and get last error";`
`1980`	`1985`	`#endif`
`1981`	`1986`	`}`
`1982`		`- if (FLAGS_enable_collect_shape) {`
`1983`		`- CollectShapeManager::Instance().CollectShapeInfo(`
`1984`		`- instr_node, value_exe_info_.get(), scope_);`
`1985`		`- }`
	`1987`	`+`
`1986`	`1988`	`if (FLAGS_check_nan_inf) {`
`1987`	`1989`	`CheckTensorHasNanOrInf(instr_node, scope_, value_exe_info_.get());`
`1988`	`1990`	`}`