[AMP OP&Test]Add fp16/bf16 support isnan/isfinite/isinf op

paddle/phi/kernels/funcs/isfinite_functor.h

@@ -45,6 +45,13 @@ struct IsNanFunctor<phi::dtype::float16, void> {
   }
 };
 
+template <>
+struct IsNanFunctor<phi::dtype::bfloat16, void> {
+  HOSTDEVICE bool operator()(const phi::dtype::bfloat16& a) const {
+    return phi::dtype::isnan(a);
+  }
+};
+
 template <typename T, class Enable = void>
 struct IsInfFunctor {
   HOSTDEVICE bool operator()(const T& a) const {
@@ -69,6 +76,13 @@ struct IsInfFunctor<phi::dtype::float16, void> {
   }
 };
 
+template <>
+struct IsInfFunctor<phi::dtype::bfloat16, void> {
+  HOSTDEVICE bool operator()(const phi::dtype::bfloat16& a) const {
+    return phi::dtype::isinf(a);
+  }
+};
+
 template <typename T, class Enable = void>
 struct IsFiniteFunctor {
   HOSTDEVICE bool operator()(const T& a) const {
@@ -94,5 +108,12 @@ struct IsFiniteFunctor<phi::dtype::float16, void> {
   }
 };
 
+template <>
+struct IsFiniteFunctor<phi::dtype::bfloat16, void> {
+  HOSTDEVICE bool operator()(const phi::dtype::bfloat16& a) const {
+    return phi::dtype::isfinite(a);
+  }
+};
+
 }  // namespace funcs
 }  // namespace phi

paddle/phi/kernels/gpu/isfinite_kernel.cu

@@ -25,6 +25,7 @@ PD_REGISTER_KERNEL(isinf,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
@@ -37,6 +38,7 @@ PD_REGISTER_KERNEL(isnan,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
@@ -49,6 +51,7 @@ PD_REGISTER_KERNEL(isfinite,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);

python/paddle/fluid/tests/unittests/test_isfinite_op.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 from eager_op_test import OpTest
+from op_test import convert_float_to_uint16
 
 from paddle.fluid import core
 
@@ -48,6 +49,23 @@ def init_dtype(self):
         self.dtype = np.float16
 
 
+# BFP16 isinf Test
+class TestInfBF16(OpTest):
+    def setUp(self):
+        self.op_type = "isinf"
+        self.dtype = np.uint16
+        x = np.random.uniform(0.1, 1, [11, 17]).astype(np.float32)
+        x[0] = np.inf
+        x[-1] = np.inf
+
+        out = np.isinf(x).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(x)}
+        self.outputs = {'Out': convert_float_to_uint16(out)}
+
+    def test_output(self):
+        self.check_output()
+
+
 class TestNAN(OpTest):
     def setUp(self):
         self.op_type = "isnan"
@@ -76,6 +94,23 @@ def init_dtype(self):
         self.dtype = np.float16
 
 
+# BFP16 isnan Test
+class TestNANBF16(OpTest):
+    def setUp(self):
+        self.op_type = "isnan"
+        self.dtype = np.uint16
+        x = np.random.uniform(0.1, 1, [11, 17]).astype(np.float32)
+        x[0] = np.nan
+        x[-1] = np.nan
+
+        out = np.isnan(x).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(x)}
+        self.outputs = {'Out': convert_float_to_uint16(out)}
+
+    def test_output(self):
+        self.check_output()
+
+
 class TestIsfinite(OpTest):
     def setUp(self):
         self.op_type = "isfinite"
@@ -105,5 +140,22 @@ def init_dtype(self):
         self.dtype = np.float16
 
 
+# BFP16 isfinite Test
+class TestIsfiniteBF16(OpTest):
+    def setUp(self):
+        self.op_type = "isfinite"
+        self.dtype = np.uint16
+        x = np.random.uniform(0.1, 1, [11, 17]).astype(np.float32)
+        x[0] = np.inf
+        x[-1] = np.nan
+
+        out = np.isfinite(x).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(x)}
+        self.outputs = {'Out': convert_float_to_uint16(out)}
+
+    def test_output(self):
+        self.check_output()
+
+
 if __name__ == '__main__':
     unittest.main()

python/paddle/tensor/math.py

@@ -3463,7 +3463,14 @@ def isfinite(x, name=None):
         check_variable_and_dtype(
             x,
             'x',
-            ['float16', 'float32', 'float64', 'int32', 'int64'],
+            [
+                'float16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'uint16',
+            ],
             'isfinite',
         )
         out = helper.create_variable_for_type_inference('bool')
@@ -3499,7 +3506,17 @@ def isinf(x, name=None):
     else:
         helper = LayerHelper("isinf_v2", **locals())
         check_variable_and_dtype(
-            x, 'x', ['float16', 'float32', 'float64', 'int32', 'int64'], 'isinf'
+            x,
+            'x',
+            [
+                'float16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'uint16',
+            ],
+            'isinf',
         )
         out = helper.create_variable_for_type_inference(dtype='bool')
         helper.append_op(type="isinf_v2", inputs={"X": x}, outputs={"Out": out})
@@ -3532,7 +3549,17 @@ def isnan(x, name=None):
     else:
         helper = LayerHelper("isnan_v2", **locals())
         check_variable_and_dtype(
-            x, 'x', ['float16', 'float32', 'float64', 'int32', 'int64'], 'isnan'
+            x,
+            'x',
+            [
+                'float16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'uint16',
+            ],
+            'isnan',
         )
         out = helper.create_variable_for_type_inference(dtype='bool')
         helper.append_op(type="isnan_v2", inputs={"X": x}, outputs={"Out": out})