refine shadow feed (#57291)

wanghuancoder · web-flow · commit 5dd02fbf111d · 2023-09-14T09:39:44.000+08:00
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -1009,13 +1009,13 @@ void AddShadowFeed(
     std::unordered_map<pir::Value, pir::OpResult>* map_value_pair) {
   bool feed_op_add_shadow_feed =
       (op_item->name() == "pd_op.feed") && platform::is_gpu_place(place);
-  bool data_op_add_shadow_feed = (op_item->name() == "pd_op.data") &&
-                                 platform::is_gpu_place(place) &&
-                                 (kernel_op->attributes()
-                                      .at("place")
-                                      .dyn_cast<dialect::PlaceAttribute>()
-                                      .data()
-                                      .GetType() != phi::AllocationType::GPU);
+  bool data_op_add_shadow_feed =
+      (op_item->name() == "pd_op.data") && platform::is_gpu_place(place) &&
+      (kernel_op->attributes()
+           .at("place")
+           .dyn_cast<dialect::PlaceAttribute>()
+           .data()
+           .GetType() == phi::AllocationType::UNDEFINED);
   bool add_shadow_feed = feed_op_add_shadow_feed || data_op_add_shadow_feed;
   if (add_shadow_feed) {
     // if shadow data op place not gpu,add shadow feed op
