forwad decomp the heaviside op and support dynamic shape for it (#66795)

zeroRains · web-flow · commit 11afd45b0bac · 2024-08-02T14:40:30.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -35,6 +35,7 @@
     "gelu",
     "hardswish",
     "hardsigmoid",
+    "heaviside",
     "group_norm",
     "index_sample",
     "index_select",
@@ -84,6 +85,7 @@
     "gelu",
     "hardswish",
     "hardsigmoid",
+    "heaviside",
     "group_norm",
     "index_sample",
     "index_select",
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -925,6 +925,27 @@ Tensor hardswish_decomp(const Tensor& x) {
   return (minimum_out * x) / full_scalar<T>(SCALE, x.dtype());
 }
 
+template <typename T>
+Tensor heaviside_decomp(const Tensor& x, const Tensor& y) {
+  Tensor zero, one;
+  if (has_dynamic_shape(x.shape())) {
+    Tensor zero_x = backend::full_with_tensor<T>(shape<T>(x), 0.0, x.dtype());
+    Tensor zero_y = backend::full_with_tensor<T>(shape<T>(y), 0.0, x.dtype());
+    zero = zero_x + zero_y;
+    one = backend::full_with_tensor<T>(shape<T>(zero), 1.0, x.dtype());
+  } else {
+    auto out_dims = phi::funcs::BroadcastTwoDims(x.dims(), y.dims());
+    zero = full<T>(phi::vectorize(out_dims), 0.0, x.dtype());
+    one = full<T>(phi::vectorize(out_dims), 1.0, x.dtype());
+  }
+  Tensor broadcast_x = x + zero;
+  Tensor broadcast_y = y + zero;
+  Tensor res = where<T>(broadcast_x > zero, one, broadcast_x);
+  res = where<T>(broadcast_x == zero, broadcast_y, res);
+  res = where<T>(broadcast_x < zero, zero, res);
+  return res;
+}
+
 template <typename T>
 Tensor leaky_relu_decomp(const Tensor& x, float negative_slope) {
   auto multiply_tmp = full_scalar<T>(negative_slope, x.dtype()) * x;
diff --git a/test/legacy_test/test_elementwise_heaviside_op.py b/test/legacy_test/test_elementwise_heaviside_op.py
@@ -38,20 +38,34 @@ def setUp(self):
         x = np.random.random((13, 17)).astype("float64")
         y = np.random.random((13, 17)).astype("float64")
         self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
         self.inputs = {'X': x, 'Y': y}
         self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X', 'Y'], 'Out', check_pir=True)
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True, check_prim_pir=True)
 
     def test_check_grad_ignore_x(self):
-        self.check_grad(['Y'], 'Out', no_grad_set=set("X"), check_pir=True)
+        self.check_grad(
+            ['Y'],
+            'Out',
+            no_grad_set=set("X"),
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad_ignore_y(self):
-        self.check_grad(['X'], 'Out', no_grad_set=set('Y'), check_pir=True)
+        self.check_grad(
+            ['X'],
+            'Out',
+            no_grad_set=set('Y'),
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestHeavisideBroadcast(unittest.TestCase):
@@ -167,19 +181,45 @@ def setUp(self):
         self.dtype = "int32"
 
 
+class TestElementwiseOp1(TestElementwiseOp):
+    def setUp(self):
+        self.op_type = "elementwise_heaviside"
+        x = np.random.random(100).astype("float64")
+        y = np.random.random((13, 100)).astype("float64")
+        self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
+        self.inputs = {'X': x, 'Y': y}
+        self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
+
+
+class TestElementwiseOp2(TestElementwiseOp):
+    def setUp(self):
+        self.op_type = "elementwise_heaviside"
+        x = np.random.random((13, 100)).astype("float64")
+        y = np.random.random(100).astype("float64")
+        self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
+        self.inputs = {'X': x, 'Y': y}
+        self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
+
+
 class TestHeavisideFP16Op(OpTest):
     def setUp(self):
         self.dtype = np.float16
         self.op_type = "elementwise_heaviside"
         self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
         self.inputs = {
             'X': np.random.uniform(1, 2, [20, 5]).astype("float16"),
             'Y': np.random.uniform(1, 2, [20, 5]).astype("float16"),
         }
         self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
@@ -189,6 +229,7 @@ def test_check_grad(self):
                 self.inputs['X'], self.inputs['Y'], 1 / self.inputs['X'].size
             ),
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
@@ -203,6 +244,8 @@ def setUp(self):
         self.np_dtype = np.float32
         self.op_type = "elementwise_heaviside"
         self.python_api = paddle.heaviside
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.heaviside
         self.inputs = {
             'X': np.random.uniform(1, 2, [20, 5]).astype(self.np_dtype),
             'Y': np.random.uniform(1, 2, [20, 5]).astype(self.np_dtype),
@@ -215,7 +258,9 @@ def setUp(self):
         self.outputs['Out'] = convert_float_to_uint16(self.outputs['Out'])
 
     def test_check_output(self):
-        self.check_output_with_place(self.place, check_pir=True)
+        self.check_output_with_place(
+            self.place, check_pir=True, check_prim_pir=True
+        )
 
     def test_check_grad(self):
         self.check_grad_with_place(
@@ -230,6 +275,7 @@ def test_check_grad(self):
                 True,
             ),
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -200,6 +200,10 @@ def kldiv_loss_net3(x, target):
     return paddle.nn.functional.kl_div(x, target, "batchmean", True)
 
 
+def heaviside_net(x, y):
+    return paddle.heaviside(x, y)
+
+
 class TestPrimBase(unittest.TestCase):
     def setUp(self):
         np.random.seed(2023)
@@ -871,5 +875,90 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimHeaviside1(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [300, 2048]
+        self.shape_y = [300, 2048]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, 2048]
+        self.init_y_shape = [None, 2048]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.net = heaviside_net
+        self.necessary_ops = "pd_op.heaviside"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimHeaviside2(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [300, 2048]
+        self.shape_y = [1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, 2048]
+        self.init_y_shape = [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)
+        self.net = heaviside_net
+        self.necessary_ops = "pd_op.heaviside"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimHeaviside3(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [300, 2048]
+        self.shape_y = [2048]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, 2048]
+        self.init_y_shape = [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)
+        self.net = heaviside_net
+        self.necessary_ops = "pd_op.heaviside"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimHeaviside4(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2048]
+        self.shape_y = [300, 2048]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None]
+        self.init_y_shape = [300, 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)
+        self.net = heaviside_net
+        self.necessary_ops = "pd_op.heaviside"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimHeaviside5(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1]
+        self.shape_y = [300, 2048]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None]
+        self.init_y_shape = [300, 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)
+        self.net = heaviside_net
+        self.necessary_ops = "pd_op.heaviside"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()
