fix replace_fetch_with_shadow_output_pass

Author: yuanlehome

paddle/fluid/inference/api/analysis_predictor.cc

@@ -1031,7 +1031,14 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
                                             sub_scope_);
     basic_pass_pm.AddPass(std::move(dead_code_elimination_pass));
   }
-  basic_pass_pm.AddPass(::pir::CreateReplaceFetchWithShadowOutputPass());
+  auto replace_fetch_with_shadow_output_pass =
+      ::pir::CreateReplaceFetchWithShadowOutputPass();
+  if (std::find(config_.deleted_passes_.begin(),
+                config_.deleted_passes_.end(),
+                replace_fetch_with_shadow_output_pass->name()) ==
+      config_.deleted_passes_.end()) {
+    basic_pass_pm.AddPass(std::move(replace_fetch_with_shadow_output_pass));
+  }
   if (!config_.glog_info_disabled()) {
     basic_pass_pm.EnablePrintStatistics();
   }
@@ -1090,6 +1097,7 @@ bool AnalysisPredictor::SaveOrLoadPirParameters(bool for_save) {
         std::string fetch_name =
             op->attribute("name").dyn_cast<pir::StrAttribute>().AsString();
         idx2fetches_[idx] = fetch_name;
+        fetch_name2shapes_[fetch_name] = pir::GetShapeFromValue(op->operand_source(0));
       }
     } else if (op->isa<paddle::dialect::DataOp>() ||
                op->isa<paddle::dialect::FeedOp>()) {
@@ -1102,6 +1110,7 @@ bool AnalysisPredictor::SaveOrLoadPirParameters(bool for_save) {
       feed_names_[data_name] = feed_idx;
       feed_idx++;
       pir_feeds_.emplace_back(op);
+      feed_name2shapes_[data_name] = pir::GetShapeFromValue(op->result(0));
     }
 
     if (op->isa<::pir::ParameterOp>()) {
@@ -2545,6 +2554,9 @@ std::vector<std::string> AnalysisPredictor::GetInputNames() {
 
 std::map<std::string, std::vector<int64_t>>
 AnalysisPredictor::GetInputTensorShape() {
+  if (load_pir_model_) {
+    return feed_name2shapes_;
+  }
   std::map<std::string, std::vector<int64_t>> input_shapes;
   std::vector<std::string> names = GetInputNames();
   for (std::string const &name : names) {
@@ -2604,6 +2616,9 @@ std::vector<std::string> AnalysisPredictor::GetOutputNames() {
 
 std::map<std::string, std::vector<int64_t>>
 AnalysisPredictor::GetOutputTensorShape() {
+  if (load_pir_model_) {
+    return fetch_name2shapes_;
+  }
   std::map<std::string, std::vector<int64_t>> output_shapes;
   std::vector<std::string> names = GetOutputNames();
   for (std::string const &name : names) {

paddle/fluid/inference/api/analysis_predictor.h

@@ -578,9 +578,11 @@ class AnalysisPredictor : public PaddlePredictor {
   std::map<std::string, size_t> feed_names_;
   // Sorted according to the idx.
   std::map<size_t, std::string> idx2feeds_;
+  std::map<std::string, std::vector<int64_t>> feed_name2shapes_;
   std::vector<framework::OpDesc *> fetches_;
   std::vector<pir::Operation *> pir_fetches_;
   std::map<size_t, std::string> idx2fetches_;
+  std::map<std::string, std::vector<int64_t>> fetch_name2shapes_;
 
   phi::DataType model_precision_{phi::DataType::FLOAT32};
 

paddle/fluid/pir/transforms/general/inplace_pass.cc

@@ -88,6 +88,10 @@ bool CanDoInplace(const std::unordered_set<pir::Value>& eager_dels,
     return false;
   }
 
+  // for (auto it = output.use_begin(); it != output.use_end(); ++it) {
+  //   if(it->owner()->isa<pir::ShadowOutputOp>() || )
+  // }
+
   if (input.type().isa<TensorType>() && output.type().isa<TensorType>()) {
     auto input_alloc_tensor_type = input.type().dyn_cast<TensorType>();
     auto output_alloc_tensor_type = output.type().dyn_cast<TensorType>();

paddle/fluid/pir/transforms/general/replace_fetch_with_shadow_output_pass.cc

@@ -17,6 +17,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/fluid/pir/utils/general_functions.h"
 #include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/core/ir_context.h"
 #include "paddle/pir/include/pass/pass.h"
 #include "paddle/pir/include/pass/pass_registry.h"
 
@@ -29,13 +30,20 @@ class ReplaceFetchWithShadowOutputPattern
   bool MatchAndRewrite(
       paddle::dialect::FetchOp op,
       pir::PatternRewriter& rewriter) const override {  // NOLINT
+    // for pd_op.data -> value -> pd_op.fetch, we insert pd_op.scale before
+    // pd_op.fetch to solve error likes [what(): (NotFound) Variable 'xxx' is
+    // not found in scope.]
     if (pir::GetDefiningOpForInput(op, 0)->HasAttribute("name")) {
-      // DataOp/FeedOp
-      return false;
+      auto scale_op = rewriter.Build<paddle::dialect::ScaleOp>(
+          op->operand_source(0), 1.0, 0.0, true);
+      rewriter.Build<pir::ShadowOutputOp>(
+          scale_op->result(0),
+          op->attributes().at("name").dyn_cast<pir::StrAttribute>().AsString());
+    } else {
+      rewriter.Build<pir::ShadowOutputOp>(
+          op->operand_source(0),
+          op->attributes().at("name").dyn_cast<pir::StrAttribute>().AsString());
     }
-    rewriter.Build<pir::ShadowOutputOp>(
-        op->operand_source(0),
-        op->attributes().at("name").dyn_cast<pir::StrAttribute>().AsString());
     rewriter.EraseOp(op);
     return true;
   }