[XPU] Fix visit_type bug in fallback process (#72782)

fangfangssj · web-flow · commit 85948606e3f8 · 2025-05-20T11:17:20.000+08:00
* fix

* fix
diff --git a/paddle/phi/core/visit_type.h b/paddle/phi/core/visit_type.h
@@ -394,6 +394,39 @@ namespace phi {
   }()
 #endif
 
+#define PD_VISIT_ALL_CPU_TYPES(TYPE, NAME, ...)                                \
+  [&] {                                                                        \
+    const auto& __dtype__ = TYPE;                                              \
+    switch (__dtype__) {                                                       \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::BOOL, bool, __VA_ARGS__)     \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT8, int8_t, __VA_ARGS__)   \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::UINT8, uint8_t, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT16, int16_t, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT32, int32_t, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT64, int64_t, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::BFLOAT16, phi::bfloat16, __VA_ARGS__)         \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::FLOAT16, phi::float16, __VA_ARGS__)           \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::FLOAT32, float, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::FLOAT64, double, __VA_ARGS__)                 \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::COMPLEX64, phi::complex64, __VA_ARGS__)       \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::COMPLEX128, phi::complex128, __VA_ARGS__)     \
+      PD_PRIVATE_CASE_TYPE(NAME,                                               \
+                           ::phi::DataType::FLOAT8_E4M3FN,                     \
+                           phi::float8_e4m3fn,                                 \
+                           __VA_ARGS__)                                        \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::FLOAT8_E5M2, phi::float8_e5m2, __VA_ARGS__)   \
+      default:                                                                 \
+        PADDLE_THROW(common::errors::InvalidArgument(                          \
+            "Invalid enum data type `%d`.", static_cast<int>(__dtype__)));     \
+    }                                                                          \
+  }()
+
 #define PD_VISIT_BOOL_AND_FLOATING_AND_COMPLEX_AND_4_TYPES(SPECIFIED_TYPE1,   \
                                                            SPECIFIED_TYPE2,   \
                                                            SPECIFIED_TYPE3,   \
diff --git a/paddle/phi/kernels/cpu/cast_kernel.cc b/paddle/phi/kernels/cpu/cast_kernel.cc
@@ -37,14 +37,15 @@ void CastKernel(const Context& dev_ctx,
   }
 
   if (out->IsSharedWith(x)) {
-    PD_VISIT_ALL_TYPES(out_dtype, "CastInplaceKernelImpl", ([&] {
-                         CastInplaceKernelImpl<T, data_t>(
-                             dev_ctx, x, out_dtype, out);
-                       }));
+    PD_VISIT_ALL_CPU_TYPES(out_dtype, "CastInplaceKernelImpl", ([&] {
+                             CastInplaceKernelImpl<T, data_t>(
+                                 dev_ctx, x, out_dtype, out);
+                           }));
   } else {
-    PD_VISIT_ALL_TYPES(out_dtype, "CastKernelImpl", ([&] {
-                         CastKernelImpl<T, data_t>(dev_ctx, x, out_dtype, out);
-                       }));
+    PD_VISIT_ALL_CPU_TYPES(out_dtype, "CastKernelImpl", ([&] {
+                             CastKernelImpl<T, data_t>(
+                                 dev_ctx, x, out_dtype, out);
+                           }));
   }
 }
 
diff --git a/paddle/phi/kernels/cpu/reduce.h b/paddle/phi/kernels/cpu/reduce.h
@@ -47,7 +47,7 @@ void Reduce(const DeviceContext& dev_ctx,
   // no need to cast dtype
   if (out_dtype == phi::DataType::UNDEFINED || out_dtype == x.dtype()) {
     // do reduce sum
-    PD_VISIT_ALL_TYPES(
+    PD_VISIT_ALL_CPU_TYPES(
         x.dtype(), "ReduceKernelImpl", ([&] {
           phi::funcs::ReduceKernelImpl<DeviceContext, T, data_t, Functor>(
               dev_ctx, x, out, dims, keep_dim, reduce_all);
@@ -58,7 +58,7 @@ void Reduce(const DeviceContext& dev_ctx,
     auto tmp_tensor = phi::Cast<T, DeviceContext>(dev_ctx, x, out_dtype);
 
     // do reduce sum
-    PD_VISIT_ALL_TYPES(
+    PD_VISIT_ALL_CPU_TYPES(
         out_dtype, "ReduceKernelImpl", ([&] {
           phi::funcs::ReduceKernelImpl<DeviceContext, T, data_t, Functor>(
               dev_ctx, tmp_tensor, out, dims, keep_dim, reduce_all);
