[PIR] pir onednn support add,add_grad,add_double_grad,add_triple_grad (#60952)

wanghuancoder · web-flow · commit adc5089f10ac · 2024-01-19T10:39:34.000+08:00
* pir onednn support add,add_grad,add_double_grad,add_triple_grad
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -168,6 +168,8 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 #include "paddle/fluid/primitive/rule/vjp/vjp.h"
 #include "paddle/pir/core/op_base.h"
 
+using namespace paddle::dialect;
+
 {input}
 
 {define_type_id}
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -1,4 +1,16 @@
 
+- op : add
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0
+
+- op : add_grad
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0
+
+- op : add_double_grad
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0
+
+- op : add_triple_grad
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0
+
 - op : conv2d
   extra_args : bool is_test=false
   data_format_tensors : input
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -101,6 +101,8 @@
     {x : X, y : Y}
   outputs :
     {out : Out}
+  attrs :
+    {scale_x : Scale_x, scale_y : Scale_y, scale_out : Scale_out}
   extra :
     attrs : [bool use_mkldnn = false, str mkldnn_data_type = "float32",
              bool use_quantizer = false, float Scale_x = 1.0f, float Scale_y = 1.0f, float Scale_out = 1.0f]
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -2989,6 +2989,7 @@ def check_grad(
         check_cinn=False,
         check_pir=False,
         check_auto_parallel=False,
+        check_pir_onednn=False,
     ):
         if hasattr(self, "use_custom_device") and self.use_custom_device:
             check_dygraph = False
@@ -3014,7 +3015,94 @@ def check_grad(
                 check_cinn=check_cinn,
                 check_pir=check_pir,
                 check_auto_parallel=check_auto_parallel,
+                check_pir_onednn=check_pir_onednn,
+            )
+
+    def check_grad_with_place_for_static(
+        self,
+        user_defined_grads,
+        inputs_to_check,
+        place,
+        output_names,
+        no_grad_set,
+        user_defined_grad_outputs,
+        numeric_place,
+        numeric_grad_delta,
+        in_place,
+        check_cinn,
+        max_relative_error,
+        atol,
+    ):
+        if user_defined_grads is None and self.is_compared_with_fp32():
+            self.enable_cal_ref_output()
+            numeric_grads = self._get_gradient(
+                inputs_to_check,
+                place,
+                output_names,
+                no_grad_set,
+                user_defined_grad_outputs,
             )
+            self.disable_cal_ref_output()
+        else:
+            numeric_grads = user_defined_grads or [
+                get_numeric_gradient(
+                    numeric_place,
+                    self.scope,
+                    self.op,
+                    self.inputs,
+                    input_to_check,
+                    output_names,
+                    delta=numeric_grad_delta,
+                    in_place=in_place,
+                )
+                for input_to_check in inputs_to_check
+            ]
+
+        analytic_grads = self._get_gradient(
+            inputs_to_check,
+            place,
+            output_names,
+            no_grad_set,
+            user_defined_grad_outputs,
+            check_cinn=check_cinn,
+        )
+        # comparison of bf16 results will happen as fp32
+        # loop over list of grads and convert bf16 to fp32
+
+        fp32_analytic_grads = []
+        for grad in analytic_grads:
+            if grad.dtype == np.uint16:
+                grad = convert_uint16_to_float(grad)
+                max_relative_error = (
+                    0.01 if max_relative_error < 0.01 else max_relative_error
+                )
+            fp32_analytic_grads.append(grad)
+        analytic_grads = fp32_analytic_grads
+
+        fp32_numeric_grads = []
+        for grad in numeric_grads:
+            if grad.dtype == np.uint16:
+                grad = convert_uint16_to_float(grad)
+                max_relative_error = (
+                    0.01 if max_relative_error < 0.01 else max_relative_error
+                )
+            fp32_numeric_grads.append(grad)
+        numeric_grads = fp32_numeric_grads
+
+        if self.is_float16_op():
+            max_relative_error = (
+                0.001 if max_relative_error < 0.001 else max_relative_error
+            )
+        self._assert_is_close(
+            numeric_grads,
+            analytic_grads,
+            inputs_to_check,
+            max_relative_error,
+            "Gradient Check On %s" % str(place),
+            atol=atol,
+        )
+
+        return numeric_grads
 
     def check_grad_with_place(
         self,
@@ -3036,6 +3124,7 @@ def check_grad_with_place(
         check_cinn=False,
         check_pir=False,
         check_auto_parallel=False,
+        check_pir_onednn=False,
     ):
         if hasattr(self, "use_custom_device") and self.use_custom_device:
             check_dygraph = False
@@ -3134,6 +3223,7 @@ def check_grad_with_place(
                         else None,
                     )
                     run_subprocess(start_command, runtime_envs, timeout=120)
+
         self.scope = core.Scope()
         op_inputs = self.inputs if hasattr(self, "inputs") else {}
         op_outputs = self.outputs if hasattr(self, "outputs") else {}
@@ -3209,74 +3299,37 @@ def check_grad_with_place(
         if numeric_place is None:
             numeric_place = place
 
-        if user_defined_grads is None and self.is_compared_with_fp32():
-            self.enable_cal_ref_output()
-            numeric_grads = self._get_gradient(
-                inputs_to_check,
-                place,
-                output_names,
-                no_grad_set,
-                user_defined_grad_outputs,
-            )
-            self.disable_cal_ref_output()
-        else:
-            numeric_grads = user_defined_grads or [
-                get_numeric_gradient(
-                    numeric_place,
-                    self.scope,
-                    self.op,
-                    self.inputs,
-                    input_to_check,
-                    output_names,
-                    delta=numeric_grad_delta,
-                    in_place=in_place,
-                )
-                for input_to_check in inputs_to_check
-            ]
-
-        analytic_grads = self._get_gradient(
+        numeric_grads = self.check_grad_with_place_for_static(
+            user_defined_grads,
             inputs_to_check,
             place,
             output_names,
             no_grad_set,
             user_defined_grad_outputs,
-            check_cinn=check_cinn,
+            numeric_place,
+            numeric_grad_delta,
+            in_place,
+            check_cinn,
+            max_relative_error,
+            atol,
         )
-        # comparison of bf16 results will happen as fp32
-        # loop over list of grads and convert bf16 to fp32
 
-        fp32_analytic_grads = []
-        for grad in analytic_grads:
-            if grad.dtype == np.uint16:
-                grad = convert_uint16_to_float(grad)
-                max_relative_error = (
-                    0.01 if max_relative_error < 0.01 else max_relative_error
-                )
-            fp32_analytic_grads.append(grad)
-        analytic_grads = fp32_analytic_grads
-
-        fp32_numeric_grads = []
-        for grad in numeric_grads:
-            if grad.dtype == np.uint16:
-                grad = convert_uint16_to_float(grad)
-                max_relative_error = (
-                    0.01 if max_relative_error < 0.01 else max_relative_error
+        if check_pir_onednn:
+            with pir_executor_guard():
+                self.check_grad_with_place_for_static(
+                    user_defined_grads,
+                    inputs_to_check,
+                    place,
+                    output_names,
+                    no_grad_set,
+                    user_defined_grad_outputs,
+                    numeric_place,
+                    numeric_grad_delta,
+                    in_place,
+                    check_cinn,
+                    max_relative_error,
+                    atol,
                 )
-            fp32_numeric_grads.append(grad)
-        numeric_grads = fp32_numeric_grads
-
-        if self.is_float16_op():
-            max_relative_error = (
-                0.001 if max_relative_error < 0.001 else max_relative_error
-            )
-        self._assert_is_close(
-            numeric_grads,
-            analytic_grads,
-            inputs_to_check,
-            max_relative_error,
-            "Gradient Check On %s" % str(place),
-            atol=atol,
-        )
 
         if check_dygraph:
             with base.dygraph.base.guard(place):
diff --git a/test/legacy_test/test_elementwise_add_op.py b/test/legacy_test/test_elementwise_add_op.py
@@ -30,6 +30,7 @@
 class TestElementwiseAddOp(OpTest):
     def init_kernel_type(self):
         self.use_mkldnn = False
+        self.check_pir_onednn = False
 
     def setUp(self):
         self.op_type = "elementwise_add"
@@ -58,6 +59,7 @@ def test_check_output(self):
         self.check_output(
             check_dygraph=self.check_dygraph(),
             check_pir=self.check_dygraph(),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_normal(self):
@@ -71,6 +73,7 @@ def test_check_grad_normal(self):
             check_prim=self.check_prim,
             check_prim_pir=self.check_dygraph(),
             check_pir=self.check_dygraph(),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_ingore_x(self):
@@ -85,6 +88,7 @@ def test_check_grad_ingore_x(self):
             check_prim=self.check_prim,
             check_prim_pir=self.check_dygraph(),
             check_pir=self.check_dygraph(),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_ingore_y(self):
@@ -99,6 +103,7 @@ def test_check_grad_ingore_y(self):
             check_prim=self.check_prim,
             check_prim_pir=self.check_dygraph(),
             check_pir=self.check_dygraph(),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def init_input_output(self):
diff --git a/test/mkldnn/test_elementwise_add_bf16_mkldnn_op.py b/test/mkldnn/test_elementwise_add_bf16_mkldnn_op.py
@@ -49,7 +49,7 @@ def generate_data(self):
         self.out = np.add(self.x, self.y)
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(core.CPUPlace(), check_pir_onednn=True)
 
     # elementwise_add grad (no braodcasting) is just passing upper gradients to either X or Y or both
     def test_check_grad_normal(self):
@@ -60,6 +60,7 @@ def test_check_grad_normal(self):
             check_dygraph=False,
             user_defined_grads=[self.x, self.x],
             user_defined_grad_outputs=[self.x_bf16],
+            check_pir_onednn=True,
         )
 
     def test_check_grad_ingore_x(self):
@@ -70,6 +71,7 @@ def test_check_grad_ingore_x(self):
             check_dygraph=False,
             user_defined_grads=[self.y],
             user_defined_grad_outputs=[self.y_bf16],
+            check_pir_onednn=True,
         )
 
     def test_check_grad_ingore_y(self):
@@ -80,6 +82,7 @@ def test_check_grad_ingore_y(self):
             check_dygraph=False,
             user_defined_grads=[self.x],
             user_defined_grad_outputs=[self.x_bf16],
+            check_pir_onednn=True,
         )
 
 
@@ -106,6 +109,7 @@ def test_check_grad_normal(self):
             check_dygraph=False,
             user_defined_grads=[self.x, self.compute_reduced_gradients(self.x)],
             user_defined_grad_outputs=[self.x_bf16],
+            check_pir_onednn=True,
         )
 
     def test_check_grad_ingore_x(self):
@@ -116,6 +120,7 @@ def test_check_grad_ingore_x(self):
             check_dygraph=False,
             user_defined_grads=[self.compute_reduced_gradients(self.x)],
             user_defined_grad_outputs=[self.x_bf16],
+            check_pir_onednn=True,
         )
 
 
diff --git a/test/mkldnn/test_elementwise_add_mkldnn_op.py b/test/mkldnn/test_elementwise_add_mkldnn_op.py
@@ -24,6 +24,7 @@
 class TestOneDNNElementwiseAddOp(TestElementwiseAddOp):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
     def init_dtype(self):
         self.dtype = np.float32
@@ -132,6 +133,7 @@ class TestInt8(TestElementwiseAddOp):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self._cpu_only = True
+        self.check_pir_onednn = True
 
     def init_dtype(self):
         self.dtype = np.int8
@@ -149,7 +151,10 @@ def init_scales(self):
     def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         self.init_scales()
-        self.check_output(check_dygraph=(not self.use_mkldnn))
+        self.check_output(
+            check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
+        )
 
     def test_check_grad_normal(self):
         pass
diff --git a/test/mkldnn/test_elementwise_add_onednn_op_rare_shape.py b/test/mkldnn/test_elementwise_add_onednn_op_rare_shape.py
@@ -31,6 +31,7 @@
 class TestOneDNNlementwiseAddSrcDifferentShape(TestElementwiseAddOp):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
     def init_dtype(self):
         self.dtype = np.float32
