add complex (#73414)

fangfangssj · web-flow · commit 61c79a12da40 · 2025-06-18T16:53:47.000+08:00
diff --git a/paddle/phi/core/visit_type.h b/paddle/phi/core/visit_type.h
@@ -336,6 +336,15 @@ namespace phi {
                  "`");                                                        \
     }                                                                         \
   }()
+
+#ifdef PADDLE_WITH_XPU_FFT
+#define PD_XPU_COMPLEX64_CASE(NAME, ...) \
+  PD_PRIVATE_CASE_TYPE(                  \
+      NAME, ::phi::DataType::COMPLEX64, phi::complex64, __VA_ARGS__)
+#else
+#define PD_XPU_COMPLEX64_CASE(NAME, ...)
+#endif
+
 #if defined(PADDLE_WITH_XPU)
 #define PD_VISIT_ALL_TYPES(TYPE, NAME, ...)                                    \
   [&] {                                                                        \
@@ -354,6 +363,7 @@ namespace phi {
       PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::FLOAT32, float, __VA_ARGS__) \
       PD_PRIVATE_CASE_TYPE(                                                    \
           NAME, ::phi::DataType::FLOAT64, double, __VA_ARGS__)                 \
+      PD_XPU_COMPLEX64_CASE(NAME, __VA_ARGS__)                                 \
       default:                                                                 \
         PADDLE_THROW(common::errors::InvalidArgument(                          \
             "Invalid enum data type `%d`.", static_cast<int>(__dtype__)));     \
diff --git a/paddle/phi/kernels/xpu/strided_copy_kernel.cc b/paddle/phi/kernels/xpu/strided_copy_kernel.cc
@@ -16,6 +16,7 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/complex_kernel.h"
 namespace phi {
 
 template <typename T, typename Context>
@@ -82,6 +83,30 @@ void StridedCopyKernel(const Context& dev_ctx,
   }
 }
 
+#ifdef PADDLE_WITH_XPU_FFT
+template <>
+void StridedCopyKernel<phi::dtype::complex<float>, XPUContext>(
+    const XPUContext& dev_ctx,
+    const DenseTensor& input,
+    const std::vector<int64_t>& dims,
+    const std::vector<int64_t>& out_stride,
+    int64_t offset,
+    DenseTensor* out) {
+  using T = phi::dtype::complex<float>;
+  dev_ctx.template Alloc<T>(out);
+  const DenseTensor real = Real<T, XPUContext>(dev_ctx, input);
+  const DenseTensor imag = Imag<T, XPUContext>(dev_ctx, input);
+  DenseTensor real_out, imag_out;
+  real_out.Resize(out->dims());
+  imag_out.Resize(out->dims());
+  StridedCopyKernel<float, XPUContext>(
+      dev_ctx, real, dims, out_stride, offset, &real_out);
+  StridedCopyKernel<float, XPUContext>(
+      dev_ctx, imag, dims, out_stride, offset, &imag_out);
+  phi::ComplexKernel<float>(dev_ctx, real_out, imag_out, out);
+}
+#endif
+
 }  // namespace phi
 
 PD_REGISTER_KERNEL(strided_copy,
@@ -96,5 +121,9 @@ PD_REGISTER_KERNEL(strided_copy,
                    int64_t,
                    float,
                    double,
+#ifdef PADDLE_WITH_XPU_FFT
+                   phi::dtype::complex<float>,
+#endif
                    ::phi::dtype::float16,
-                   ::phi::dtype::bfloat16) {}
+                   ::phi::dtype::bfloat16) {
+}
