From 7a8635d82300278916dbf91e812b87fdfc9a30c7 Mon Sep 17 00:00:00 2001
From: co63oc <co63oc@users.noreply.github.com>
Date: Wed, 10 May 2023 11:42:43 +0800
Subject: [PATCH] Add elementwise_heaviside tests (#53549)

---
 .../kernels/gpu/elementwise_grad_kernel.cu    |  1 +
 paddle/phi/kernels/kps/elementwise_kernel.cu  |  1 +
 .../test_elementwise_heaviside_op.py          | 54 +++++++++++++++++--
 3 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu b/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
index efd15adc383..c0646e71a9a 100644
--- a/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
+++ b/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
@@ -139,6 +139,7 @@ PD_REGISTER_KERNEL(heaviside_grad,
                    double,
                    int,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int64_t) {}
 
 PD_REGISTER_KERNEL(elementwise_pow_grad,
diff --git a/paddle/phi/kernels/kps/elementwise_kernel.cu b/paddle/phi/kernels/kps/elementwise_kernel.cu
index e56cb4d8459..e1eb5f85288 100644
--- a/paddle/phi/kernels/kps/elementwise_kernel.cu
+++ b/paddle/phi/kernels/kps/elementwise_kernel.cu
@@ -177,5 +177,6 @@ PD_REGISTER_KERNEL(heaviside,
                    double,
                    int,
                    float16,
+                   bfloat16,
                    int64_t) {}
 #endif
diff --git a/python/paddle/fluid/tests/unittests/test_elementwise_heaviside_op.py b/python/paddle/fluid/tests/unittests/test_elementwise_heaviside_op.py
index 06439b4962c..a034c74fb84 100644
--- a/python/paddle/fluid/tests/unittests/test_elementwise_heaviside_op.py
+++ b/python/paddle/fluid/tests/unittests/test_elementwise_heaviside_op.py
@@ -15,15 +15,19 @@
 import unittest
 
 import numpy as np
-from eager_op_test import OpTest
+from eager_op_test import OpTest, convert_float_to_uint16
 
 import paddle
+from paddle.fluid import core
 
 
-def Heaviside_grad(x, y, dout):
-    tmp = np.zeros(x.shape).astype("float16")
+def Heaviside_grad(x, y, dout, astype="float16", is_bfloat16=False):
+    tmp = np.zeros(x.shape).astype(astype)
     dx = np.multiply(tmp, dout)
-    dy = np.multiply(np.equal(x, 0), dout).astype("float16")
+    dy = np.multiply(np.equal(x, 0), dout).astype(astype)
+    if is_bfloat16:
+        dx = convert_float_to_uint16(dx)
+        dy = convert_float_to_uint16(dy)
     return dx, dy
 
 
@@ -161,7 +165,7 @@ class TestHeavisideAPI_int32(TestHeavisideAPI_float64):
         self.dtype = "int32"
 
 
-class TestHeavisideAPI_float16(OpTest):
+class TestHeavisideFP16Op(OpTest):
     def setUp(self):
         self.dtype = np.float16
         self.op_type = "elementwise_heaviside"
@@ -185,6 +189,46 @@ class TestHeavisideAPI_float16(OpTest):
         )
 
 
+@unittest.skipIf(
+    not core.is_compiled_with_cuda()
+    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
+    "core is not compiled with CUDA or not support bfloat16",
+)
+class TestHeavisideBF16Op(OpTest):
+    def setUp(self):
+        self.dtype = np.uint16
+        self.np_dtype = np.float32
+        self.op_type = "elementwise_heaviside"
+        self.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),
+        }
+        self.outputs = {'Out': np.heaviside(self.inputs['X'], self.inputs['Y'])}
+
+        self.place = core.CUDAPlace(0)
+        self.inputs['X'] = convert_float_to_uint16(self.inputs['X'])
+        self.inputs['Y'] = convert_float_to_uint16(self.inputs['Y'])
+        self.outputs['Out'] = convert_float_to_uint16(self.outputs['Out'])
+
+    def test_check_output(self):
+        self.check_output_with_place(self.place)
+
+    def test_check_grad(self):
+        self.check_grad_with_place(
+            self.place,
+            ['X', 'Y'],
+            'Out',
+            user_defined_grads=Heaviside_grad(
+                self.inputs['X'],
+                self.inputs['Y'],
+                1 / self.inputs['X'].size,
+                self.np_dtype,
+                True,
+            ),
+        )
+
+
 class TestHeavisideError(unittest.TestCase):
     def test_input(self):
         paddle.disable_static()
-- 
GitLab