[Inference] support collect shape in sub block (#60451)

yuanlehome · web-flow · commit c7a3f63d19ef · 2024-01-03T11:22:25.000+08:00
* support collect shape in sub block

* udpate

* udpate
diff --git a/paddle/fluid/framework/naive_executor.cc b/paddle/fluid/framework/naive_executor.cc
@@ -101,8 +101,9 @@ void NaiveExecutor::Run() {
       func(op.get(), scope_);
     }
 
-    if (op->Type() == "while") {
+    if (op->Type() == "while" || op->Type() == "conditional_block") {
       op->SetOutputHooks(output_hookfuncs_);
+      op->SetInputHooks(input_hookfuncs_);
     }
 
 #ifdef PADDLE_WITH_NVTX
diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
@@ -635,16 +635,12 @@ void BuildOpFuncList(const platform::Place& place,
         hook(op, local_scope);
       }
 
-      if (op->Type() == "while") {
+      if (op->Type() == "while" || op->Type() == "conditional_block") {
         op->SetInputHooks(input_hookfuncs);
         op->SetOutputHooks(output_hookfuncs);
         auto runtime_attrs = op->RuntimeAttrs();
         runtime_attrs.insert(std::make_pair("used_for_inference", true));
         op->SetRuntimeAttributeMap(runtime_attrs);
-      } else if (op->Type() == "conditional_block") {
-        auto runtime_attrs = op->RuntimeAttrs();
-        runtime_attrs.insert(std::make_pair("used_for_inference", true));
-        op->SetRuntimeAttributeMap(runtime_attrs);
       }
     }
 
diff --git a/paddle/fluid/framework/new_executor/program_interpreter.cc b/paddle/fluid/framework/new_executor/program_interpreter.cc
@@ -915,16 +915,12 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {
       hook(op, local_scope);
     }
 
-    if (op->Type() == "while") {
+    if (op->Type() == "while" || op->Type() == "conditional_block") {
       op->SetInputHooks(input_hookfuncs_);
       op->SetOutputHooks(output_hookfuncs_);
       auto runtime_attrs = op->RuntimeAttrs();
       runtime_attrs.insert(std::make_pair("used_for_inference", true));
       op->SetRuntimeAttributeMap(runtime_attrs);
-    } else if (op->Type() == "conditional_block") {
-      auto runtime_attrs = op->RuntimeAttrs();
-      runtime_attrs.insert(std::make_pair("used_for_inference", true));
-      op->SetRuntimeAttributeMap(runtime_attrs);
     }
   }
 
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -2391,9 +2391,15 @@ bool AnalysisPredictor::ExpRunWithExternalStream(const gpuStream_t stream) {
 #endif
 
 void AnalysisPredictor::HookCollectShapeRangeInfo() {
+  if (config_.new_executor_enabled()) {
+    LOG_FIRST_N(WARNING, 1)
+        << "When collecting shapes, it is recommended to run multiple loops to "
+           "obtain more accurate shape information.";
+  }
+
   auto hook = [&](const std::string &op_type,
                   const std::string &input_name,
-                  const paddle::Tensor &var) -> void {
+                  const paddle::Tensor &input_tensor) -> void {
     paddle::platform::DeviceContextPool &pool =
         paddle::platform::DeviceContextPool::Instance();
     if (config_.use_gpu()) {
@@ -2408,26 +2414,22 @@ void AnalysisPredictor::HookCollectShapeRangeInfo() {
 #endif
     }
 
-    auto *new_var = sub_scope_->GetVar(input_name);
-    if (!new_var) return;
-    if (!new_var->IsType<phi::DenseTensor>()) {
-      return;
-    }
-    auto tensor = new_var->Get<phi::DenseTensor>();
-    if (!tensor.initialized()) return;
-    framework::DDim dim = tensor.dims();
+    if (!input_tensor.is_dense_tensor()) return;
+    auto tensor =
+        std::dynamic_pointer_cast<phi::DenseTensor>(input_tensor.impl()).get();
+    framework::DDim dim = tensor->dims();
     std::vector<int32_t> shape(dim.size());
     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_name].emplace_back(shape);
-    } else if (tensor.numel() > 0) {
+    } else if (tensor->numel() > 0) {
       // This must be a zero dimension tensor.
-      PADDLE_ENFORCE_EQ(tensor.numel(),
+      PADDLE_ENFORCE_EQ(tensor->numel(),
                         1UL,
                         platform::errors::PreconditionNotMet(
                             "This tensor must have one element, but got %ld.",
-                            tensor.numel()));
+                            tensor->numel()));
       std::vector<int32_t> zero_shape(1, 1);
       shape_info_[input_name].emplace_back(zero_shape);
     }
@@ -2437,34 +2439,34 @@ void AnalysisPredictor::HookCollectShapeRangeInfo() {
     // assumption that all shape tensors in the model have numbers <= 8.
     // This is a simple method to identify all shape tensors with some
     // mistakes, but it doesn't matter.
-    auto is_shape_tensor = tensor.numel() <= 8 && tensor.numel() >= 1;
-    if ((tensor.dtype() == phi::DataType::INT32 ||
-         tensor.dtype() == phi::DataType::INT64) &&
+    auto is_shape_tensor = tensor->numel() <= 8 && tensor->numel() >= 1;
+    if ((tensor->dtype() == phi::DataType::INT32 ||
+         tensor->dtype() == phi::DataType::INT64) &&
         is_shape_tensor) {
-      std::vector<int> int32_host(tensor.numel());
+      std::vector<int> int32_host(tensor->numel());
 
-      if (platform::is_cpu_place(tensor.place())) {
-        auto &int32_tensor = tensor;
-        if (tensor.dtype() == phi::DataType::INT64) {
+      if (platform::is_cpu_place(tensor->place())) {
+        auto &int32_tensor = *tensor;
+        if (tensor->dtype() == phi::DataType::INT64) {
           auto *cpu_ctx = pool.Get(platform::CPUPlace());
           int32_tensor = phi::funcs::TransDataType(
               reinterpret_cast<const phi::CPUContext &>(*cpu_ctx),
-              tensor,
+              *tensor,
               DataType::INT32);
         }
         paddle::memory::Copy(platform::CPUPlace(),
                              int32_host.data(),
                              platform::CPUPlace(),
                              int32_tensor.data<int>(),
                              int32_tensor.numel() * sizeof(int));
-      } else if (platform::is_gpu_place(tensor.place())) {
+      } else if (platform::is_gpu_place(tensor->place())) {
 #if defined(PADDLE_WITH_CUDA)
-        auto *dev_ctx = pool.Get(tensor.place());
-        auto &int32_tensor = tensor;
-        if (tensor.dtype() == phi::DataType::INT64) {
+        auto *dev_ctx = pool.Get(tensor->place());
+        auto &int32_tensor = *tensor;
+        if (tensor->dtype() == phi::DataType::INT64) {
           int32_tensor = phi::funcs::TransDataType(
               reinterpret_cast<const phi::GPUContext &>(*dev_ctx),
-              tensor,
+              *tensor,
               DataType::INT32);
         }
         paddle::memory::Copy(platform::CPUPlace(),
diff --git a/paddle/fluid/operators/controlflow/conditional_block_op.cc b/paddle/fluid/operators/controlflow/conditional_block_op.cc
@@ -116,6 +116,8 @@ class ConditionalBlockOp : public ConditionalOp {
 #endif
         core_.reset(new InterpreterCore(
             dev_place, *block, &cur_scope, execution_config));
+        core_->SetOutputHooks(output_hookfuncs_);
+        core_->SetInputHooks(input_hookfuncs_);
         VLOG(10) << "[interpreterCore] created:" << core_;
       } else {
         BuildScopeForControlFlowOp(*core_, *block, &cur_scope);

Original file line number	Diff line number	Diff line change
`@@ -101,8 +101,9 @@ void NaiveExecutor::Run() {`
`101`	`101`	`func(op.get(), scope_);`
`102`	`102`	`}`
`103`	`103`
`104`		`- if (op->Type() == "while") {`
	`104`	`+ if (op->Type() == "while" \|\| op->Type() == "conditional_block") {`
`105`	`105`	`op->SetOutputHooks(output_hookfuncs_);`
	`106`	`+ op->SetInputHooks(input_hookfuncs_);`
`106`	`107`	`}`
`107`	`108`
`108`	`109`	`#ifdef PADDLE_WITH_NVTX`
Original file line number	Diff line number	Diff line change
`@@ -635,16 +635,12 @@ void BuildOpFuncList(const platform::Place& place,`
`635`	`635`	`hook(op, local_scope);`
`636`	`636`	`}`
`637`	`637`
`638`		`- if (op->Type() == "while") {`
	`638`	`+ if (op->Type() == "while" \|\| op->Type() == "conditional_block") {`
`639`	`639`	`op->SetInputHooks(input_hookfuncs);`
`640`	`640`	`op->SetOutputHooks(output_hookfuncs);`
`641`	`641`	`auto runtime_attrs = op->RuntimeAttrs();`
`642`	`642`	`runtime_attrs.insert(std::make_pair("used_for_inference", true));`
`643`	`643`	`op->SetRuntimeAttributeMap(runtime_attrs);`
`644`		`- } else if (op->Type() == "conditional_block") {`
`645`		`- auto runtime_attrs = op->RuntimeAttrs();`
`646`		`- runtime_attrs.insert(std::make_pair("used_for_inference", true));`
`647`		`- op->SetRuntimeAttributeMap(runtime_attrs);`
`648`	`644`	`}`
`649`	`645`	`}`
`650`	`646`
Original file line number	Diff line number	Diff line change
`@@ -915,16 +915,12 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {`
`915`	`915`	`hook(op, local_scope);`
`916`	`916`	`}`
`917`	`917`
`918`		`- if (op->Type() == "while") {`
	`918`	`+ if (op->Type() == "while" \|\| op->Type() == "conditional_block") {`
`919`	`919`	`op->SetInputHooks(input_hookfuncs_);`
`920`	`920`	`op->SetOutputHooks(output_hookfuncs_);`
`921`	`921`	`auto runtime_attrs = op->RuntimeAttrs();`
`922`	`922`	`runtime_attrs.insert(std::make_pair("used_for_inference", true));`
`923`	`923`	`op->SetRuntimeAttributeMap(runtime_attrs);`
`924`		`- } else if (op->Type() == "conditional_block") {`
`925`		`- auto runtime_attrs = op->RuntimeAttrs();`
`926`		`- runtime_attrs.insert(std::make_pair("used_for_inference", true));`
`927`		`- op->SetRuntimeAttributeMap(runtime_attrs);`
`928`	`924`	`}`
`929`	`925`	`}`
`930`	`926`