support bool/int32/int64 for paddle.allclose kernel (#69012)

HydrogenSulfate · web-flow · commit 77b2af500aef · 2024-10-29T15:12:19.000+08:00
* support int32/int64 for paddle.allclose kernel

* fix

* support bool and add test

* fix isnan for integral

* fix place
diff --git a/paddle/phi/kernels/cpu/allclose_kernel.cc b/paddle/phi/kernels/cpu/allclose_kernel.cc
@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/allclose_kernel.h"
 
 #include <cmath>
+#include <type_traits>
 
 #include "glog/logging.h"
 #include "paddle/phi/core/enforce.h"
@@ -56,24 +57,46 @@ void AllCloseKernel(const Context& dev_ctx,
   *out_data = true;
 
   auto num = x.numel();
-  for (int64_t i = 0; i < num; ++i) {
-    const T a = in_a[i], b = in_b[i];
-    bool val = false;
-    if (std::isnan(a) || std::isnan(b)) {
-      val = equal_nan && std::isnan(a) == std::isnan(b);
-    } else {
-      T left = (a > b ? a - b : b - a);
-      T right = atol_v + (b > 0 ? rtol_v * b : (-rtol_v) * b);
-      T diff = (left > right ? left - right : right - left);
-      val = a == b || left <= right || diff <= 1e-15;
+  if (std::is_same<T, int32_t>::value || std::is_same<T, int64_t>::value ||
+      std::is_same<T, bool>::value) {
+    for (int64_t i = 0; i < num; ++i) {
+      const double a = static_cast<double>(in_a[i]),
+                   b = static_cast<double>(in_b[i]);
+      double left = (a > b ? a - b : b - a);
+      double right = atol_v + (b > 0 ? rtol_v * b : (-rtol_v) * b);
+      double diff = (left > right ? left - right : right - left);
+      bool val = a == b || left <= right || diff <= 1e-15;
+      *out_data &= val;
+    }
+  } else {
+    for (int64_t i = 0; i < num; ++i) {
+      const T a = in_a[i], b = in_b[i];
+      bool val = false;
+      if (std::isnan(static_cast<double>(a)) ||
+          std::isnan(static_cast<double>(b))) {
+        val = equal_nan && std::isnan(static_cast<double>(a)) ==
+                               std::isnan(static_cast<double>(b));
+      } else {
+        T left = (a > b ? a - b : b - a);
+        T right = atol_v + (b > 0 ? rtol_v * b : (-rtol_v) * b);
+        T diff = (left > right ? left - right : right - left);
+        val = a == b || left <= right || diff <= 1e-15;
+      }
+      *out_data &= val;
     }
-    *out_data &= val;
   }
 }
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    allclose, CPU, ALL_LAYOUT, phi::AllCloseKernel, float, double) {
+PD_REGISTER_KERNEL(allclose,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::AllCloseKernel,
+                   float,
+                   double,
+                   bool,
+                   int,
+                   int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
diff --git a/paddle/phi/kernels/gpu/allclose_kernel.cu b/paddle/phi/kernels/gpu/allclose_kernel.cu
@@ -14,6 +14,7 @@
 
 #include "paddle/phi/kernels/allclose_kernel.h"
 
+#include <type_traits>
 #include "glog/logging.h"
 
 #include "paddle/phi/common/amp_type_traits.h"
@@ -33,7 +34,15 @@ __global__ void AllcloseCUDAKernel(const T* in_data,
                                    bool* out_data) {
   unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
   bool val;
-  using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
+  using BaseMPType = typename phi::dtype::MPTypeTrait<T>::Type;
+
+  using MPType =
+      typename std::conditional<std::is_same<T, int32_t>::value ||
+                                    std::is_same<T, int64_t>::value ||
+                                    std::is_same<T, bool>::value,
+                                double,
+                                BaseMPType>::type;
+
   for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
     const MPType a = static_cast<MPType>(in_data[i]);
     const MPType b = static_cast<MPType>(other_data[i]);
@@ -102,6 +111,9 @@ PD_REGISTER_KERNEL(allclose,
                    phi::AllCloseKernel,
                    float,
                    double,
+                   bool,
+                   int,
+                   int64_t,
                    phi::dtype::float16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
@@ -499,10 +499,16 @@ def allclose(
         return _C_ops.allclose(x, y, rtol, atol, equal_nan)
     elif in_pir_mode():
         check_variable_and_dtype(
-            x, "input", ['float16', 'float32', 'float64'], 'allclose'
+            x,
+            "input",
+            ['bool', 'int32', 'int64', 'float16', 'float32', 'float64'],
+            'allclose',
         )
         check_variable_and_dtype(
-            y, "input", ['float16', 'float32', 'float64'], 'allclose'
+            y,
+            "input",
+            ['bool', 'int32', 'int64', 'float16', 'float32', 'float64'],
+            'allclose',
         )
         if not isinstance(rtol, (float, paddle.pir.Value)):
             raise TypeError(
@@ -516,10 +522,16 @@ def allclose(
         return _C_ops.allclose(x, y, rtol, atol, equal_nan)
     else:
         check_variable_and_dtype(
-            x, "input", ['float16', 'float32', 'float64'], 'allclose'
+            x,
+            "input",
+            ['bool', 'int32', 'int64', 'float16', 'float32', 'float64'],
+            'allclose',
         )
         check_variable_and_dtype(
-            y, "input", ['float16', 'float32', 'float64'], 'allclose'
+            y,
+            "input",
+            ['bool', 'int32', 'int64', 'float16', 'float32', 'float64'],
+            'allclose',
         )
         check_type(rtol, 'rtol', float, 'allclose')
         check_type(atol, 'atol', float, 'allclose')
diff --git a/test/legacy_test/test_allclose_op.py b/test/legacy_test/test_allclose_op.py
@@ -16,6 +16,7 @@
 
 import numpy as np
 from op_test import OpTest
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle.base import core
@@ -133,7 +134,9 @@ def test_x_dtype():
             with paddle.static.program_guard(
                 paddle.static.Program(), paddle.static.Program()
             ):
-                x = paddle.static.data(name='x', shape=[10, 10], dtype='int32')
+                x = paddle.static.data(
+                    name='x', shape=[10, 10], dtype='complex32'
+                )
                 y = paddle.static.data(
                     name='y', shape=[10, 10], dtype='float64'
                 )
@@ -148,7 +151,9 @@ def test_y_dtype():
                 x = paddle.static.data(
                     name='x', shape=[10, 10], dtype='float64'
                 )
-                y = paddle.static.data(name='y', shape=[10, 10], dtype='int32')
+                y = paddle.static.data(
+                    name='y', shape=[10, 10], dtype='complex32'
+                )
                 result = paddle.allclose(x, y)
 
         self.assertRaises(TypeError, test_y_dtype)
@@ -226,6 +231,228 @@ def set_args(self):
         self.equal_nan = False
 
 
+class TestAllcloseOpBool(unittest.TestCase):
+    def test_close_True(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([1]).astype("bool")
+                self.other = np.array([1]).astype("bool")
+                self.rtol = np.array([0.0]).astype("float32")
+                self.atol = np.array([0.0]).astype("float32")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    True,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='bool')
+                    y = paddle.static.data(shape=[1], name='y', dtype='bool')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], True)
+
+    def test_close_False(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([0]).astype("bool")
+                self.other = np.array([1]).astype("bool")
+                self.rtol = np.array([0.0]).astype("float32")
+                self.atol = np.array([0.0]).astype("float32")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    False,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='bool')
+                    y = paddle.static.data(shape=[1], name='y', dtype='bool')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], False)
+
+
+class TestAllcloseOpInt32(unittest.TestCase):
+    def test_close_True(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([100]).astype("int32")
+                self.other = np.array([1]).astype("int32")
+                self.rtol = np.array([50.0]).astype("float32")
+                self.atol = np.array([49]).astype("float32")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    True,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='int32')
+                    y = paddle.static.data(shape=[1], name='y', dtype='int32')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], True)
+
+    def test_close_False(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([100]).astype("int32")
+                self.other = np.array([1]).astype("int32")
+                self.rtol = np.array([50.0]).astype("float32")
+                self.atol = np.array([48]).astype("float32")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    False,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='int32')
+                    y = paddle.static.data(shape=[1], name='y', dtype='int32')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], False)
+
+
+class TestAllcloseOpInt64(unittest.TestCase):
+    def test_close_True(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([100]).astype("int64")
+                self.other = np.array([1]).astype("int64")
+                self.rtol = np.array([50.0]).astype("float64")
+                self.atol = np.array([49]).astype("float64")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    True,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='int64')
+                    y = paddle.static.data(shape=[1], name='y', dtype='int64')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], True)
+
+    def test_close_False(self):
+        places = [paddle.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            with dygraph_guard():
+                # absolute(a−b)≤(atol+rtol×absolute(b))
+                self.input = np.array([100]).astype("int64")
+                self.other = np.array([1]).astype("int64")
+                self.rtol = np.array([50.0]).astype("float64")
+                self.atol = np.array([48]).astype("float64")
+                self.equal_nan = False
+                input = paddle.to_tensor(self.input, place=place)
+                other = paddle.to_tensor(self.other, place=place)
+                self.assertEqual(
+                    paddle.allclose(
+                        input, other, self.rtol, self.atol, self.equal_nan
+                    ).item(),
+                    False,
+                )
+
+            with static_guard():
+                with paddle.static.program_guard(paddle.static.Program()):
+                    x = paddle.static.data(shape=[1], name='x', dtype='int64')
+                    y = paddle.static.data(shape=[1], name='y', dtype='int64')
+                    out = paddle.allclose(
+                        x, y, self.rtol.item(), self.atol.item(), self.equal_nan
+                    )
+                    exe = paddle.static.Executor(place)
+                    exe.run(paddle.static.default_startup_program())
+                    out = exe.run(
+                        feed={'x': self.input, 'y': self.other},
+                        fetch_list=[out],
+                    )
+                    self.assertEqual(out[0], False)
+
+
 class TestAllcloseOpLargeDimInput(TestAllcloseOp):
     def set_args(self):
         self.input = np.array(np.zeros([2048, 1024])).astype("float64")
