cinn(ops): Add symbolic isfinite/inf/nan (#63581)

hxzd5568 · web-flow · commit 12223edc788e · 2024-04-17T21:13:52.000+08:00
diff --git a/paddle/cinn/hlir/op/elementwise.cc b/paddle/cinn/hlir/op/elementwise.cc
@@ -1524,6 +1524,9 @@ CINN_REGISTER_HELPER(elementwise_ops) {
       .set_num_outputs(1)                                                     \
       .set_attr<cinn::hlir::framework::StrategyFunction>(                     \
           "CINNStrategy", cinn::hlir::op::StrategyFor##op_strategy__)         \
+      .set_attr<cinn::hlir::framework::StrategyFunctionSymbolic>(             \
+          "CINNStrategySymbolic",                                             \
+          cinn::hlir::op::StrategyFor##op_strategy__##Symbolic)               \
       .set_attr("infershape",                                                 \
                 MakeOpFunction(cinn::hlir::op::InferShapeForElementwise))     \
       .set_attr("inferdtype",                                                 \
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -77,6 +77,8 @@ OP_SAME_OPERANDS_AND_RESULT(Floor_)
 OP_SAME_OPERANDS_AND_RESULT(Imag)
 OP_SAME_OPERANDS_AND_RESULT(Increment)
 OP_SAME_OPERANDS_AND_RESULT(Increment_)
+OP_SAME_OPERANDS_AND_RESULT(Isfinite)
+OP_SAME_OPERANDS_AND_RESULT(IsfiniteSr)
 OP_SAME_OPERANDS_AND_RESULT(Isinf)
 OP_SAME_OPERANDS_AND_RESULT(IsinfSr)
 OP_SAME_OPERANDS_AND_RESULT(Isnan)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -68,6 +68,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Floor_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Imag)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Increment)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Increment_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isfinite)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IsfiniteSr)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isinf)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IsinfSr)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isnan)
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -1528,6 +1528,7 @@
   kernel :
     func : isfinite {dense -> dense},
            isfinite_sr {selected_rows -> selected_rows}
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : isinf
   args : (Tensor x)
diff --git a/test/ir/pir/cinn/symbolic/test_cinn_elementwise_symbolic.py b/test/ir/pir/cinn/symbolic/test_cinn_elementwise_symbolic.py
@@ -28,6 +28,18 @@ def tril(x):
     return paddle.tril(x)
 
 
+def isinf(x):
+    return paddle.isinf(x)
+
+
+def isfinite(x):
+    return paddle.isfinite(x)
+
+
+def isnan(x):
+    return paddle.isnan(x)
+
+
 def tril_diag_neg(x):
     return paddle.tril(x, -1)
 
@@ -190,5 +202,113 @@ def test_eval_symbolic(self):
         np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
 
 
+class TestCinnSubGrapIsInf(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x_shape = [32, 32]
+        self.x = paddle.randn(self.x_shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 1)
+
+    def eval_symbolic(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSubGraphNet(isinf)
+        input_spec = [
+            InputSpec(shape=[None, 32], dtype='float32'),
+        ]
+        net = utils.apply_to_static(net, use_cinn, input_spec)
+        net.eval()
+        out = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return out
+
+    def test_eval_symbolic(self):
+        cinn_out = self.eval_symbolic(use_cinn=True)
+        dy_out = self.eval_symbolic(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
+class TestCinnSubGrapIsFinite(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x_shape = [32, 32]
+        self.x = paddle.randn(self.x_shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 1)
+
+    def eval_symbolic(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSubGraphNet(isfinite)
+        input_spec = [
+            InputSpec(shape=[None, 32], dtype='float32'),
+        ]
+        net = utils.apply_to_static(net, use_cinn, input_spec)
+        net.eval()
+        out = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return out
+
+    def test_eval_symbolic(self):
+        cinn_out = self.eval_symbolic(use_cinn=True)
+        dy_out = self.eval_symbolic(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
+class TestCinnSubGrapIsNan(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x_shape = [32, 32]
+        self.x = paddle.randn(self.x_shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 1)
+
+    def eval_symbolic(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSubGraphNet(isnan)
+        input_spec = [
+            InputSpec(shape=[None, 32], dtype='float32'),
+        ]
+        net = utils.apply_to_static(net, use_cinn, input_spec)
+        net.eval()
+        out = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return out
+
+    def test_eval_symbolic(self):
+        cinn_out = self.eval_symbolic(use_cinn=True)
+        dy_out = self.eval_symbolic(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
 if __name__ == '__main__':
     unittest.main()
