【Paddle Tensor 规范化第二期】paddle.cumsum supports complex (#72391)

fangfangssj · web-flow · commit 0221eeeaaef1 · 2025-04-23T17:52:32.000+08:00
* add complex

* add kernel

* add test

* fix ci
diff --git a/paddle/phi/common/complex.h b/paddle/phi/common/complex.h
@@ -64,7 +64,7 @@ struct PADDLE_ALIGN(sizeof(T) * 2) complex {
   complex& operator=(complex<T>&& o) = default;
   ~complex() = default;
 
-  HOSTDEVICE complex(T real, T imag) : real(real), imag(imag) {}
+  HOSTDEVICE constexpr complex(T real, T imag) : real(real), imag(imag) {}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 
diff --git a/paddle/phi/kernels/cpu/cum_grad_kernel.cc b/paddle/phi/kernels/cpu/cum_grad_kernel.cc
@@ -56,4 +56,6 @@ PD_REGISTER_KERNEL(cumsum_grad,
                    double,
                    int16_t,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/cum_kernel.cc b/paddle/phi/kernels/cpu/cum_kernel.cc
@@ -271,7 +271,9 @@ PD_REGISTER_KERNEL(cumsum,
                    double,
                    int16_t,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(
     logcumsumexp, CPU, ALL_LAYOUT, phi::LogcumsumexpKernel, float, double) {}
diff --git a/paddle/phi/kernels/gpu/cum_grad_kernel.cu b/paddle/phi/kernels/gpu/cum_grad_kernel.cu
@@ -85,5 +85,7 @@ PD_REGISTER_KERNEL(cumsum_grad,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 #endif
diff --git a/paddle/phi/kernels/gpu/cum_kernel.cu b/paddle/phi/kernels/gpu/cum_kernel.cu
@@ -143,6 +143,14 @@ struct LogAddExp {
   }
 };
 
+struct ComplexSum {
+  template <typename T>
+  __host__ __device__ __forceinline__ T operator()(const T& a,
+                                                   const T& b) const {
+    return a + b;
+  }
+};
+
 template <typename T, typename op>
 struct Identity;
 
@@ -156,6 +164,11 @@ struct Identity<T, LogAddExp> {
   static constexpr T value = std::numeric_limits<T>::lowest();
 };
 
+template <typename T>
+struct Identity<T, ComplexSum> {
+  static constexpr T value = {0, 0};
+};
+
 template <typename T, int BLOCK_THREADS, int ITEMS_PER_THREAD, typename Op>
 __global__ void BlockScanKernel(T* d_out,
                                 const T* d_in,
@@ -410,7 +423,11 @@ void CumsumKernel(const Context& dev_ctx,
                   bool exclusive,
                   bool reverse,
                   DenseTensor* out) {
-  using Op = cub::Sum;
+  using Op = typename std::conditional<
+      std::is_same<T, phi::dtype::complex<float>>::value ||
+          std::is_same<T, phi::dtype::complex<double>>::value,
+      ComplexSum,
+      cub::Sum>::type;
   auto op = Op();
   ScanKernel<T, Context, Op>(
       dev_ctx, x, axis.to<int>(), flatten, exclusive, reverse, op, out);
@@ -457,7 +474,9 @@ PD_REGISTER_KERNEL(cumsum,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(logcumsumexp,
                    GPU,
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4280,7 +4280,7 @@ def cumsum(
     Args:
         x (Tensor): The input tensor needed to be cumsumed.
         axis (int, optional): The dimension to accumulate along. -1 means the last dimension. The default (None) is to compute the cumsum over the flattened array.
-        dtype (str|paddle.dtype|np.dtype|None, optional): The data type of the output tensor, can be bfloat16, float16, float32, float64, int32, int64. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. The default value is None.
+        dtype (str|paddle.dtype|np.dtype|None, optional): The data type of the output tensor, can be bfloat16, float16, float32, float64, int32, int64, complex64, complex128. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. The default value is None.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4331,7 +4331,16 @@ def cumsum(
         check_variable_and_dtype(
             x,
             'x',
-            ['float16', 'uint16', 'float32', 'float64', 'int32', 'int64'],
+            [
+                'float16',
+                'uint16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'complex64',
+                'complex128',
+            ],
             'cumsum',
         )
         check_type(x, 'x', (Variable), 'cumsum')
diff --git a/test/legacy_test/test_cumsum_op.py b/test/legacy_test/test_cumsum_op.py
@@ -212,6 +212,95 @@ def set_attrs_input_output(self):
         self.out = self.x.cumsum(axis=0)
 
 
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp1(TestSumOp1):
+    def set_attrs_input_output(self):
+        self.attrs = {'axis': 2}
+        x_real = np.random.random((5, 6, 10)).astype(self.dtype_)
+        x_imag = np.random.random((5, 6, 10)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum(axis=2)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp2(TestSumOp1):
+    def set_attrs_input_output(self):
+        self.attrs = {'axis': -1, 'reverse': True}
+        x_real = np.random.random((5, 6, 10)).astype(self.dtype_)
+        x_imag = np.random.random((5, 6, 10)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = np.flip(np.flip(self.x, axis=2).cumsum(axis=2), axis=2)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp3(TestSumOp1):
+    def set_attrs_input_output(self):
+        self.attrs = {'axis': 1}
+        x_real = np.random.random((5, 6, 10)).astype(self.dtype_)
+        x_imag = np.random.random((5, 6, 10)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum(axis=1)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp4(TestSumOp1):
+    def set_attrs_input_output(self):
+        self.attrs = {'axis': 0}
+        x_real = np.random.random((5, 6, 10)).astype(self.dtype_)
+        x_imag = np.random.random((5, 6, 10)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum(axis=0)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp5(TestSumOp1):
+    def set_attrs_input_output(self):
+        x_real = np.random.random((5, 20)).astype(self.dtype_)
+        x_imag = np.random.random((5, 20)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum(axis=1)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp6(TestSumOp1):
+    def set_attrs_input_output(self):
+        self.attrs = {'axis': -1, 'flatten': True}
+        x_real = np.random.random((5, 6, 5)).astype(self.dtype_)
+        x_imag = np.random.random((5, 6, 5)).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum()
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestSumComplexOp7(TestSumOp1):
+    def set_attrs_input_output(self):
+        x_real = np.random.random(100).astype(self.dtype_)
+        x_imag = np.random.random(100).astype(self.dtype_)
+        self.x = x_real + 1j * x_imag
+        self.out = self.x.cumsum(axis=0)
+
+
 class TestCumsumFP16(unittest.TestCase):
     def check_main(self, x_np, dtype):
         paddle.disable_static()
