pir onednn support softmax swish (#61667)

wanghuancoder · web-flow · commit f458a92419d4 · 2024-02-07T16:35:51.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -237,9 +237,13 @@
 
 # - op : slice_grad
 
-# - op : softmax
+- op : softmax
+  extra_args : str data_format="AnyLayout", str mkldnn_data_type="float32", bool is_test=false
+  data_format_tensors : x
 
-# - op : softmax_grad
+- op : softmax_grad
+  extra_args : str data_format="AnyLayout", str mkldnn_data_type="float32", bool is_test=false
+  data_format_tensors : out, out_grad
 
 - op : softplus
 
@@ -267,9 +271,11 @@
 
 - op : sum_grad
 
-# - op : swish
+- op : swish
+  extra_args : float beta=1.0
 
-# - op : swish_grad
+- op : swish_grad
+  extra_args : float beta=1.0
 
 - op : tanh
 
diff --git a/test/legacy_test/test_softmax_op.py b/test/legacy_test/test_softmax_op.py
@@ -91,10 +91,14 @@ def test_check_output(self):
                 check_prim=True,
                 check_pir=True,
                 check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         else:
             self.check_output(
-                check_prim=True, check_pir=True, check_prim_pir=True
+                check_prim=True,
+                check_pir=True,
+                check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
     def test_check_grad(self):
@@ -110,6 +114,7 @@ def test_check_grad(self):
                     check_dygraph=(not self.use_mkldnn),
                     check_pir=True,
                     check_prim_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
         else:
             self.check_grad(
@@ -120,6 +125,7 @@ def test_check_grad(self):
                 check_prim=True,
                 check_pir=True,
                 check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
 
@@ -158,11 +164,18 @@ def test_check_output(self):
         if self.use_cudnn:
             place = core.CUDAPlace(0)
             self.check_output_with_place(
-                place, atol=1e-5, check_pir=True, check_prim_pir=True
+                place,
+                atol=1e-5,
+                check_pir=True,
+                check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         else:
             self.check_output(
-                check_prim=True, check_pir=True, check_prim_pir=True
+                check_prim=True,
+                check_pir=True,
+                check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
 
@@ -203,10 +216,14 @@ def test_check_output(self):
                 atol=1e-5,
                 check_pir=True,
                 check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         else:
             self.check_output(
-                check_prim=True, check_pir=True, check_prim_pir=True
+                check_prim=True,
+                check_pir=True,
+                check_prim_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
 
@@ -387,6 +404,7 @@ def test_check_output(self):
                     check_prim=True,
                     check_pir=True,
                     check_prim_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
     # FIXME: If the x_shape is [10, 10], gradient failed.
@@ -420,6 +438,7 @@ def test_check_output(self):
                     check_prim=True,
                     check_pir=True,
                     check_prim_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
 
@@ -472,6 +491,7 @@ def test_check_output(self):
             check_prim=True,
             check_pir=(not self.use_mkldnn),
             check_prim_pir=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad(self):
@@ -485,6 +505,7 @@ def test_check_grad(self):
             check_prim=True,
             check_pir=(not self.use_mkldnn),
             check_prim_pir=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
 
diff --git a/test/mkldnn/test_activation_bf16_mkldnn_op.py b/test/mkldnn/test_activation_bf16_mkldnn_op.py
@@ -217,6 +217,7 @@ def set_attrs(self):
 class TestMKLDNNSwishBF16Op(MKLDNNBF16ActivationOp, TestActivation):
     def config(self):
         self.op_type = "swish"
+        self.check_pir_onednn = True
 
     def expit(self, val):
         return 1 / (1 + np.exp(-self.beta * val))
diff --git a/test/mkldnn/test_activation_mkldnn_op.py b/test/mkldnn/test_activation_mkldnn_op.py
@@ -233,6 +233,7 @@ def setUp(self):
         super().setUp()
 
         self.attrs["use_mkldnn"] = True
+        self.check_pir_onednn = True
 
     def init_dtype(self):
         self.dtype = np.float32
@@ -244,6 +245,7 @@ def setUp(self):
 
         self.attrs["use_mkldnn"] = True
         self.check_eager = False
+        self.check_pir_onednn = True
 
     def init_dtype(self):
         self.dtype = np.float32
diff --git a/test/mkldnn/test_softmax_bf16_mkldnn_op.py b/test/mkldnn/test_softmax_bf16_mkldnn_op.py
@@ -64,7 +64,7 @@ def setUp(self):
         self.attrs = {'axis': self.axis, 'use_mkldnn': self.use_mkldnn}
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(core.CPUPlace(), check_pir_onednn=True)
 
     def test_check_grad(self):
         pass
@@ -76,26 +76,31 @@ def init_kernel_type(self):
 class TestSoftmaxMKLDNNOp2(TestSoftmaxOp2):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp3(TestSoftmaxOp3):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp4(TestSoftmaxOp4):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp5(TestSoftmaxOp5):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp6(TestSoftmaxOp6):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
 
 if __name__ == '__main__':
diff --git a/test/mkldnn/test_softmax_mkldnn_op.py b/test/mkldnn/test_softmax_mkldnn_op.py
@@ -26,6 +26,7 @@
     TestSoftmaxOp6,
     TestSoftmaxOp_ZeroDim1,
 )
+from utils import compare_legacy_with_pt
 
 import paddle
 from paddle.base import core
@@ -71,9 +72,11 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         if self.use_cudnn:
             place = core.CUDAPlace(0)
-            self.check_output_with_place(place, check_dygraph=False)
+            self.check_output_with_place(
+                place, check_dygraph=False, check_pir_onednn=True
+            )
         else:
-            self.check_output(check_dygraph=False)
+            self.check_output(check_dygraph=False, check_pir_onednn=True)
 
     def test_check_grad(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
@@ -86,10 +89,15 @@ def test_check_grad(self):
                     "Out",
                     max_relative_error=0.01,
                     check_dygraph=False,
+                    check_pir_onednn=True,
                 )
         else:
             self.check_grad(
-                ["X"], "Out", max_relative_error=0.01, check_dygraph=False
+                ["X"],
+                "Out",
+                max_relative_error=0.01,
+                check_dygraph=False,
+                check_pir_onednn=True,
             )
 
     def init_kernel_type(self):
@@ -101,36 +109,42 @@ def init_kernel_type(self):
         self.use_mkldnn = True
         # oneDNN doesn't support float64 dtype
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp3(TestSoftmaxOp3):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp4(TestSoftmaxOp4):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp5(TestSoftmaxOp5):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp6(TestSoftmaxOp6):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestSoftmaxMKLDNNOp_ZeroDim(TestSoftmaxOp_ZeroDim1):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 # Check if primitives already exist in backward
@@ -149,6 +163,7 @@ def setUp(self):
     def __softmax_bwd(self, out, out_grad):
         return out * (out_grad - np.dot(out, out_grad))
 
+    @compare_legacy_with_pt
     def test_check(self):
         check_if_mkldnn_primitives_exist_in_bwd(
             self, self.op_type, self.x, self.out, self.out_grad, self.x_grad
