add cpu kernel

PaddlePaddle · Feb 9, 2022 · f907a86 · f907a86
1 parent 160725a
commit f907a86
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 31 deletions.
diff --git a/paddle/fluid/operators/squeeze_op.cc b/paddle/fluid/operators/squeeze_op.cc
@@ -393,7 +393,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::SqueezeKernel<paddle::platform::CPUDeviceContext,
                        paddle::platform::complex<float>>,
     ops::SqueezeKernel<paddle::platform::CPUDeviceContext,
-                       paddle::platform::complex<double>>);
+                       paddle::platform::complex<double>>,
+    ops::SqueezeKernel<paddle::platform::CPUDeviceContext,
+                       paddle::platform::bfloat16>);
 REGISTER_OP_CPU_KERNEL(
     squeeze_grad,
     ops::SqueezeGradKernel<paddle::platform::CPUDeviceContext, float>,
@@ -406,7 +408,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::SqueezeGradKernel<paddle::platform::CPUDeviceContext,
                            paddle::platform::complex<float>>,
     ops::SqueezeGradKernel<paddle::platform::CPUDeviceContext,
-                           paddle::platform::complex<double>>);
+                           paddle::platform::complex<double>>,
+    ops::SqueezeGradKernel<paddle::platform::CPUDeviceContext,
+                           paddle::platform::bfloat16>);
 
 REGISTER_OP_CPU_KERNEL(
     squeeze2, ops::Squeeze2Kernel<paddle::platform::CPUDeviceContext, float>,
@@ -419,7 +423,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::Squeeze2Kernel<paddle::platform::CPUDeviceContext,
                         paddle::platform::complex<float>>,
     ops::Squeeze2Kernel<paddle::platform::CPUDeviceContext,
-                        paddle::platform::complex<double>>);
+                        paddle::platform::complex<double>>,
+    ops::Squeeze2Kernel<paddle::platform::CPUDeviceContext,
+                        paddle::platform::bfloat16>);
 
 REGISTER_OP_CPU_KERNEL(
     squeeze2_grad,
@@ -433,4 +439,6 @@ REGISTER_OP_CPU_KERNEL(
     ops::Squeeze2GradKernel<paddle::platform::CPUDeviceContext,
                             paddle::platform::complex<float>>,
     ops::Squeeze2GradKernel<paddle::platform::CPUDeviceContext,
-                            paddle::platform::complex<double>>);
+                            paddle::platform::complex<double>>,
+    ops::Squeeze2GradKernel<paddle::platform::CPUDeviceContext,
+                            paddle::platform::bfloat16>);
diff --git a/paddle/fluid/operators/stack_op.cc b/paddle/fluid/operators/stack_op.cc
@@ -173,13 +173,16 @@ REGISTER_OPERATOR(stack, ops::StackOp, ops::StackOpMaker,
                   ops::StackGradOpMaker<paddle::imperative::OpBase>);
 REGISTER_OPERATOR(stack_grad, ops::StackOpGrad);
 
-REGISTER_OP_CPU_KERNEL(stack, ops::StackKernel<plat::CPUDeviceContext, float>,
-                       ops::StackKernel<plat::CPUDeviceContext, double>,
-                       ops::StackKernel<plat::CPUDeviceContext, int>,
-                       ops::StackKernel<plat::CPUDeviceContext, int64_t>);
-
-REGISTER_OP_CPU_KERNEL(stack_grad,
-                       ops::StackGradKernel<plat::CPUDeviceContext, float>,
-                       ops::StackGradKernel<plat::CPUDeviceContext, double>,
-                       ops::StackGradKernel<plat::CPUDeviceContext, int>,
-                       ops::StackGradKernel<plat::CPUDeviceContext, int64_t>);
+REGISTER_OP_CPU_KERNEL(
+    stack, ops::StackKernel<plat::CPUDeviceContext, float>,
+    ops::StackKernel<plat::CPUDeviceContext, double>,
+    ops::StackKernel<plat::CPUDeviceContext, int>,
+    ops::StackKernel<plat::CPUDeviceContext, int64_t>,
+    ops::StackKernel<plat::CPUDeviceContext, paddle::platform::bfloat16>);
+
+REGISTER_OP_CPU_KERNEL(
+    stack_grad, ops::StackGradKernel<plat::CPUDeviceContext, float>,
+    ops::StackGradKernel<plat::CPUDeviceContext, double>,
+    ops::StackGradKernel<plat::CPUDeviceContext, int>,
+    ops::StackGradKernel<plat::CPUDeviceContext, int64_t>,
+    ops::StackGradKernel<plat::CPUDeviceContext, paddle::platform::bfloat16>);
diff --git a/paddle/fluid/operators/unsqueeze_op.cc b/paddle/fluid/operators/unsqueeze_op.cc
@@ -366,7 +366,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
                          paddle::platform::complex<float>>,
     ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
-                         paddle::platform::complex<double>>);
+                         paddle::platform::complex<double>>,
+    ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
+                         paddle::platform::bfloat16>);
 REGISTER_OP_CPU_KERNEL(
     unsqueeze_grad,
     ops::UnsqueezeGradKernel<paddle::platform::CPUDeviceContext, float>,
@@ -379,7 +381,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::UnsqueezeGradKernel<paddle::platform::CPUDeviceContext,
                              paddle::platform::complex<float>>,
     ops::UnsqueezeGradKernel<paddle::platform::CPUDeviceContext,
-                             paddle::platform::complex<double>>);
+                             paddle::platform::complex<double>>,
+    ops::UnsqueezeGradKernel<paddle::platform::CPUDeviceContext,
+                             paddle::platform::bfloat16>);
 REGISTER_OP_CPU_KERNEL(
     unsqueeze2, ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext, float>,
     ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext, double>,
@@ -391,7 +395,9 @@ REGISTER_OP_CPU_KERNEL(
     ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
                          paddle::platform::complex<float>>,
     ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
-                         paddle::platform::complex<double>>);
+                         paddle::platform::complex<double>>,
+    ops::UnsqueezeKernel<paddle::platform::CPUDeviceContext,
+                         paddle::platform::bfloat16>);
 REGISTER_OP_CPU_KERNEL(
     unsqueeze2_grad,
     ops::Unsqueeze2GradKernel<paddle::platform::CPUDeviceContext, float>,
@@ -404,4 +410,6 @@ REGISTER_OP_CPU_KERNEL(
     ops::Unsqueeze2GradKernel<paddle::platform::CPUDeviceContext,
                               paddle::platform::complex<float>>,
     ops::Unsqueeze2GradKernel<paddle::platform::CPUDeviceContext,
-                              paddle::platform::complex<double>>);
+                              paddle::platform::complex<double>>,
+    ops::Unsqueeze2GradKernel<paddle::platform::CPUDeviceContext,
+                              paddle::platform::bfloat16>);
diff --git a/python/paddle/fluid/tests/unittests/test_squeeze_op.py b/python/paddle/fluid/tests/unittests/test_squeeze_op.py
@@ -62,12 +62,10 @@ def setUp(self):
         self.outputs = {"Out": convert_float_to_uint16(out)}
 
     def test_check_output(self):
-        if core.is_compiled_with_cuda():
-            self.check_output_with_place(core.CUDAPlace(0))
+        self.check_output()
 
     def test_check_grad(self):
-        if core.is_compiled_with_cuda():
-            self.check_grad_with_place(core.CUDAPlace(0), ["X"], "Out")
+        self.check_grad(["X"], "Out")
 
     def init_test_case(self):
         self.ori_shape = (1, 3, 1, 40)

diff --git a/python/paddle/fluid/tests/unittests/test_stack_op.py b/python/paddle/fluid/tests/unittests/test_stack_op.py
@@ -128,13 +128,10 @@ def setUp(self):
         self.attrs = {'axis': self.axis}
 
     def test_check_output(self):
-        if core.is_compiled_with_cuda():
-            self.check_output_with_place(core.CUDAPlace(0))
+        self.check_output()
 
     def test_check_grad(self):
-        if core.is_compiled_with_cuda():
-            self.check_grad_with_place(
-                core.CUDAPlace(0), self.get_x_names(), 'Y')
+        self.check_grad(self.get_x_names(), 'Y')
 
 
 class TestStackAPIWithLoDTensorArray(unittest.TestCase):

diff --git a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py
@@ -61,12 +61,10 @@ def setUp(self):
         self.outputs = {"Out": convert_float_to_uint16(out)}
 
     def test_check_output(self):
-        if core.is_compiled_with_cuda():
-            self.check_output_with_place(core.CUDAPlace(0))
+        self.check_output()
 
     def test_check_grad(self):
-        if core.is_compiled_with_cuda():
-            self.check_grad_with_place(core.CUDAPlace(0), ["X"], "Out")
+        self.check_grad(["X"], "Out")
 
     def init_test_case(self):
         self.ori_shape = (3, 40)