[Dy2St] Remove ShadowFeedOp for DataOp downstream (#73337)

0x3878f · web-flow · commit 8ebe012b4186 · 2025-06-15T15:49:11.000+08:00
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -468,6 +468,9 @@ inline void PirRunProgramAPI(
     // Step 3. create new interpretercore
     auto passed_kernel_program = paddle::framework::ApplyIrPass(
         forward_program.get(), place, no_need_buffer_name_set);
+    const auto &new_block = passed_kernel_program->block();
+    passed_kernel_program = paddle::framework::ApplyRemoveShadowFeedPass(
+        std::move(passed_kernel_program), new_block, place, global_inner_scope);
     interpreter_core = paddle::framework::CreatePirInterpreterCoreInfoToCache(
         std::move(passed_kernel_program),
         place,
diff --git a/paddle/fluid/pir/transforms/general/remove_shadow_feed_pass.cc b/paddle/fluid/pir/transforms/general/remove_shadow_feed_pass.cc
@@ -78,11 +78,24 @@ class RemoveShadowFeedPattern
   bool IsSamePlaceShadowFeed(paddle::dialect::PhiKernelOp op) const {
     if (op.op_name() == "pd_op.shadow_feed") {
       auto in = op.operand_source(0);
-      if (!kwargs_map_.count(in)) {
-        return false;
-      }
-      auto in_name = kwargs_map_.at(in);
-      auto *var = scope_->FindVar(in_name);
+      auto *var = [&]() -> paddle::framework::Variable * {
+        auto *defined_op = in.defining_op();
+        if (defined_op && defined_op->isa<paddle::dialect::PhiKernelOp>()) {
+          if (defined_op->dyn_cast<paddle::dialect::PhiKernelOp>()
+                  .kernel_name() != "data")
+            return nullptr;
+          const auto &name = defined_op->attributes()
+                                 .at("name")
+                                 .dyn_cast<pir::StrAttribute>()
+                                 .AsString();
+          return scope_->FindVar(name);
+        }
+        if (kwargs_map_.count(in)) {
+          const auto &name = kwargs_map_.at(in);
+          return scope_->FindVar(name);
+        }
+        return nullptr;
+      }();
       if (!var) {
         return false;
       }
