Decomp allclose

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -31,6 +31,7 @@
 GENERATE_IMPL_DECOMP = [
     "add_n",
     "addmm",
+    "allclose",
     "any",
     "bce_loss",
     "bmm",

paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h

@@ -1482,6 +1482,36 @@ Tensor diag_decomp(const Tensor& x,
   return ConvertToOrig<T>(res, x.dtype());
 }
 
+template <typename T>
+Tensor allclose_decomp(const Tensor& x,
+                       const Tensor& y,
+                       const paddle::Scalar& rtol,
+                       const paddle::Scalar& atol,
+                       const bool equal_nan) {
+  Tensor left = abs<T>(x - y);
+  Tensor min_diff_tensor;
+  if (has_dynamic_shape(y.shape())) {
+    min_diff_tensor =
+        backend::full_with_tensor<T>(shape64<T>(y), 1e-15, y.dtype());
+  } else {
+    min_diff_tensor = full<T>(y.shape(), 1e-15, y.dtype());
+  }
+  Tensor rtol_tensor = full_scalar<T>(rtol.to<double>(), y.dtype());
+  Tensor atol_tensor = full_scalar<T>(atol.to<double>(), y.dtype());
+  Tensor right = atol_tensor + rtol_tensor * y;
+  Tensor diff = abs<T>(right - left);
+  Tensor res_tmp = backend::logical_or<T>(less_equal<T>(left, right),
+                                          less_equal<T>(diff, min_diff_tensor));
+  Tensor res = backend::logical_or<T>(equal<T>(x, y), res_tmp);
+  if (equal_nan) {
+    Tensor x_nan = isnan<T>(x);
+    Tensor y_nan = isnan<T>(y);
+    res = backend::logical_or<T>(
+        res, backend::logical_or<T>(backend::logical_not<T>(x_nan), y_nan));
+  }
+  return backend::all<T>(res);
+}
+
 }  // namespace details
 
 }  // namespace primitive

test/legacy_test/test_allclose_op.py

@@ -33,14 +33,18 @@ def set_args(self):
     def setUp(self):
         self.set_args()
         self.op_type = "allclose"
+        self.prim_op_type = "comp"
         self.python_api = paddle.allclose
+        self.public_python_api = paddle.allclose
         self.inputs = {
             'Input': self.input,
             'Other': self.other,
+        }
+        self.attrs = {
             "Rtol": self.rtol,
             "Atol": self.atol,
+            'equal_nan': self.equal_nan,
         }
-        self.attrs = {'equal_nan': self.equal_nan}
         self.outputs = {
             'Out': np.array(
                 np.allclose(
@@ -54,7 +58,7 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
 
 class TestAllcloseOpException(TestAllcloseOp):
@@ -226,7 +230,7 @@ class TestAllcloseOpFloat64(TestAllcloseOp):
     def set_args(self):
         self.input = np.array([10.1]).astype("float64")
         self.other = np.array([10]).astype("float64")
-        self.rtol = np.array([0.01]).astype("float64")
+        self.rtol = np.array([0.001]).astype("float64")
         self.atol = np.array([0]).astype("float64")
         self.equal_nan = False
 

test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py

@@ -41,6 +41,14 @@ def log_softmax_net(x):
     return F.log_softmax(x)
 
 
+def allclose_net1(x, y):
+    return paddle.allclose(x, y, 1e-05, 1e-08, False)
+
+
+def allclose_net2(x, y):
+    return paddle.allclose(x, y, 1e-05, 1e-08, True)
+
+
 def any_net(x):
     return paddle.any(x)
 
@@ -938,6 +946,54 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimAllclose1(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [30, 50]
+        self.shape_y = [30, 50]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        num_nan = int(0.1 * np.prod(self.shape_x))
+        indices = np.random.choice(
+            np.prod(self.shape_x), num_nan, replace=False
+        )
+        self.x.ravel()[indices] = np.nan
+        self.y.ravel()[indices] = np.nan
+        self.net = allclose_net1
+        self.necessary_ops = "pd_op.allclose"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAllclose2(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [50, 20]
+        self.shape_y = [50, 20]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        num_nan = int(0.2 * np.prod(self.shape_x))
+        indices = np.random.choice(
+            np.prod(self.shape_x), num_nan, replace=False
+        )
+        self.x.ravel()[indices] = np.nan
+        self.y.ravel()[indices] = np.nan
+        self.net = allclose_net2
+        self.necessary_ops = "pd_op.allclose"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimLerp1(TestPrimThree):
     def setUp(self):
         np.random.seed(2023)