add fp16 support for tril_triu op (#30186)

paddle/fluid/operators/tril_triu_op.cc

@@ -99,6 +99,7 @@ class TrilTriuGradOpMaker : public framework::SingleGradOpMaker<T> {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
+namespace plat = paddle::platform;
 REGISTER_OPERATOR(tril_triu, ops::TrilTriuOp, ops::TrilTriuOpMaker,
                   ops::TrilTriuGradOpMaker<paddle::framework::OpDesc>,
                   ops::TrilTriuGradOpMaker<paddle::imperative::OpBase>);
@@ -107,10 +108,13 @@ REGISTER_OP_CPU_KERNEL(
     tril_triu, ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, float>,
     ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, double>,
     ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, int>,
-    ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, int64_t>);
+    ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, int64_t>,
+    ops::TrilTriuOpKernel<paddle::platform::CPUDeviceContext, plat::float16>);
 REGISTER_OP_CPU_KERNEL(
     tril_triu_grad,
     ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext, float>,
     ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext, double>,
     ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext, int>,
-    ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext, int64_t>);
+    ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext, int64_t>,
+    ops::TrilTriuGradOpKernel<paddle::platform::CPUDeviceContext,
+                              plat::float16>);

paddle/fluid/operators/tril_triu_op.cu

@@ -15,16 +15,20 @@ limitations under the License. */
 #include "paddle/fluid/operators/tril_triu_op.h"
 
 namespace ops = paddle::operators;
+namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(
     tril_triu,
     ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, float>,
     ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, double>,
     ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, int>,
-    ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, int64_t>);
+    ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, int64_t>,
+    ops::TrilTriuOpKernel<paddle::platform::CUDADeviceContext, plat::float16>);
 REGISTER_OP_CUDA_KERNEL(
     tril_triu_grad,
     ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext, float>,
     ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext, double>,
     ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext, int>,
-    ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext, int64_t>);
+    ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext, int64_t>,
+    ops::TrilTriuGradOpKernel<paddle::platform::CUDADeviceContext,
+                              plat::float16>);

paddle/fluid/operators/tril_triu_op.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/platform/float16.h"
 #include "paddle/fluid/platform/for_range.h"
 
 namespace paddle {

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

@@ -16,8 +16,10 @@ from __future__ import print_function
 import unittest
 import numpy as np
 from op_test import OpTest
+import paddle
 import paddle.fluid as fluid
 import paddle.tensor as tensor
+from paddle.fluid.framework import Program, program_guard
 
 
 class TrilTriuOpDefaultTest(OpTest):
@@ -68,6 +70,8 @@ def case_generator(op_type, Xshape, diagonal, expected):
 
     class FailureCase(unittest.TestCase):
         def test_failure(self):
+            paddle.enable_static()
+
             data = fluid.data(shape=Xshape, dtype='float64', name=cls_name)
             with self.assertRaisesRegexp(
                     eval(expected.split(':')[-1]), errmsg[expected]):
@@ -75,6 +79,8 @@ def case_generator(op_type, Xshape, diagonal, expected):
 
     class SuccessCase(TrilTriuOpDefaultTest):
         def initTestCase(self):
+            paddle.enable_static()
+
             self.real_op_type = op_type
             self.diagonal = diagonal
             self.X = np.random.random(Xshape).astype("float64")
@@ -120,8 +126,15 @@ class TestTrilTriuOpAPI(unittest.TestCase):
     """
 
     def test_api(self):
-        data = np.random.random([1, 9, 9, 4]).astype('float32')
-        x = fluid.data(shape=[1, 9, -1, 4], dtype='float32', name='x')
+        paddle.enable_static()
+
+        dtypes = ['float16', 'float32']
+        for dtype in dtypes:
+            prog = Program()
+            startup_prog = Program()
+            with program_guard(prog, startup_prog):
+                data = np.random.random([1, 9, 9, 4]).astype(dtype)
+                x = fluid.data(shape=[1, 9, -1, 4], dtype=dtype, name='x')
                 tril_out, triu_out = tensor.tril(x), tensor.triu(x)
 
                 place = fluid.CUDAPlace(0) if fluid.core.is_compiled_with_cuda(
@@ -135,16 +148,28 @@ class TestTrilTriuOpAPI(unittest.TestCase):
                 self.assertTrue(np.allclose(triu_out, np.triu(data)))
 
     def test_api_with_dygraph(self):
+        paddle.disable_static()
+
+        dtypes = ['float16', 'float32']
+        for dtype in dtypes:
             with fluid.dygraph.guard():
-            data = np.random.random([1, 9, 9, 4]).astype('float32')
+                data = np.random.random([1, 9, 9, 4]).astype(dtype)
                 x = fluid.dygraph.to_variable(data)
-            tril_out, triu_out = tensor.tril(x).numpy(), tensor.triu(x).numpy()
+                tril_out, triu_out = tensor.tril(x).numpy(), tensor.triu(
+                    x).numpy()
                 self.assertTrue(np.allclose(tril_out, np.tril(data)))
                 self.assertTrue(np.allclose(triu_out, np.triu(data)))
 
     def test_fluid_api(self):
-        data = np.random.random([1, 9, 9, 4]).astype('float32')
-        x = fluid.data(shape=[1, 9, -1, 4], dtype='float32', name='x')
+        paddle.enable_static()
+
+        dtypes = ['float16', 'float32']
+        for dtype in dtypes:
+            prog = Program()
+            startup_prog = Program()
+            with program_guard(prog, startup_prog):
+                data = np.random.random([1, 9, 9, 4]).astype(dtype)
+                x = fluid.data(shape=[1, 9, -1, 4], dtype=dtype, name='x')
                 triu_out = fluid.layers.triu(x)
 
                 place = fluid.CUDAPlace(0) if fluid.core.is_compiled_with_cuda(

python/paddle/tensor/creation.py

@@ -558,8 +558,8 @@ def _tril_triu_op(helper):
     x = helper.kwargs.get('x', None)
 
     assert x is not None, 'x cannot be None in {}'.format(op_type)
-    check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int64'],
-                             op_type)
+    check_variable_and_dtype(
+        x, 'x', ['float16', 'float32', 'float64', 'int32', 'int64'], op_type)
     if len(x.shape) < 2:
         raise ValueError("x shape in {} must be at least 2-D".format(op_type))
     diagonal = helper.kwargs.get('diagonal', 0)