PaddlePaddle
diff --git a/‎paddle/fluid/framework/executor_cache.cc
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions b/‎paddle/fluid/framework/executor_cache.cc
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+25-27Lines changed: 25 additions & 27 deletions b/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+25-27Lines changed: 25 additions & 27 deletions
diff --git a/‎paddle/phi/api/lib/data_transform.h
Copy file name to clipboard
+19Lines changed: 19 additions & 0 deletions b/‎paddle/phi/api/lib/data_transform.h
Copy file name to clipboard
+19Lines changed: 19 additions & 0 deletions
diff --git a/‎paddle/phi/common/backend.h
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions b/‎paddle/phi/common/backend.h
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions
diff --git a/‎paddle/phi/common/place.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions b/‎paddle/phi/common/place.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎paddle/phi/common/place.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/phi/common/place.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/cpu/data_kernel.cc renamed to ‎paddle/phi/kernels/data_kernel.cc
Copy file name to clipboard
+9-48Lines changed: 9 additions & 48 deletions b/‎paddle/phi/kernels/cpu/data_kernel.cc renamed to ‎paddle/phi/kernels/data_kernel.cc
Copy file name to clipboard
+9-48Lines changed: 9 additions & 48 deletions
diff --git a/‎paddle/phi/kernels/data_kernel.h
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/phi/kernels/data_kernel.h
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
@@ -247,6 +247,11 @@ std::unique_ptr<::pir::Program> ConstructForwardIrProgram(
     // TODO(phlrain) : using tensor dtype
     op_desc->SetAttr("dtype", 0);
     op_desc->SetAttr("place", static_cast<int>(p));
+    if (p == phi::AllocationType::CUSTOM) {
+      op_desc->SetAttr("place_device_id", in_t.place().GetDeviceId());
+      op_desc->SetAttr("place_device_type", in_t.place().GetDeviceType());
+    }
+
     op_desc->SetAttr("name", name);
     op_desc->SetOutput("out", {name});
   }
@@ -264,6 +269,10 @@ std::unique_ptr<::pir::Program> ConstructForwardIrProgram(
     // TODO(phlrain) : using tensor dtype
     op_desc->SetAttr("dtype", 0);
     op_desc->SetAttr("place", static_cast<int>(p));
+    if (p == phi::AllocationType::CUSTOM) {
+      op_desc->SetAttr("place_device_id", param.place().GetDeviceId());
+      op_desc->SetAttr("place_device_type", param.place().GetDeviceType());
+    }
     op_desc->SetAttr("name", name);
     op_desc->SetOutput("out", {name});
 
@@ -344,6 +353,11 @@ std::unique_ptr<::pir::Program> ConstructBackwardIrProgram(
       // TODO(phlrain) : using tensor dtype
       op_desc->SetAttr("dtype", 0);
       op_desc->SetAttr("place", static_cast<int>(p));
+      if (p == phi::AllocationType::CUSTOM) {
+        op_desc->SetAttr("place_device_id", tensor.place().GetDeviceId());
+        op_desc->SetAttr("place_device_type", tensor.place().GetDeviceType());
+      }
+
       op_desc->SetAttr("name", var_name);
       op_desc->SetOutput("out", {var_name});
     }
 
@@ -1527,6 +1527,19 @@ struct DataOpTranscriber : public FeedOpTranscriber {
              ctx, phi::Place(static_cast<phi::AllocationType>(allocate_type)))},
     };
 
+    if (static_cast<phi::AllocationType>(allocate_type) ==
+        phi::AllocationType::CUSTOM) {
+      int place_device_id =
+          PADDLE_GET_CONST(int, op_desc.GetAttr("place_device_id"));
+      std::string place_device_type =
+          PADDLE_GET_CONST(std::string, op_desc.GetAttr("place_device_type"));
+      attribute_map["place"] = paddle::dialect::PlaceAttribute::get(
+          ctx,
+          phi::Place(static_cast<phi::AllocationType>(allocate_type),
+                     place_device_id,
+                     place_device_type));
+    }
+
     return attribute_map;
   }
 };
 
@@ -211,6 +211,7 @@ static bool NeedFallBackCpu(const pir::Operation* op,
   return false;
 }
 
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 static bool NeedFallBackFromGPUDNN2GPU(pir::Operation* op,
                                        const std::string& kernel_name,
                                        const phi::KernelKey kernel_key) {
@@ -275,6 +276,7 @@ static bool NeedFallBackFromGPUDNN2GPU(pir::Operation* op,
 
   return false;
 }
+#endif
 
 bool CanRunOnCpuKernel(const std::vector<::pir::Value>& vec_inputs,
                        ::pir::Operation* op) {
@@ -326,7 +328,7 @@ static phi::Backend DeriveBackend(const std::string& op,
   // NOTE: Parameters are initialized on executor place defined
   if ((op == pir::SetParameterOp::name() ||
        op == pir::ShadowOutputOp::name()) &&
-      place.GetType() == phi::AllocationType::GPU) {
+      phi::is_accelerat_allocation_type(place.GetType())) {
     return phi::TransToPhiBackend(place);
   }
   // Tensor Attribute should on cpu backend for better performance
@@ -486,6 +488,9 @@ static pir::Value AddPlaceTransferOp(pir::Value in,
       case phi::AllocationType::CUSTOM:
         new_backend = phi::Backend::CUSTOM;
         break;
+      case phi::AllocationType::IPU:
+        new_backend = phi::Backend::IPU;
+        break;
       default:
         new_backend = phi::Backend::CPU;
         break;
@@ -494,18 +499,14 @@ static pir::Value AddPlaceTransferOp(pir::Value in,
   };
   std::unordered_map<std::string, pir::Attribute> op_attribute;
   if ((src_place.GetType() == phi::AllocationType::CPU) &&
-      (dst_place.GetType() == phi::AllocationType::GPU ||
-       dst_place.GetType() == phi::AllocationType::XPU ||
-       dst_place.GetType() == phi::AllocationType::CUSTOM)) {
+      phi::is_accelerat_allocation_type(dst_place.GetType())) {
     copy_kernel_key.set_backend(place2backend(dst_place.GetType()));
     op_attribute = {
         {"op_name", pir::StrAttribute::get(ctx, "pd_op.memcpy_h2d")},
         {"kernel_name", pir::StrAttribute::get(ctx, "memcpy_h2d")},
         {"kernel_key", KernelAttribute::get(ctx, copy_kernel_key)},
         {"dst_place_type", pir::Int32Attribute::get(ctx, 1)}};
-  } else if ((src_place.GetType() == phi::AllocationType::GPU ||
-              src_place.GetType() == phi::AllocationType::XPU ||
-              src_place.GetType() == phi::AllocationType::CUSTOM) &&
+  } else if (phi::is_accelerat_allocation_type(src_place.GetType()) &&
              (dst_place.GetType() == phi::AllocationType::CPU)) {
     copy_kernel_key.set_backend(place2backend(src_place.GetType()));
 
@@ -1351,10 +1352,11 @@ phi::KernelKey GetKernelKey(
   phi::KernelKey res(kernel_backend, kernel_layout, kernel_dtype);
 
   // kernel backend infered incorrectly from memcpy op operands,
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   // case that place from (not GPU) to GPU.
   // We handle this special case by following code to fix up the problem.
   // This could be further improved if we had another method.
-  if (!phi::is_gpu_place(place)) {
+  if (!phi::is_accelerat_place(place)) {
     if (op->isa<MemcpyOp>()) {
       VLOG(6) << "MemcpyOp need a special handle";
       int dst_place_type = op->attribute("dst_place_type")
@@ -1365,6 +1367,7 @@ phi::KernelKey GetKernelKey(
       }
     }
   }
+#endif
 
   if (op->isa<LoadCombineOp>()) {
     res.set_dtype(phi::DataType::FLOAT32);
@@ -1387,11 +1390,13 @@ phi::KernelKey GetKernelKey(
     VLOG(8) << "kernel backend must be on CPU when need fallback";
   }
 
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   if (NeedFallBackFromGPUDNN2GPU(op, kernel_fn_str, res)) {
     res.set_backend(phi::Backend::GPU);
     VLOG(8) << "kernel backend must be on GPU when need fallback from GPUDNN "
                "to GPU";
   }
+#endif
 
 #ifdef PADDLE_WITH_DNNL
   std::regex reg(",");
@@ -1651,13 +1656,7 @@ void AddShadowFeedForValue(
     pir::IrContext* ctx,
     std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
     std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
-  phi::Backend backend = phi::Backend::GPU;
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-  backend = phi::Backend::GPU;
-#endif
-#if defined(PADDLE_WITH_XPU)
-  backend = phi::Backend::XPU;
-#endif
+  phi::Backend backend = paddle::experimental::get_accelerat_backend();
 
   if (op_item->result(index).type().isa<DenseTensorType>()) {
     phi::KernelKey shadow_key{
@@ -1758,8 +1757,7 @@ void AddShadowFeedForTuplePopOp(
   }
 
   // if value place not gpu, add shadow feed op
-  if ((phi::is_gpu_place(place) || phi::is_xpu_place(place)) &&
-      add_shadow_feed) {
+  if (phi::is_accelerat_place(place) && add_shadow_feed) {
     for (size_t i = 0; i < op_item->num_results(); ++i) {
       AddShadowFeedForValue(i,
                             op_item,
@@ -2154,9 +2152,10 @@ void HandleForSpecialOp(
           auto out_place = phi::TransToPhiPlace(dst_backend);
           auto out_type =
               AllocatedDenseTensorType::get(ctx, out_place, value_type);
-          phi::KernelKey kernel_key(phi::Backend::GPU,
-                                    phi::DataLayout::ANY,
-                                    TransToPhiDataType(value_type.dtype()));
+          phi::KernelKey kernel_key(
+              paddle::experimental::get_accelerat_backend(),
+              phi::DataLayout::ANY,
+              TransToPhiDataType(value_type.dtype()));
           new_in = AddPlaceTransferOp(
               new_in, out_type, in_place, out_place, kernel_key, block);
         }
@@ -2998,11 +2997,9 @@ void AddShadowFeedOpForDataOrFeed(
     std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
     std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
   bool feed_op_add_shadow_feed =
-      (op_item->isa<FeedOp>()) &&
-      (phi::is_gpu_place(place) || phi::is_xpu_place(place));
+      (op_item->isa<FeedOp>()) && phi::is_accelerat_place(place);
   bool data_op_add_shadow_feed =
-      (op_item->isa<DataOp>()) &&
-      (phi::is_gpu_place(place) || phi::is_xpu_place(place)) &&
+      (op_item->isa<DataOp>()) && phi::is_accelerat_place(place) &&
       (kernel_op->attributes()
            .at("place")
            .dyn_cast<PlaceAttribute>()
@@ -3275,12 +3272,13 @@ void ProcessBlock(
           AllocatedDenseTensorType::get(ctx, phi::Place(), dense_tensor_type));
     }
   }
-  if (phi::is_gpu_place(place) || phi::is_xpu_place(place)) {
+  if (phi::is_accelerat_place(place)) {
     for (auto& [keyword, arg] : block->kwargs()) {
       if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
         auto dtype = dense_tensor_type.dtype();
-        phi::KernelKey shadow_key{
-            phi::Backend::GPU, phi::DataLayout::ANY, TransToPhiDataType(dtype)};
+        phi::KernelKey shadow_key{paddle::experimental::get_accelerat_backend(),
+                                  phi::DataLayout::ANY,
+                                  TransToPhiDataType(dtype)};
         std::unordered_map<std::string, pir::Attribute> attr_map{
             {"op_name", pir::StrAttribute::get(ctx, "pd_op.shadow_feed")},
             {"kernel_name", pir::StrAttribute::get(ctx, "shadow_feed")},
 
@@ -174,10 +174,29 @@ inline bool NeedTransformPlace(const phi::Place& src_place,
   if (!transform_flag.need_trans_backend()) {
     return false;
   }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   bool ret = src_place.GetType() == AllocationType::GPUPINNED ||
              (target != Backend::ALL_BACKEND &&
               phi::TransToPhiBackend(src_place) !=
                   (target != Backend::GPUDNN ? target : Backend::GPU));
+#elif defined(PADDLE_WITH_XPU)
+  bool ret = target != Backend::ALL_BACKEND &&
+             phi::TransToPhiBackend(src_place) != target;
+#elif defined(PADDLE_WITH_IPU)
+  bool ret = target != Backend::ALL_BACKEND &&
+             phi::TransToPhiBackend(src_place) != target;
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+  bool ret = target != Backend::ALL_BACKEND;
+  if (target == Backend::CUSTOM) {
+    ret = ret && !is_custom_place(src_place);
+  } else {
+    ret = ret && phi::TransToPhiBackend(src_place) != target;
+  }
+#else
+  bool ret = false;
+#endif
+
 #ifdef PADDLE_WITH_DNNL
   if (target == Backend::ONEDNN) {
     ret = src_place.GetType() != AllocationType::CPU;
 
@@ -193,6 +193,20 @@ inline std::string BackendToString(const Backend& backend) {
   }
 }
 
+inline Backend get_accelerat_backend() {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  return Backend::GPU;
+#elif defined(PADDLE_WITH_XPU)
+  return Backend::XPU;
+#elif defined(PADDLE_WITH_IPU)
+  return Backend::IPU;
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+  return Backend::CUSTOM;
+#else
+  return Backend::UNDEFINED;
+#endif
+}
+
 }  // namespace experimental
 }  // namespace paddle
 
 
@@ -179,6 +179,17 @@ bool is_custom_place(const Place &p) {
   return p.GetType() == phi::AllocationType::CUSTOM;
 }
 
+bool is_accelerat_place(const Place &p) {
+  return is_gpu_place(p) || is_xpu_place(p) || is_ipu_place(p) ||
+         is_custom_place(p);
+}
+
+bool is_accelerat_allocation_type(AllocationType type) {
+  return type == phi::AllocationType::GPU || type == phi::AllocationType::XPU ||
+         type == phi::AllocationType::IPU ||
+         type == phi::AllocationType::CUSTOM;
+}
+
 bool places_are_same_class(const Place &p1, const Place &p2) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   if (is_custom_place(p1) && is_custom_place(p2)) {
 
@@ -208,9 +208,10 @@ bool is_ipu_place(const Place&);
 TEST_API bool is_cpu_place(const Place&);
 bool is_cuda_pinned_place(const Place&);
 bool is_custom_place(const Place& p);
+bool is_accelerat_place(const Place& p);
 bool places_are_same_class(const Place&, const Place&);
 bool is_same_place(const Place&, const Place&);
-
+bool is_accelerat_allocation_type(AllocationType type);
 }  // namespace phi
 
 namespace paddle {
 
@@ -53,53 +53,14 @@ PD_REGISTER_KERNEL(data,
                    phi::complex64,
                    phi::complex128) {}
 
-PD_REGISTER_KERNEL(shadow_feed,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ShadowFeedKernel,
-                   bool,
-                   uint8_t,
-                   float,
-                   int8_t,
-                   int16_t,
-                   int32_t,
-                   int64_t,
-                   double,
-                   phi::float16,
-                   phi::bfloat16,
-                   phi::complex64,
-                   phi::complex128) {}
-
-PD_REGISTER_KERNEL(shadow_feed_tensors,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ShadowFeedTensorsKernel,
-                   bool,
-                   uint8_t,
-                   float,
-                   int8_t,
-                   int16_t,
-                   int32_t,
-                   int64_t,
-                   double,
-                   phi::float16,
-                   phi::bfloat16,
-                   phi::complex64,
-                   phi::complex128) {}
-
-PD_REGISTER_KERNEL(print_kernel,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::PrintKernel,
-                   bool,
-                   float,
-                   int32_t,
-                   int64_t,
-                   double,
-                   phi::float16,
-                   phi::bfloat16,
-                   phi::complex64,
-                   phi::complex128) {}
-
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(shadow_feed,
+                                         ALL_LAYOUT,
+                                         phi::ShadowFeedKernel) {}
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(shadow_feed_tensors,
+                                         ALL_LAYOUT,
+                                         phi::ShadowFeedTensorsKernel) {}
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(print_kernel,
+                                         ALL_LAYOUT,
+                                         phi::PrintKernel) {}
 PD_REGISTER_KERNEL(
     shadow_output, CPU, ALL_LAYOUT, phi::ShadowOutputKernel, float) {}
@@ -31,19 +31,19 @@ void ShadowOutputKernel(const Context& ctx,
                         const DenseTensor& x,
                         DenseTensor* out);
 
-template <typename T, typename Context>
+template <typename Context>
 void ShadowFeedKernel(const Context& ctx,
                       const DenseTensor& x,
                       int dst_place_type,
                       DenseTensor* out);
 
-template <typename T, typename Context>
+template <typename Context>
 void ShadowFeedTensorsKernel(const Context& ctx,
                              const std::vector<const DenseTensor*>& xs,
                              int dst_place_type,
                              std::vector<DenseTensor*> outs);
 
-template <typename T, typename Context>
+template <typename Context>
 void PrintKernel(const Context& ctx,
                  const DenseTensor& x,
                  int first_n,