pir onednn support relu (#61218)

wanghuancoder · web-flow · commit a8328fd35c3d · 2024-01-29T19:16:27.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -1642,7 +1642,10 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
                 )
 
                 if dialect_name == "onednn_op":
-                    if len(op_info.onednn_extra_args) > 0:
+                    if (
+                        op_info.onednn_extra_args is not None
+                        and len(op_info.onednn_extra_args) > 0
+                    ):
                         args_name = []
                         for arg in op_info.onednn_extra_args:
                             args_name.append(arg["name"])
@@ -1933,7 +1936,14 @@ def OpGenerator(
                 op_name = op['op']
                 item = {}
                 item["is_onednn_only"] = False
-                item["extra_args"] = parse_extra_args(op_name, op['extra_args'])
+                if 'extra_args' in op:
+                    item["extra_args"] = parse_extra_args(
+                        op_name, op['extra_args']
+                    )
+                    item["attrs"] = parse_extra_args(op_name, op['extra_args'])
+                else:
+                    item["extra_args"] = None
+                    item["attrs"] = None
                 if 'data_format_tensors' in op:
                     item["data_format_tensors"] = parse_data_format_tensors(
                         op_name, op['data_format_tensors']
@@ -1944,7 +1954,6 @@ def OpGenerator(
                     item["dynamic_fallback"] = op['dynamic_fallback']
                 else:
                     item["dynamic_fallback"] = False
-                item["attrs"] = parse_extra_args(op_name, op['extra_args'])
                 ops_onednn_extra_map[op_name] = item
         op_yaml_files.insert(0, onednn_yaml_file)
 
@@ -2008,7 +2017,8 @@ def OpGenerator(
                         "data_format_tensors"
                     ]
                     op["dynamic_fallback"] = onednn_item["dynamic_fallback"]
-                    op["attrs"] = op["attrs"] + onednn_item["attrs"]
+                    if onednn_item["attrs"] is not None:
+                        op["attrs"] = op["attrs"] + onednn_item["attrs"]
                 else:
                     continue
             item = OpInfoParser(op, op_compat_item)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -56,6 +56,10 @@
   data_format_tensors : x
   dynamic_fallback : True
 
+- op : relu
+
+- op : relu_grad
+
 # - op : matmul
 #   extra_args : str mkldnn_data_type="float32"
 #   layout_transform :
diff --git a/test/mkldnn/test_activation_bf16_mkldnn_op.py b/test/mkldnn/test_activation_bf16_mkldnn_op.py
@@ -58,7 +58,9 @@ def calculate_grads(self):
         self.dx = self.op_grad(self.out, self.x)
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(
+            core.CPUPlace(), check_pir_onednn=self.check_pir_onednn
+        )
 
     def test_check_grad(self):
         self.calculate_grads()
@@ -68,6 +70,7 @@ def test_check_grad(self):
             "Out",
             user_defined_grads=[self.dx],
             user_defined_grad_outputs=[convert_float_to_uint16(self.out)],
+            check_pir_onednn=self.check_pir_onednn,
         )
 
 
@@ -151,6 +154,7 @@ def init_data(self):
 class TestMKLDNNReluBF16Op(MKLDNNBF16ActivationOp, TestActivation):
     def config(self):
         self.op_type = "relu"
+        self.check_pir_onednn = True
 
     def op_forward(self, x):
         return np.maximum(x, 0)
