为paddle.less_than、paddle.less_equal、paddle.greater_than、paddle.greater_equal添加复数类型支持 (#72619)

aquagull · web-flow · commit 876366b2c851 · 2025-05-09T14:22:20.000+08:00
* update code

* update api

* fix less_than
diff --git a/paddle/phi/kernels/cpu/compare_kernel.cc b/paddle/phi/kernels/cpu/compare_kernel.cc
@@ -124,28 +124,10 @@ PD_REGISTER_KERNEL(equal_all,
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
   }
 
-#define PD_REGISTER_COMPARE_KERNEL(name, func)            \
-  PD_REGISTER_KERNEL(name,                                \
-                     CPU,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##Kernel,                   \
-                     bool,                                \
-                     int,                                 \
-                     uint8_t,                             \
-                     int8_t,                              \
-                     int16_t,                             \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }
-
-PD_REGISTER_COMPARE_KERNEL(less_than, LessThan)
-PD_REGISTER_COMPARE_KERNEL(less_equal, LessEqual)
-PD_REGISTER_COMPARE_KERNEL(greater_than, GreaterThan)
-PD_REGISTER_COMPARE_KERNEL(greater_equal, GreaterEqual)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(less_than, LessThan)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(greater_equal, GreaterEqual)
 
 PD_REGISTER_COMPLEX_COMPARE_KERNEL(equal, Equal)
 PD_REGISTER_COMPLEX_COMPARE_KERNEL(not_equal, NotEqual)
diff --git a/paddle/phi/kernels/funcs/compare_functors.h b/paddle/phi/kernels/funcs/compare_functors.h
@@ -26,11 +26,49 @@ namespace funcs {
     }                                                            \
   };
 
+#define COMPARE_COMPLEX_FUNCTOR(func_name, op)                              \
+  template <typename T>                                                     \
+  struct func_name<phi::dtype::complex<T>> {                                \
+    HOSTDEVICE bool operator()(const phi::dtype::complex<T> a,              \
+                               const phi::dtype::complex<T> b) const {      \
+      if (isnan(a.real) || isnan(a.imag) || isnan(b.real) || isnan(b.imag)) \
+        return false;                                                       \
+      T ar = a.real;                                                        \
+      T br = b.real;                                                        \
+      T ai = a.imag;                                                        \
+      T bi = b.imag;                                                        \
+      return (ar op br) || (ar == br && ai op bi);                          \
+    }                                                                       \
+  };
+
+#define COMPARE_COMPLEX_EQUAL_FUNCTOR(func_name, op_equal, op)              \
+  template <typename T>                                                     \
+  struct func_name<phi::dtype::complex<T>> {                                \
+    HOSTDEVICE bool operator()(const phi::dtype::complex<T> a,              \
+                               const phi::dtype::complex<T> b) const {      \
+      if (isnan(a.real) || isnan(a.imag) || isnan(b.real) || isnan(b.imag)) \
+        return false;                                                       \
+      T ar = a.real;                                                        \
+      T br = b.real;                                                        \
+      T ai = a.imag;                                                        \
+      T bi = b.imag;                                                        \
+      return (ar op br) || (ar == br && ai op_equal bi);                    \
+    }                                                                       \
+  };
+
 COMPARE_FUNCTOR(LessThanFunctor, <)
 COMPARE_FUNCTOR(LessEqualFunctor, <=)
 COMPARE_FUNCTOR(GreaterThanFunctor, >)
 COMPARE_FUNCTOR(GreaterEqualFunctor, >=)
+
+COMPARE_COMPLEX_FUNCTOR(LessThanFunctor, <)
+COMPARE_COMPLEX_FUNCTOR(GreaterThanFunctor, >)
+COMPARE_COMPLEX_EQUAL_FUNCTOR(LessEqualFunctor, <=, <)
+COMPARE_COMPLEX_EQUAL_FUNCTOR(GreaterEqualFunctor, >=, >)
+
 #undef COMPARE_FUNCTOR
+#undef COMPARE_COMPLEX_FUNCTOR
+#undef COMPARE_COMPLEX_EQUAL_FUNCTOR
 
 template <typename InT, typename OutT = bool>
 struct EqualFunctor {
diff --git a/paddle/phi/kernels/kps/compare_kernel.cu b/paddle/phi/kernels/kps/compare_kernel.cu
@@ -145,24 +145,6 @@ PD_REGISTER_KERNEL(equal_all,
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
-#define PD_REGISTER_COMPARE_KERNEL(name, func)            \
-  PD_REGISTER_KERNEL(name,                                \
-                     KPS,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##Kernel,                   \
-                     bool,                                \
-                     int,                                 \
-                     uint8_t,                             \
-                     int8_t,                              \
-                     int16_t,                             \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }
-
 #define PD_REGISTER_COMPLEX_COMPARE_KERNEL(name, func)    \
   PD_REGISTER_KERNEL(name,                                \
                      KPS,                                 \
@@ -183,11 +165,10 @@ PD_REGISTER_KERNEL(equal_all,
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
   }
 
-PD_REGISTER_COMPARE_KERNEL(less_than, LessThan)
-PD_REGISTER_COMPARE_KERNEL(less_equal, LessEqual)
-PD_REGISTER_COMPARE_KERNEL(greater_than, GreaterThan)
-PD_REGISTER_COMPARE_KERNEL(greater_equal, GreaterEqual)
-
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(less_than, LessThan)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPLEX_COMPARE_KERNEL(greater_equal, GreaterEqual)
 PD_REGISTER_COMPLEX_COMPARE_KERNEL(equal, Equal)
 PD_REGISTER_COMPLEX_COMPARE_KERNEL(not_equal, NotEqual)
 
diff --git a/paddle/phi/kernels/legacy/cpu/compare_kernel.cc b/paddle/phi/kernels/legacy/cpu/compare_kernel.cc
@@ -120,31 +120,15 @@ PD_REGISTER_KERNEL(less_than_raw,
                    int16_t,
                    int,
                    int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>,
                    float,
                    double,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
-#define PD_REGISTER_COMPARE_RAW_KERNEL(name, func)        \
-  PD_REGISTER_KERNEL(name##_raw,                          \
-                     CPU,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##RawKernel,                \
-                     bool,                                \
-                     uint8_t,                             \
-                     int8_t,                              \
-                     int16_t,                             \
-                     int,                                 \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }
-
 #define PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name##_raw,                           \
                      CPU,                                  \
@@ -165,9 +149,9 @@ PD_REGISTER_KERNEL(less_than_raw,
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);  \
   }
 
-PD_REGISTER_COMPARE_RAW_KERNEL(less_equal, LessEqual)
-PD_REGISTER_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
-PD_REGISTER_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
 
 PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(equal, Equal)
 PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(not_equal, NotEqual)
diff --git a/paddle/phi/kernels/legacy/kps/compare_kernel.cu b/paddle/phi/kernels/legacy/kps/compare_kernel.cu
@@ -146,31 +146,15 @@ PD_REGISTER_KERNEL(less_than_raw,
                    int16_t,
                    int,
                    int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>,
                    float,
                    double,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
-#define PD_REGISTER_COMPARE_RAW_KERNEL(name, func)        \
-  PD_REGISTER_KERNEL(name##_raw,                          \
-                     KPS,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##RawKernel,                \
-                     bool,                                \
-                     uint8_t,                             \
-                     int16_t,                             \
-                     int,                                 \
-                     int8_t,                              \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }
-
 #define PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name##_raw,                           \
                      KPS,                                  \
@@ -191,10 +175,9 @@ PD_REGISTER_KERNEL(less_than_raw,
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);  \
   }
 
-PD_REGISTER_COMPARE_RAW_KERNEL(less_equal, LessEqual)
-PD_REGISTER_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
-PD_REGISTER_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
-
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
 PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(equal, Equal)
 PD_REGISTER_COMPLEX_COMPARE_RAW_KERNEL(not_equal, NotEqual)
 
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
@@ -668,8 +668,8 @@ def greater_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
+        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
         name (str|None, optional): The default value is None.  Normally there is no need for
             user to set this property.  For more information, please refer to :ref:`api_guide_Name`.
     Returns:
@@ -704,6 +704,8 @@ def greater_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "greater_equal",
         )
@@ -721,6 +723,8 @@ def greater_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "greater_equal",
         )
@@ -758,8 +762,8 @@ def greater_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
+        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
         name (str|None, optional): The default value is None.  Normally there is no need for
             user to set this property.  For more information, please refer to :ref:`api_guide_Name`.
     Returns:
@@ -794,6 +798,8 @@ def greater_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "greater_than",
         )
@@ -811,6 +817,8 @@ def greater_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "greater_than",
         )
@@ -848,8 +856,8 @@ def less_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
+        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
         name (str|None, optional): The default value is None.  Normally there is no need for
             user to set this property.  For more information, please refer to :ref:`api_guide_Name`.
 
@@ -885,6 +893,8 @@ def less_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "less_equal",
         )
@@ -902,6 +912,8 @@ def less_equal(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "less_equal",
         )
@@ -939,8 +951,8 @@ def less_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
+        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, bfloat16, float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
         name (str|None, optional): The default value is None.  Normally there is no need for
             user to set this property.  For more information, please refer to :ref:`api_guide_Name`.
 
@@ -976,6 +988,8 @@ def less_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "less_than",
         )
@@ -993,6 +1007,8 @@ def less_than(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
                 "int32",
                 "int64",
                 "uint16",
+                "complex64",
+                "complex128",
             ],
             "less_than",
         )
diff --git a/test/legacy_test/test_compare_complex_op.py b/test/legacy_test/test_compare_complex_op.py
