[PIR] Update output place for memcpy op (#70201)

SigureMo · web-flow · commit 33cf3fd30ac6 · 2024-12-13T18:59:53.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
@@ -176,6 +176,19 @@ const std::unordered_map<std::string, uint32_t> NoBufferRelatedOps = {
     {paddle::dialect::BatchNorm_Op::name(), /*reserve_space*/ 5U},
 };
 
+// Please keep the consistency with paddle/phi/kernels/memcpy_kernel.cc
+const std::unordered_map<int, phi::Place> MemcpyOpAttr2Place = {
+    {0, phi::CPUPlace()},
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    {1, phi::GPUPlace()},
+    {2, phi::GPUPinnedPlace()},
+#elif defined(PADDLE_WITH_XPU)
+    {3, phi::XPUPlace()},
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+    {4, phi::CustomPlace()}
+#endif
+};
+
 static bool NeedSkipPlaceTransfer(const pir::Operation* op) {
   bool need_skip = false;
   if (op->isa<paddle::dialect::FetchOp>()) {
@@ -2179,8 +2192,6 @@ void HandleForSpecialOp(
   // only deal with single output
   if (op_item->num_results() > 0) {
     for (size_t i = 0; i < op_item->num_results(); ++i) {
-      VLOG(6) << "2816:" << op_item->result(i).type();
-      VLOG(6) << "2817:" << op->result(i).type();
       (*map_value_pair)[op_item->result(i)] = op->result(i);
     }
   }
@@ -2508,6 +2519,13 @@ std::vector<pir::Type> BuildOutputs(
           (!IsLegacyOp(op_item->name())) && phi_kernel.IsValid()) {
         out_place = phi::TransToPhiPlace(output_defs[i].backend);
       }
+      if (op_item->isa<MemcpyOp>()) {
+        // If the op is MemcpyOp, the output type is determined by the
+        // attribute "dst_place_type".
+        out_place = MemcpyOpAttr2Place.at(op_item->attribute("dst_place_type")
+                                              .dyn_cast<pir::Int32Attribute>()
+                                              .data());
+      }
       PushBackOutputTypes(ctx,
                           op_item,
                           op_item->result(i).type(),
diff --git a/paddle/phi/kernels/memcpy_kernel.cc b/paddle/phi/kernels/memcpy_kernel.cc
@@ -106,12 +106,12 @@ void MemcpyKernel(const Context& dev_ctx,
   PADDLE_ENFORCE_GE(
       dst_place_type,
       0,
-      errors::OutOfRange("dst_place_type only support 0-2, but got: %d",
+      errors::OutOfRange("dst_place_type only support 0-4, but got: %d",
                          dst_place_type));
   PADDLE_ENFORCE_LE(
       dst_place_type,
-      2,
-      errors::OutOfRange("dst_place_type only support 0-2, but got: %d",
+      4,
+      errors::OutOfRange("dst_place_type only support 0-4, but got: %d",
                          dst_place_type));
   switch (dst_place_type) {
     case 0: /* CPUPlace */
diff --git a/test/dygraph_to_static/test_tensor_memcpy_on_gpu.py b/test/dygraph_to_static/test_tensor_memcpy_on_gpu.py
@@ -45,6 +45,12 @@ def tensor_copy_to_cuda_with_warning(x, device_id=None, blocking=True):
     return y
 
 
+def tensor_copy_to_cpu_with_compute(x):
+    x = paddle.to_tensor(x)
+    y = x.cpu()
+    return y + 1
+
+
 class TestTensorCopyToCpuOnDefaultGPU(Dy2StTestBase):
     def _run(self):
         x1 = paddle.ones([1, 2, 3])
@@ -124,5 +130,27 @@ def test_with_warning_on_gpu(self):
         self.assertIn('math_op_patch.py', cm.filename)
 
 
+class TestTensorCopyToCPUWithComputeOnDefaultGPU(Dy2StTestBase):
+    def _run(self):
+        x1 = paddle.ones([1, 2, 3])
+        x2 = paddle.jit.to_static(tensor_copy_to_cpu_with_compute)(x1)
+        return x1.place, x2.place, x2.numpy()
+
+    def test_tensor_cpu_with_compute_on_default_gpu(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+        place = paddle.CUDAPlace(int(os.environ.get('FLAGS_selected_gpus', 0)))
+        paddle.framework._set_expected_place(place)
+        with enable_to_static_guard(False):
+            dygraph_x1_place, dygraph_place, dygraph_res = self._run()
+
+        static_x1_place, static_place, static_res = self._run()
+        np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
+        self.assertTrue(dygraph_x1_place.is_gpu_place())
+        self.assertTrue(static_x1_place.is_gpu_place())
+        self.assertTrue(dygraph_place.is_cpu_place())
+        self.assertTrue(static_place.is_cpu_place())
+
+
 if __name__ == '__main__':
     unittest.main()
