refine

Author: wanghuancoder

paddle/common/flags.cc

@@ -1820,3 +1820,8 @@ PHI_DEFINE_EXPORTED_double(accuracy_check_atol_bf16,
 PHI_DEFINE_EXPORTED_double(accuracy_check_rtol_bf16,
                            1e-3,
                            "It controls the rtol of accuracy_check op");
+
+PHI_DEFINE_EXPORTED_bool(
+    pinned_memory_as_cpu_backend,
+    false,
+    "Whether use CPU backend, when tensor is pinned_memory.");

paddle/phi/core/compat/convert_utils.cc

@@ -14,16 +14,18 @@ limitations under the License. */
 
 #include "paddle/phi/core/compat/convert_utils.h"
 
+#include "paddle/common/flags.h"
 #include "paddle/phi/backends/gpu/gpu_info.h"
 #include "paddle/phi/backends/xpu/xpu_info.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/compat/op_utils.h"
 #include "paddle/phi/core/enforce.h"
-
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/phi/backends/device_manager.h"
 #endif
 
+COMMON_DECLARE_bool(pinned_memory_as_cpu_backend);
+
 namespace phi {
 
 Backend TransToPhiBackend(const phi::Place& place) {
@@ -33,8 +35,13 @@ Backend TransToPhiBackend(const phi::Place& place) {
       return Backend::GPU;
     case AllocationType::CPU:
       return Backend::CPU;
-    case AllocationType::GPUPINNED:
-      return Backend::CPU;
+    case AllocationType::GPUPINNED: {
+      if (FLAGS_pinned_memory_as_cpu_backend) {
+        return Backend::CPU;
+      } else {
+        return Backend::GPU;
+      }
+    }
     case AllocationType::XPU:
       return Backend::XPU;
     case AllocationType::IPU: