【Hackathon 5th No.13】【关联 PR】Added uint8&int8&int16 support for compare_kernel -part (#58209)

jjyaoao · web-flow · commit d700b4cb9bae · 2023-11-23T14:18:42.000+08:00
* 【Hackathon 5th No.13】【关联 PR】Added int8 support for less_than

* update 1.0

* update test

* update v1.3

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py,test=document_fix

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py

* Update test_compare_op.py

* Update compare_kernel.cc

* Update compare_kernel.cu

* Update compare_kernel.cc

* Update compare_kernel.cu

* Update test_compare_op.py

* Update logic.py

* Update compare_kernel.cc

* Update compare_kernel.cc
diff --git a/paddle/phi/kernels/cpu/compare_kernel.cc b/paddle/phi/kernels/cpu/compare_kernel.cc
@@ -110,15 +110,18 @@ PD_REGISTER_KERNEL(equal_all,
                      ALL_LAYOUT,                          \
                      phi::func##Kernel,                   \
                      bool,                                \
-                     int16_t,                             \
                      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)
diff --git a/paddle/phi/kernels/kps/compare_kernel.cu b/paddle/phi/kernels/kps/compare_kernel.cu
@@ -150,8 +150,10 @@ PD_REGISTER_KERNEL(equal_all,
                      ALL_LAYOUT,                          \
                      phi::func##Kernel,                   \
                      bool,                                \
-                     int16_t,                             \
                      int,                                 \
+                     uint8_t,                             \
+                     int8_t,                              \
+                     int16_t,                             \
                      int64_t,                             \
                      float,                               \
                      double,                              \
diff --git a/paddle/phi/kernels/legacy/cpu/compare_kernel.cc b/paddle/phi/kernels/legacy/cpu/compare_kernel.cc
@@ -115,6 +115,8 @@ PD_REGISTER_KERNEL(less_than_raw,
                    ALL_LAYOUT,
                    phi::LessThanRawKernel,
                    bool,
+                   uint8_t,
+                   int8_t,
                    int16_t,
                    int,
                    int64_t,
@@ -131,6 +133,8 @@ PD_REGISTER_KERNEL(less_than_raw,
                      ALL_LAYOUT,                          \
                      phi::func##RawKernel,                \
                      bool,                                \
+                     uint8_t,                             \
+                     int8_t,                              \
                      int16_t,                             \
                      int,                                 \
                      int64_t,                             \
diff --git a/paddle/phi/kernels/legacy/kps/compare_kernel.cu b/paddle/phi/kernels/legacy/kps/compare_kernel.cu
@@ -139,6 +139,8 @@ PD_REGISTER_KERNEL(less_than_raw,
                    ALL_LAYOUT,
                    phi::LessThanRawKernel,
                    bool,
+                   uint8_t,
+                   int8_t,
                    int16_t,
                    int,
                    int64_t,
@@ -155,8 +157,10 @@ PD_REGISTER_KERNEL(less_than_raw,
                      ALL_LAYOUT,                          \
                      phi::func##RawKernel,                \
                      bool,                                \
+                     uint8_t,                             \
                      int16_t,                             \
                      int,                                 \
+                     int8_t,                              \
                      int64_t,                             \
                      float,                               \
                      double,                              \
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
@@ -512,8 +512,8 @@ def equal(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): Tensor, data type is bool, float16, float32, float64, int32, int64.
-        y (Tensor): Tensor, data type is bool, float16, float32, float64, int32, int64.
+        x (Tensor): Tensor, data type is bool, float16, float32, float64, uint8, int8, int16, int32, int64.
+        y (Tensor): Tensor, data type is bool, float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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`.
 
@@ -551,6 +551,9 @@ def equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -565,6 +568,9 @@ def equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -609,8 +615,8 @@ def greater_equal(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, 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, float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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:
@@ -639,6 +645,9 @@ def greater_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -653,6 +662,9 @@ def greater_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -697,8 +709,8 @@ def greater_than(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, 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, float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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:
@@ -727,6 +739,9 @@ def greater_than(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -741,6 +756,9 @@ def greater_than(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -785,8 +803,8 @@ def less_equal(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, 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, float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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`.
 
@@ -816,6 +834,9 @@ def less_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -830,6 +851,9 @@ def less_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -874,8 +898,8 @@ def less_than(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, float16, float32, float64, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, 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, float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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`.
 
@@ -905,6 +929,9 @@ def less_than(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -919,6 +946,9 @@ def less_than(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -963,8 +993,8 @@ def not_equal(x, y, name=None):
         The output has no gradient.
 
     Args:
-        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, float32, float64, int32, int64.
-        y (Tensor): Second input to compare which is N-D tensor. The input data type should be bool, float32, float64, int32, int64.
+        x (Tensor): First input to compare which is N-D tensor. The input data type should be bool, 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, float32, float64, uint8, int8, int16, int32, int64.
         name (str, 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`.
 
@@ -994,6 +1024,9 @@ def not_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
@@ -1008,6 +1041,9 @@ def not_equal(x, y, name=None):
                 "float16",
                 "float32",
                 "float64",
+                "uint8",
+                "int8",
+                "int16",
                 "int32",
                 "int64",
                 "uint16",
diff --git a/test/legacy_test/test_compare_op.py b/test/legacy_test/test_compare_op.py
@@ -38,24 +38,35 @@ def setUp(self):
         def test_output(self):
             self.check_output(check_cinn=True, check_pir=check_pir)
 
-        def test_errors(self):
+        def test_int16_support(self):
             paddle.enable_static()
             with paddle.static.program_guard(
                 paddle.static.Program(), paddle.static.Program()
             ):
-                x = paddle.static.data(name='x', shape=[-1, 2], dtype='int32')
-                y = paddle.static.data(name='y', shape=[-1, 2], dtype='int32')
                 a = paddle.static.data(name='a', shape=[-1, 2], dtype='int16')
+                b = paddle.static.data(name='b', shape=[-1, 2], dtype='int16')
                 op = eval("paddle.%s" % self.op_type)
-                self.assertRaises(TypeError, op, x=x, y=a)
-                self.assertRaises(TypeError, op, x=a, y=y)
+
+                try:
+                    result = op(x=a, y=b)
+                except TypeError:
+                    self.fail("TypeError should not be raised for int16 inputs")
 
     cls_name = f"{op_type}_{typename}"
     Cls.__name__ = cls_name
     globals()[cls_name] = Cls
 
 
-for _type_name in {'float32', 'float64', 'int32', 'int64', 'float16'}:
+for _type_name in {
+    'float32',
+    'float64',
+    'uint8',
+    'int8',
+    'int16',
+    'int32',
+    'int64',
+    'float16',
+}:
     if _type_name == 'float64' and core.is_compiled_with_rocm():
         _type_name = 'float32'
     if _type_name == 'float16' and (not core.is_compiled_with_cuda()):
@@ -513,7 +524,7 @@ def test_check_output(self):
 
 
 class TestCompareOpError(unittest.TestCase):
-    def test_errors(self):
+    def test_int16_support(self):
         paddle.enable_static()
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
