【PIR API adaptor No.80、81】 Migrate fused_layer_norm and FusedDropoutAdd into pir (#59420)

MarioLulab · web-flow · commit 75a775982087 · 2023-12-07T11:03:48.000+08:00
diff --git a/python/paddle/incubate/nn/functional/fused_layer_norm.py b/python/paddle/incubate/nn/functional/fused_layer_norm.py
@@ -15,7 +15,7 @@
 
 import paddle
 from paddle import _C_ops
-from paddle.framework import LayerHelper, in_dynamic_mode
+from paddle.framework import LayerHelper, in_dynamic_mode, in_pir_mode
 
 
 def fused_layer_norm(
@@ -84,6 +84,22 @@ def fused_layer_norm(
             quant_max_bound,
             quant_min_bound,
         )
+    elif in_pir_mode():
+        out, residual_out, _, _ = _C_ops.fused_bias_residual_layernorm(
+            x,
+            bias,
+            residual,
+            norm_weight,
+            norm_bias,
+            epsilon,
+            residual_alpha,
+            begin_norm_axis,
+            quant_scale,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+        )
+        return (out, residual_out) if residual is not None else out
 
     helper = LayerHelper('fused_layernorm', **locals())
     out = None
diff --git a/test/legacy_test/test_fused_dropout_add_op.py b/test/legacy_test/test_fused_dropout_add_op.py
@@ -17,10 +17,10 @@
 import numpy as np
 
 import paddle
-from paddle import base
 from paddle.base import core
 from paddle.incubate.nn.functional import fused_dropout_add
 from paddle.incubate.nn.layer.fused_dropout_add import FusedDropoutAdd
+from paddle.pir_utils import test_with_pir_api
 
 
 def paddle_dropout_add(x, y, p=0.5, training=True, mode="upscale_in_train"):
@@ -127,6 +127,7 @@ def setUp(self):
         self.shape = (2, 80, 8, 2)
         self.dtype = 'float16'
 
+    @test_with_pir_api
     def test_static_op(self):
         paddle.disable_static()
         paddle.seed(312)
@@ -152,7 +153,7 @@ def test_static_op(self):
 
             outs = fused_dropout_add(xs, ys, p=0.5, training=True)
 
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "xs": x_data.astype('float16'),
diff --git a/test/legacy_test/test_fused_layernorm_op.py b/test/legacy_test/test_fused_layernorm_op.py
@@ -16,8 +16,8 @@
 import numpy as np
 
 import paddle
-from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def quant_helper(
@@ -421,10 +421,10 @@ def check_layernorm(self, x_np, gamma_np, beta_np, dtype):
                 name="x_static", shape=[self.batch, self.cols], dtype=dtype
             )
             gamma_static = paddle.static.data(
-                name="gamma_static", shape=[self.cols], dtype=paddle.float32
+                name="gamma_static", shape=[self.cols], dtype='float32'
             )
             beta_static = paddle.static.data(
-                name="beta_static", shape=[self.cols], dtype=paddle.float32
+                name="beta_static", shape=[self.cols], dtype='float32'
             )
             outs = paddle.incubate.nn.functional.fused_layer_norm(
                 x_static,
@@ -433,7 +433,7 @@ def check_layernorm(self, x_np, gamma_np, beta_np, dtype):
                 self.epsilon,
                 begin_norm_axis=1,
             )
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "x_static": x_np.astype(dtype),
@@ -467,10 +467,10 @@ def check_layernorm_int8(self, x_np, gamma_np, beta_np, dtype):
                 name="x_static", shape=[self.batch, self.cols], dtype=dtype
             )
             gamma_static = paddle.static.data(
-                name="gamma_static", shape=[self.cols], dtype=paddle.float32
+                name="gamma_static", shape=[self.cols], dtype='float32'
             )
             beta_static = paddle.static.data(
-                name="beta_static", shape=[self.cols], dtype=paddle.float32
+                name="beta_static", shape=[self.cols], dtype='float32'
             )
             outs = paddle.incubate.nn.functional.fused_layer_norm(
                 x_static,
@@ -483,7 +483,7 @@ def check_layernorm_int8(self, x_np, gamma_np, beta_np, dtype):
                 quant_max_bound=self.quant_max_bound,
                 quant_min_bound=self.quant_min_bound,
             )
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "x_static": x_np.astype(dtype),
@@ -532,14 +532,16 @@ def check_residual_bias_add(self, x_np, residual_np, bias_np, dtype):
                 quant_min_bound=self.quant_min_bound,
             )
 
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "x_static": x_np.astype(dtype),
                     "residual_static": residual_np.astype(dtype),
                     "bias_static": bias_np.astype(dtype),
                 },
-                fetch_list=[outs],
+                fetch_list=[
+                    outs[0]
+                ],  # NOTE: Only fetch `out`, because `residual_out` will not be initialized if both `norm_weight` and `norm_bias` are None.
             )
         return out_s, paddle_naive_residual_out
 
@@ -574,10 +576,10 @@ def check_residual_bias_layernorm(
                 name="bias_static", shape=[self.cols], dtype=dtype
             )
             gamma_static = paddle.static.data(
-                name="gamma_static", shape=[self.cols], dtype=paddle.float32
+                name="gamma_static", shape=[self.cols], dtype='float32'
             )
             beta_static = paddle.static.data(
-                name="beta_static", shape=[self.cols], dtype=paddle.float32
+                name="beta_static", shape=[self.cols], dtype='float32'
             )
             outs = paddle.incubate.nn.functional.fused_layer_norm(
                 x_static,
@@ -590,7 +592,7 @@ def check_residual_bias_layernorm(
                 residual=residual_static,
             )
 
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "x_static": x_np.astype(dtype),
@@ -644,10 +646,10 @@ def check_residual_bias_layernorm_int8(
                 name="bias_static", shape=[self.cols], dtype=dtype
             )
             gamma_static = paddle.static.data(
-                name="gamma_static", shape=[self.cols], dtype=paddle.float32
+                name="gamma_static", shape=[self.cols], dtype='float32'
             )
             beta_static = paddle.static.data(
-                name="beta_static", shape=[self.cols], dtype=paddle.float32
+                name="beta_static", shape=[self.cols], dtype='float32'
             )
             outs = paddle.incubate.nn.functional.fused_layer_norm(
                 x_static,
@@ -664,7 +666,7 @@ def check_residual_bias_layernorm_int8(
                 quant_min_bound=self.quant_min_bound,
             )
 
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             out_s = exe.run(
                 feed={
                     "x_static": x_np.astype(dtype),
@@ -677,6 +679,7 @@ def check_residual_bias_layernorm_int8(
             )
         return out_s, paddle_naive_layernorm_out, paddle_naive_residual_out
 
+    @test_with_pir_api
     def test_layernorm_fp16(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -691,6 +694,7 @@ def test_layernorm_fp16(self):
             atol=1e-3,
         )
 
+    @test_with_pir_api
     def test_layernorm_int8(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -704,6 +708,7 @@ def test_layernorm_int8(self):
             atol=2,
         )
 
+    @test_with_pir_api
     def test_residual_bias_add(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -724,6 +729,7 @@ def test_residual_bias_add(self):
             atol=1e-3,
         )
 
+    @test_with_pir_api
     def test_residual_bias_add_layernorm_fp16(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -754,6 +760,7 @@ def test_residual_bias_add_layernorm_fp16(self):
             atol=1e-3,
         )
 
+    @test_with_pir_api
     def test_residual_bias_add_layernorm_int8(self):
         if not paddle.is_compiled_with_cuda():
             return
