added sqrt bf16 fwd/bwd (#38599)

PaddlePaddle · Jan 4, 2022 · 2d2609e · 2d2609e
1 parent 1e3f01e
commit 2d2609e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc
@@ -282,7 +282,6 @@ namespace ops = paddle::operators;
   __macro(swish, SwishMKLDNNFunctor, SwishMKLDNNGradFunctor);              \
   __macro(hard_swish, HardSwishMKLDNNFunctor, HardSwishMKLDNNGradFunctor); \
   __macro(tanh, TanhMKLDNNFunctor, TanhMKLDNNGradFunctor);                 \
-  __macro(sqrt, SqrtMKLDNNFunctor, SqrtMKLDNNGradFunctor);                 \
   __macro(abs, AbsMKLDNNFunctor, AbsMKLDNNGradFunctor);                    \
   __macro(elu, EluMKLDNNFunctor, EluMKLDNNGradFunctor);
 
@@ -293,6 +292,8 @@ REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(gelu, GeluMKLDNNFunctor,
                                        GeluMKLDNNGradFunctor);
 REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(sigmoid, SigmoidMKLDNNFunctor,
                                        SigmoidMKLDNNGradFunctor);
+REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(sqrt, SqrtMKLDNNFunctor,
+                                       SqrtMKLDNNGradFunctor);
 
 namespace ops = paddle::operators;
 REGISTER_OP_KERNEL(

diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py
@@ -82,6 +82,20 @@ def op_grad(self, dout, x):
         return dout * self.op_forward(x) * (1 - self.op_forward(x))
 
 
+class TestMKLDNNSqrtBF16Op(MKLDNNBF16ActivationOp, TestActivation):
+    def config(self):
+        self.op_type = "sqrt"
+
+    def init_data(self):
+        self.x = np.random.uniform(1, 2, [2, 4, 3, 5]).astype(np.float32)
+
+    def op_forward(self, x):
+        return np.sqrt(x)
+
+    def op_grad(self, dout, x):
+        return dout / (2 * np.sqrt(x))
+
+
 class TestMKLDNNGeluErfBF16Op(MKLDNNBF16ActivationOp, TestActivation):
     def config(self):
         self.op_type = "gelu"