【PIR API adaptor No.87】Migrate fused_rotary_position_embedding into pir

python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py

@@ -15,7 +15,7 @@
 
 from paddle import _C_ops
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_or_pir_mode
 
 
 def fused_rotary_position_embedding(
@@ -87,7 +87,7 @@ def fused_rotary_position_embedding(
               [[ 0.07116699, -0.90966797],
                [-0.03628540, -0.20202637]]]])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.fused_rotary_position_embedding(
             q, k, v, sin, cos, position_ids, use_neox_rotary_style
         )

python/paddle/incubate/nn/functional/fused_transformer.py

@@ -18,7 +18,11 @@
 from paddle.base.data_feeder import check_dtype, check_variable_and_dtype
 from paddle.base.framework import default_main_program
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import in_dynamic_mode, in_dynamic_or_pir_mode
+from paddle.framework import (
+    in_dynamic_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
 
 __all__ = []
 
@@ -1189,6 +1193,38 @@ def fused_multi_transformer(
             'ring_id',
             ring_id,
         )
+    elif in_pir_mode():
+        cache_kv_out, final_out = _C_ops.fused_multi_transformer(
+            x,
+            ln_scales,
+            ln_biases,
+            qkv_weights,
+            qkv_biases,
+            cache_kvs,
+            pre_caches,
+            rotary_embs,
+            time_step,
+            seq_lens,
+            attn_mask,
+            linear_weights,
+            linear_biases,
+            ffn_ln_scales,
+            ffn_ln_biases,
+            ffn1_weights,
+            ffn1_biases,
+            ffn2_weights,
+            ffn2_biases,
+            cache_kvs,
+            pre_layer_norm,
+            epsilon,
+            dropout_rate,
+            rotary_emb_dims,
+            not training,
+            mode,
+            activation,
+            trans_qkvw,
+            ring_id,
+        )
         if cache_kvs is not None:
             return final_out, cache_kv_out
         return final_out

test/legacy_test/test_fused_multi_transformer_op.py

@@ -28,6 +28,7 @@
 from paddle.nn.layer.common import Dropout, Linear
 from paddle.nn.layer.norm import LayerNorm
 from paddle.nn.layer.transformer import _convert_attention_mask
+from paddle.pir_utils import test_with_pir_api
 
 seed = 42
 
@@ -804,6 +805,7 @@ def GetFusedMultiTransformerOut(self):
 
         return final_out
 
+    @test_with_pir_api
     def GetFusedMultiTransformerOutStatic(self):
         paddle.enable_static()
         x = paddle.static.data('x', self.query.shape, self.query.dtype)
@@ -1395,6 +1397,7 @@ def config(self):
             initializer=paddle.paddle.nn.initializer.Constant(0.0)
         )
 
+    @test_with_pir_api
     def test_fused_multi_transformer_op(self):
         self.has_pre_cache = True
         self.remove_padding = False
@@ -1425,6 +1428,7 @@ def test_fused_multi_transformer_op(self):
 # Starts the name of this test with 'Z' to make this test
 # run after others. If not, it will make other tests fail.
 class ZTestFusedMultiAttentionAPIError(unittest.TestCase):
+    @test_with_pir_api
     def test_errors(self):
         def test_invalid_input_dim():
             array = np.array([1.9], dtype=np.float32)
@@ -1438,6 +1442,7 @@ def test_invalid_input_dim():
 
 
 class ZTestFusedMultiTransformerAPIError(unittest.TestCase):
+    @test_with_pir_api
     def test_errors(self):
         def test_invalid_input_dim():
             array = np.array([], dtype=np.float32)

test/legacy_test/test_fused_rotary_position_embedding.py

@@ -19,6 +19,7 @@
 import paddle
 from paddle.base import core
 from paddle.incubate.nn.functional import fused_rotary_position_embedding
+from paddle.pir_utils import test_with_pir_api
 
 
 def deal_qkv(init_q, init_k, init_v):
@@ -281,6 +282,7 @@ def test_fused_rope_position_ids(self):
                 p_bw[i].numpy(), f_bw[i].numpy(), rtol=1e-05
             )
 
+    @test_with_pir_api
     def test_static(self):
         tensor_q, tensor_k, tensor_v, tensor_sin, tensor_cos = self.get_inputs(
             self.seed, True