[XPU] Support complex and 0-size tensor for add (#72752)

fangfangssj · web-flow · commit 6f25d857696a · 2025-05-18T17:06:11.000+08:00
* add complex

* fix test

* fix

* fix

* remove conj test
diff --git a/paddle/fluid/imperative/gradient_accumulator.cc b/paddle/fluid/imperative/gradient_accumulator.cc
@@ -187,6 +187,9 @@ void TensorAdd(const VarType& src, VarType* dst) {
     PADDLE_TENSOR_ADD(double, phi::XPUContext);
     PADDLE_TENSOR_ADD(phi::dtype::float16, phi::XPUContext);
     PADDLE_TENSOR_ADD(phi::dtype::bfloat16, phi::XPUContext);
+#ifdef PADDLE_WITH_XPU_FFT
+    PADDLE_TENSOR_ADD(phi::dtype::complex<float>, phi::XPUContext);
+#endif
 #endif
   }
 
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -407,13 +407,19 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
+#ifdef PADDLE_WITH_XPU_FFT
+                     phi::DataType::COMPLEX64,
+#endif
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"elementwise_add",
        XPUKernelSet({phi::DataType::FLOAT64,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
+#ifdef PADDLE_WITH_XPU_FFT
+                     phi::DataType::COMPLEX64,
+#endif
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"elementwise_div_grad",
diff --git a/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc b/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc
@@ -23,6 +23,8 @@
 #include "paddle/phi/backends/xpu/xpu_info.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/complex_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 
 namespace phi {
@@ -34,6 +36,25 @@ void AddGradKernel(const Context& dev_ctx,
                    int axis,
                    DenseTensor* dx,
                    DenseTensor* dy) {
+  if (dout.numel() == 0) {
+    if (dx) {
+      if (dx->numel() == 0) {
+        dev_ctx.template Alloc<T>(dx);
+      } else {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      if (dy->numel() == 0) {
+        dev_ctx.template Alloc<T>(dy);
+      } else {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   funcs::ElementwiseGradPreProcess(dout, dx);
   auto* dz = &dout;
@@ -97,6 +118,75 @@ void AddGradKernel(const Context& dev_ctx,
     }
   }
 }
+#ifdef PADDLE_WITH_XPU_FFT
+template <>
+void AddGradKernel<phi::dtype::complex<float>, XPUContext>(
+    const XPUContext& dev_ctx,
+    const DenseTensor& x,
+    const DenseTensor& y,
+    const DenseTensor& dout,
+    int axis,
+    DenseTensor* dx,
+    DenseTensor* dy) {
+  using T = phi::dtype::complex<float>;
+  const bool compute_dx = (dx != nullptr);
+  const bool compute_dy = (dy != nullptr);
+
+  // The current complex number implementation uses separate real/imaginary
+  // parts,resulting in redundant operations and performance
+  // penalties.Optimization should address this in future iterations.
+  DenseTensor x_real = Real<T, XPUContext>(dev_ctx, x);
+  DenseTensor x_imag = Imag<T, XPUContext>(dev_ctx, x);
+  DenseTensor y_real = Real<T, XPUContext>(dev_ctx, y);
+  DenseTensor y_imag = Imag<T, XPUContext>(dev_ctx, y);
+  DenseTensor dout_real = Real<T, XPUContext>(dev_ctx, dout);
+  DenseTensor dout_imag = Imag<T, XPUContext>(dev_ctx, dout);
+
+  if (compute_dx || compute_dy) {
+    DenseTensor dx_real, dx_imag, dy_real, dy_imag;
+
+    if (compute_dx) {
+      dx_real.Resize(x_real.dims());
+      dx_imag.Resize(x_imag.dims());
+      dev_ctx.template Alloc<float>(&dx_real);
+      dev_ctx.template Alloc<float>(&dx_imag);
+    }
+
+    if (compute_dy) {
+      dy_real.Resize(y_real.dims());
+      dy_imag.Resize(y_imag.dims());
+      dev_ctx.template Alloc<float>(&dy_real);
+      dev_ctx.template Alloc<float>(&dy_imag);
+    }
+
+    AddGradKernel<float, XPUContext>(dev_ctx,
+                                     x_real,
+                                     y_real,
+                                     dout_real,
+                                     axis,
+                                     compute_dx ? &dx_real : nullptr,
+                                     compute_dy ? &dy_real : nullptr);
+
+    AddGradKernel<float, XPUContext>(dev_ctx,
+                                     x_imag,
+                                     y_imag,
+                                     dout_imag,
+                                     axis,
+                                     compute_dx ? &dx_imag : nullptr,
+                                     compute_dy ? &dy_imag : nullptr);
+
+    if (compute_dx) {
+      dev_ctx.template Alloc<T>(dx);
+      phi::ComplexKernel<float>(dev_ctx, dx_real, dx_imag, dx);
+    }
+
+    if (compute_dy) {
+      dev_ctx.template Alloc<T>(dy);
+      phi::ComplexKernel<float>(dev_ctx, dy_real, dy_imag, dy);
+    }
+  }
+}
+#endif
 }  // namespace phi
 
 PD_REGISTER_KERNEL(add_grad,
@@ -105,6 +195,10 @@ PD_REGISTER_KERNEL(add_grad,
                    phi::AddGradKernel,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
+#ifdef PADDLE_WITH_XPU_FFT
+                   phi::dtype::complex<float>,
+#endif
                    float,
                    int,
-                   int64_t) {}
+                   int64_t) {
+}
diff --git a/paddle/phi/kernels/xpu/elementwise_add_kernel.cc b/paddle/phi/kernels/xpu/elementwise_add_kernel.cc
@@ -24,6 +24,7 @@
 #include "paddle/phi/backends/xpu/xpu_info.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/complex_kernel.h"
 #include "paddle/phi/kernels/impl/elementwise_kernel_impl.h"
 #include "paddle/phi/kernels/xpu/elementwise.h"
 
@@ -34,6 +35,10 @@ void AddKernel(const Context& dev_ctx,
                const DenseTensor& x,
                const DenseTensor& y,
                DenseTensor* out) {
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   if (x.dtype() == phi::DataType::FLOAT32 &&
       (y.dtype() == phi::DataType::BFLOAT16 ||
        y.dtype() == phi::DataType::FLOAT16)) {
@@ -111,6 +116,46 @@ void GradAddXPUKernel(const Context& dev_ctx,
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_add");
 }
 
+#ifdef PADDLE_WITH_XPU_FFT
+template <>
+void AddKernel<phi::dtype::complex<float>, XPUContext>(
+    const XPUContext& dev_ctx,
+    const DenseTensor& x,
+    const DenseTensor& y,
+    DenseTensor* out) {
+  using T = phi::dtype::complex<float>;
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+  auto f = [](xpu::Context* ctx,
+              const float* x,
+              const float* y,
+              float* z,
+              const std::vector<int64_t>& xshape,
+              const std::vector<int64_t>& yshape) {
+    return xpu::broadcast_add<float>(ctx, x, y, z, xshape, yshape);
+  };
+  // The current complex number implementation uses separate real/imaginary
+  // parts,resulting in redundant operations and performance
+  // penalties.Optimization should address this in future iterations.
+  const DenseTensor x_real = Real<T, XPUContext>(dev_ctx, x);
+  const DenseTensor x_imag = Imag<T, XPUContext>(dev_ctx, x);
+  const DenseTensor y_real = Real<T, XPUContext>(dev_ctx, y);
+  const DenseTensor y_imag = Imag<T, XPUContext>(dev_ctx, y);
+  DenseTensor real_out, imag_out;
+  real_out.Resize(out->dims());
+  imag_out.Resize(out->dims());
+  dev_ctx.template Alloc<float>(&real_out);
+  dev_ctx.template Alloc<float>(&imag_out);
+
+  XPUElementwise<float, float>(dev_ctx, x_real, y_real, -1, &real_out, f);
+  XPUElementwise<float, float>(dev_ctx, x_imag, y_imag, -1, &imag_out, f);
+  phi::ComplexKernel<float>(dev_ctx, real_out, imag_out, out);
+}
+
+#endif
+
 }  // namespace phi
 
 PD_REGISTER_KERNEL(grad_add,
@@ -127,6 +172,10 @@ PD_REGISTER_KERNEL(add,
                    double,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
+#ifdef PADDLE_WITH_XPU_FFT
+                   phi::dtype::complex<float>,
+#endif
                    float,
                    int,
-                   int64_t) {}
+                   int64_t) {
+}
diff --git a/test/xpu/test_elementwise_add_op_xpu.py b/test/xpu/test_elementwise_add_op_xpu.py
@@ -121,6 +121,30 @@ def init_input_output(self):
             self.y = np.random.uniform(-1, 1, []).astype(self.dtype)
             self.out = self.x + self.y
 
+    class TestElementwiseAddOp_ZeroSize1(TestElementwiseAddOp):
+        def init_input_output(self):
+            self.x = np.random.rand(0, 3, 4).astype(self.dtype)
+            self.y = np.random.rand(1).astype(self.dtype)
+            self.out = self.x + self.y
+
+    class TestElementwiseAddOp_ZeroSize2(TestElementwiseAddOp):
+        def init_input_output(self):
+            self.x = np.random.rand(0, 3, 4).astype(self.dtype)
+            self.y = np.random.rand(0, 3, 4).astype(self.dtype)
+            self.out = self.x + self.y
+
+    class TestElementwiseAddOp_ZeroSize3(TestElementwiseAddOp):
+        def init_input_output(self):
+            self.x = np.random.rand(3, 0, 4).astype(self.dtype)
+            self.y = np.random.rand(1, 1, 4).astype(self.dtype)
+            self.out = self.x + self.y
+
+    class TestElementwiseAddOp_ZeroSize4(TestElementwiseAddOp):
+        def init_input_output(self):
+            self.x = np.random.rand(3, 0, 4).astype(self.dtype)
+            self.y = np.random.rand(1, 0, 4).astype(self.dtype)
+            self.out = self.x + self.y
+
     @skip_check_grad_ci(
         reason="[skip shape check] Use y_shape(1) to test broadcast."
     )
@@ -343,6 +367,91 @@ def test_dygraph(self):
     create_test_class(globals(), XPUTestElementwiseAddOp, stype)
 
 
+class TestElementwiseAddOpComplex(XPUOpTest):
+    def setUp(self):
+        self.op_type = "elementwise_add"
+        self.init_dtype()
+        self.init_input_output()
+        self.init_axis()
+        self.inputs = {
+            'X': OpTest.np_dtype_to_base_dtype(self.x),
+            'Y': OpTest.np_dtype_to_base_dtype(self.y),
+        }
+        self.attrs = {'axis': self.axis, 'use_mkldnn': self.use_mkldnn}
+        self.outputs = {'Out': self.out}
+
+    def test_check_output(self):
+        if paddle.is_compiled_with_xpu():
+            place = paddle.XPUPlace(0)
+            self.check_output_with_place(place)
+
+    def test_check_grad_normal(self):
+        if paddle.is_compiled_with_xpu():
+            place = paddle.XPUPlace(0)
+            self.check_grad_with_place(
+                place,
+                ['X', 'Y'],
+                'Out',
+            )
+
+    def test_check_grad_ignore_x(self):
+        if paddle.is_compiled_with_xpu():
+            place = paddle.XPUPlace(0)
+            self.check_grad_with_place(
+                place,
+                ['Y'],
+                'Out',
+                no_grad_set=set("X"),
+            )
+
+    def test_check_grad_ignore_y(self):
+        if paddle.is_compiled_with_xpu():
+            place = paddle.XPUPlace(0)
+            self.check_grad_with_place(
+                place,
+                ['X'],
+                'Out',
+                no_grad_set=set("Y"),
+            )
+
+    def init_input_output(self):
+        self.x = (
+            np.random.rand(2, 3, 4) + 1j * np.random.rand(2, 3, 4)
+        ).astype(self.dtype)
+        self.y = (
+            np.random.rand(2, 3, 4) + 1j * np.random.rand(2, 3, 4)
+        ).astype(self.dtype)
+        self.out = self.x + self.y
+
+    def init_dtype(self):
+        self.dtype = np.complex64
+
+    def init_axis(self):
+        self.axis = -1
+
+
+class TestElementwiseAddOpComplex2(TestElementwiseAddOpComplex):
+    def init_input_output(self):
+        self.x = (
+            np.random.rand(2, 3, 4) + 1j * np.random.rand(2, 3, 4)
+        ).astype(self.dtype)
+        self.y = (np.random.rand(1, 1) + 1j * np.random.rand(1, 1)).astype(
+            self.dtype
+        )
+        self.out = self.x + self.y
+
+
+class TestElementwiseAddOpComplex3(TestElementwiseAddOpComplex):
+    def init_input_output(self):
+        self.x = (
+            np.random.rand(100, 2, 3) + 1j * np.random.rand(100, 2, 3)
+        ).astype(self.dtype)
+        self.y = (np.random.rand(1, 1) + 1j * np.random.rand(1, 1)).astype(
+            self.dtype
+        )
+        self.out = self.x + self.y
+
+
 @unittest.skipIf(
     core.get_xpu_device_version(0) != core.XPUVersion.XPU3,
     "only supported on XPU3",

Original file line number	Diff line number	Diff line change
`@@ -187,6 +187,9 @@ void TensorAdd(const VarType& src, VarType* dst) {`
`187`	`187`	`PADDLE_TENSOR_ADD(double, phi::XPUContext);`
`188`	`188`	`PADDLE_TENSOR_ADD(phi::dtype::float16, phi::XPUContext);`
`189`	`189`	`PADDLE_TENSOR_ADD(phi::dtype::bfloat16, phi::XPUContext);`
	`190`	`+#ifdef PADDLE_WITH_XPU_FFT`
	`191`	`+ PADDLE_TENSOR_ADD(phi::dtype::complex<float>, phi::XPUContext);`
	`192`	`+#endif`
`190`	`193`	`#endif`
`191`	`194`	`}`
`192`	`195`