chore: import upstream snapshot with attribution

test/prim/CMakeLists.txt

@@ -0,0 +1,16 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()
+
+add_subdirectory(prim)
+add_subdirectory(model)
+add_subdirectory(composite_ops)
+add_subdirectory(process)
+add_subdirectory(pir_prim)
+add_subdirectory(high_order)

test/prim/README.txt

@@ -0,0 +1,34 @@
+# Prim
+
+Test composite and primitive API and related autodiff rules.
+
+
+## Package Structure
+
+- comp(abbr for composite):
+
+  Test composite API which is composed of primitive api, but not include the autodiff rules for primitive api.
+
+- prim(abbr for primitive):
+
+  Test primitive api and autodiff rules.
+
+    - vjp: Test vjp rules.
+
+        - eager: Test vjp rules in eager mode.
+        - static: Test vjp rules in static mode.
+
+    - jvp(TODO): Test jvp rules.
+
+
+## How to
+
+- Forward API and First-order Autodiff
+
+  Compare numerical value with raw phi operators using tools such as `np.testing.assert_allclose`
+
+
+- Higher-order Autodiff
+
+  Compare numerical value with the result computed by finite difference. Tool
+  used for computing Higher-order finite difference will be provided in the future.

test/prim/composite_ops/CMakeLists.txt

@@ -0,0 +1,10 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()

test/prim/composite_ops/utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# default tolerance
+TOLERANCE = {
+    "float16": {
+        "forward": {"rtol": 1e-3, "atol": 1e-3},
+        "backward": {"rtol": 1e-3, "atol": 1e-3},
+        "prim_backward": {"rtol": 1e-3, "atol": 1e-3},
+    },
+    "float32": {
+        "forward": {"rtol": 1e-6, "atol": 1e-6},
+        "backward": {"rtol": 1e-6, "atol": 1e-6},
+        "prim_backward": {"rtol": 1e-6, "atol": 1e-6},
+    },
+    "float64": {
+        "forward": {"rtol": 1e-15, "atol": 1e-15},
+        "backward": {"rtol": 1e-15, "atol": 1e-15},
+        "prim_backward": {"rtol": 1e-15, "atol": 1e-15},
+    },
+}
+
+# this tolerance is for big composite ops like batch_norm.
+SUB_TOLERANCE = {
+    "float16": {
+        "forward": {"rtol": 1e-2, "atol": 1e-2},
+        "backward": {"rtol": 1e-2, "atol": 1e-2},
+        "prim_backward": {"rtol": 1e-2, "atol": 1e-2},
+    },
+    "float32": {
+        "forward": {"rtol": 1e-5, "atol": 1e-5},
+        "backward": {"rtol": 1e-5, "atol": 1e-5},
+        "prim_backward": {"rtol": 1e-5, "atol": 1e-5},
+    },
+    "float64": {
+        "forward": {"rtol": 1e-13, "atol": 1e-13},
+        "backward": {"rtol": 1e-13, "atol": 1e-13},
+        "prim_backward": {"rtol": 1e-13, "atol": 1e-13},
+    },
+}

test/prim/high_order/CMakeLists.txt

@@ -0,0 +1,12 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS FLAGS_enable_pir_api=true
+                  FLAGS_prim_enable_dynamic=true)
+endforeach()
+
+set_tests_properties(test_high_order_derivative PROPERTIES TIMEOUT 90)

test/prim/high_order/test_high_order_derivative.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+
+import numpy as np
+
+sys.path.append("../../legacy_test")
+
+import autodiff_checker_helper as ad_checker
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+class TestMulHigherOrderAD(unittest.TestCase):
+    @param.parameterized.expand(
+        [
+            (
+                paddle.multiply,
+                [2, 3, 2],
+                [2, 3, 2],
+                'float32',
+                2,
+                ('cpu', 'gpu'),
+            ),
+            (
+                paddle.multiply,
+                [2, 3, 2],
+                [2, 3, 2],
+                'float64',
+                2,
+                ('cpu', 'gpu'),
+            ),
+            (
+                paddle.multiply,
+                [2, 3, 2, 1],
+                [2, 3, 2, 4],
+                'float64',
+                2,
+                ('cpu', 'gpu'),
+            ),
+        ]
+    )
+    def test_high_order_autodiff(
+        self, func, shape1, shape2, dtype, order, places
+    ):
+        var_1 = np.random.randn(*shape1).astype(dtype)
+        var_2 = np.random.randn(*shape2).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            var2 = paddle.to_tensor(var_2, place=place)
+            var1, var2 = paddle.broadcast_tensors(input=[var1, var2])
+            ad_checker.check_vjp(
+                func, [var1, var2], argnums=(0, 1), order=order
+            )
+
+
+class TestSinHigherOrderAD(unittest.TestCase):
+    @param.parameterized.expand(
+        [
+            (paddle.sin, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.sin, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+class TestCosHigherOrderAD(TestSinHigherOrderAD):
+    @param.parameterized.expand(
+        [
+            (paddle.cos, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.cos, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+class TestTanhHigherOrderAD(TestSinHigherOrderAD):
+    @param.parameterized.expand(
+        [
+            (paddle.tanh, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.tanh, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/model/CMakeLists.txt

@@ -0,0 +1,25 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+list(REMOVE_ITEM TEST_OP test_resnet_cinn test_resnet_prim_cinn
+     test_prim_simplenet_cinn)
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()
+
+set_tests_properties(test_resnet_prim PROPERTIES TIMEOUT 850)
+set_tests_properties(test_bert_prim PROPERTIES TIMEOUT 500)
+
+if(WITH_CINN)
+  set_tests_properties(test_bert_cinn PROPERTIES TIMEOUT 500)
+  set_tests_properties(test_bert_prim_cinn PROPERTIES TIMEOUT 500)
+
+  set_tests_properties(test_resnet_prim PROPERTIES LABELS "RUN_TYPE=CINN")
+  set_tests_properties(test_bert_prim PROPERTIES LABELS "RUN_TYPE=CINN")
+  set_tests_properties(test_bert_cinn PROPERTIES LABELS "RUN_TYPE=CINN")
+  set_tests_properties(test_bert_prim_cinn PROPERTIES LABELS "RUN_TYPE=CINN")
+endif()

test/prim/model/bert.py

@@ -0,0 +1,919 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import functools
+import warnings
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle import Tensor, nn
+from paddle.base.data_feeder import convert_dtype
+from paddle.distributed.fleet.utils import recompute
+from paddle.io import DataLoader, Dataset
+from paddle.nn import MultiHeadAttention
+
+try:
+    from paddle.incubate.nn import FusedTransformerEncoderLayer
+except ImportError:
+    FusedTransformerEncoderLayer = None
+
+
+VOCAB_SIZE = 30522
+
+
+class Stack:
+    def __init__(self, axis=0, dtype=None):
+        self._axis = axis
+        self._dtype = dtype
+
+    def __call__(self, data):
+        data = (
+            np.stack(data, axis=self._axis).astype(self._dtype)
+            if self._dtype
+            else np.stack(data, axis=self._axis)
+        )
+        return data
+
+
+def is_tensor(x):
+    if isinstance(x, paddle.Tensor):
+        return True
+
+    return isinstance(x, np.ndarray)
+
+
+class BertConfig:
+    def __init__(self):
+        self.attention_probs_dropout_prob = 0.1
+        self.fuse = False
+        self.hidden_act = 'gelu'
+        self.hidden_dropout_prob = 0.1
+        # Decrease config to speed up unittest
+        # self.hidden_size = 768
+        self.hidden_size = 60
+        self.initializer_range = 0.02
+        self.intermediate_size = 3072
+        self.layer_norm_eps = 1e-12
+        self.max_position_embeddings = 512
+        self.model_type = 'bert'
+        # self.num_attention_heads = 12
+        self.num_attention_heads = 6
+        # self.num_hidden_layers = 12
+        self.num_hidden_layers = 6
+        self.pad_token_id = 0
+        self.paddlenlp_version = None
+        self.pool_act = 'tanh'
+        self.type_vocab_size = 2
+        self.vocab_size = VOCAB_SIZE
+        self.use_return_dict = False
+        self.output_hidden_states = False
+        self.output_attentions = False
+        self.use_cache = False
+
+
+class BertLMPredictionHead(nn.Layer):
+    def __init__(self, config: BertConfig, embedding_weights=None):
+        super().__init__()
+
+        self.transform = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = getattr(nn.functional, config.hidden_act)
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+        self.decoder_weight = (
+            self.create_parameter(
+                shape=[config.vocab_size, config.hidden_size],
+                dtype=self.transform.weight.dtype,
+                is_bias=False,
+            )
+            if embedding_weights is None
+            else embedding_weights
+        )
+
+        self.decoder_bias = self.create_parameter(
+            shape=[config.vocab_size],
+            dtype=self.decoder_weight.dtype,
+            is_bias=True,
+        )
+
+    def forward(self, hidden_states, masked_positions=None):
+        if masked_positions is not None:
+            hidden_states = paddle.reshape(
+                hidden_states, [-1, hidden_states.shape[-1]]
+            )
+            hidden_states = paddle.tensor.gather(
+                hidden_states, masked_positions
+            )
+        # gather masked tokens might be more quick
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = (
+            paddle.tensor.matmul(
+                hidden_states, self.decoder_weight, transpose_y=True
+            )
+            + self.decoder_bias
+        )
+        return hidden_states
+
+
+class BertPretrainingHeads(nn.Layer):
+    def __init__(self, config: BertConfig, embedding_weights=None):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config, embedding_weights)
+        self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, sequence_output, pooled_output, masked_positions=None):
+        prediction_scores = self.predictions(sequence_output, masked_positions)
+        seq_relationship_score = self.seq_relationship(pooled_output)
+        return prediction_scores, seq_relationship_score
+
+
+class BertEmbeddings(nn.Layer):
+    def __init__(self, config: BertConfig):
+        super().__init__()
+
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size
+        )
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        token_type_ids: Tensor | None = None,
+        position_ids: Tensor | None = None,
+        past_key_values_length: int | None = None,
+    ):
+        if position_ids is None:
+            ones = paddle.ones_like(input_ids, dtype="int64")
+            seq_length = paddle.cumsum(ones, axis=-1)
+
+            position_ids = seq_length - ones
+            if past_key_values_length is not None:
+                position_ids += past_key_values_length
+            position_ids.stop_gradient = True
+        if token_type_ids is None:
+            token_type_ids = paddle.zeros_like(input_ids, dtype="int64")
+
+        input_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = (
+            input_embeddings + position_embeddings + token_type_embeddings
+        )
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertPooler(nn.Layer):
+    def __init__(self, config: BertConfig):
+        super().__init__()
+
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+        self.pool_act = config.pool_act
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        if self.pool_act == "tanh":
+            pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertModel(nn.Layer):
+    def __init__(self, config: BertConfig, to_static, enable_cinn):
+        super().__init__()
+        self.config = config
+        self.pad_token_id = config.pad_token_id
+        self.initializer_range = config.initializer_range
+        self.embeddings = BertEmbeddings(config)
+        if config.fuse and FusedTransformerEncoderLayer is None:
+            warnings.warn(
+                "FusedTransformerEncoderLayer is not supported by the running Paddle. "
+                "The flag fuse_transformer will be ignored. Try Paddle >= 2.3.0"
+            )
+        self.fuse = config.fuse and FusedTransformerEncoderLayer is not None
+        if self.fuse:
+            self.encoder = nn.LayerList(
+                [
+                    FusedTransformerEncoderLayer(
+                        config.hidden_size,
+                        config.num_attention_heads,
+                        config.intermediate_size,
+                        dropout_rate=config.hidden_dropout_prob,
+                        activation=config.hidden_act,
+                        attn_dropout_rate=config.attention_probs_dropout_prob,
+                        act_dropout_rate=0.0,
+                    )
+                    for _ in range(config.num_hidden_layers)
+                ]
+            )
+        else:
+            encoder_layer = nn.TransformerEncoderLayer(
+                config.hidden_size,
+                config.num_attention_heads,
+                config.intermediate_size,
+                dropout=config.hidden_dropout_prob,
+                activation=config.hidden_act,
+                attn_dropout=config.attention_probs_dropout_prob,
+                act_dropout=0,
+            )
+            self.encoder = nn.TransformerEncoder(
+                encoder_layer, config.num_hidden_layers
+            )
+            if to_static:
+                backend = "CINN" if enable_cinn else None
+                self.encoder = paddle.jit.to_static(
+                    self.encoder, None, backend=backend, full_graph=True
+                )
+        self.pooler = BertPooler(config)
+        # self.apply(self.init_weights)
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        token_type_ids: Tensor | None = None,
+        position_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        past_key_values: tuple[tuple[Tensor]] | None = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ):
+        return_dict = (
+            return_dict
+            if return_dict is not None
+            else self.config.use_return_dict
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        use_cache = (
+            use_cache if use_cache is not None else self.config.use_cache
+        )
+
+        past_key_values_length = None
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+        if attention_mask is None:
+            attention_mask = paddle.unsqueeze(
+                (input_ids == self.pad_token_id).astype(
+                    self.pooler.dense.weight.dtype
+                )
+                * -1e4,
+                axis=[1, 2],
+            )
+            if past_key_values is not None:
+                batch_size = past_key_values[0][0].shape[0]
+                past_mask = paddle.zeros(
+                    [batch_size, 1, 1, past_key_values_length],
+                    dtype=attention_mask.dtype,
+                )
+                attention_mask = paddle.concat(
+                    [past_mask, attention_mask], axis=-1
+                )
+
+        else:
+            if attention_mask.ndim == 2:
+                # attention_mask [batch_size, sequence_length] -> [batch_size, 1, 1, sequence_length]
+                attention_mask = attention_mask.unsqueeze(axis=[1, 2]).astype(
+                    paddle.get_default_dtype()
+                )
+                attention_mask = (1.0 - attention_mask) * -1e4
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            past_key_values_length=past_key_values_length,
+        )
+        if self.fuse:
+            assert not output_attentions, (
+                "Not support attentions output currently."
+            )
+            assert past_key_values is None, (
+                "Not support past_key_values currently."
+            )
+            hidden_states = embedding_output
+            all_hidden_states = [] if output_hidden_states else None
+            for layer in self.encoder:
+                hidden_states = layer(hidden_states, attention_mask)
+                if output_hidden_states:
+                    all_hidden_states.append(hidden_states)
+            pooled_output = self.pooler(hidden_states)
+
+            return (
+                (hidden_states, pooled_output, all_hidden_states)
+                if output_hidden_states
+                else (hidden_states, pooled_output)
+            )
+        else:
+            self.encoder._use_cache = use_cache  # To be consistent with HF
+            encoder_outputs = self.encoder(
+                embedding_output,
+                src_mask=attention_mask,
+                cache=past_key_values,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+            if isinstance(encoder_outputs, type(embedding_output)):
+                sequence_output = encoder_outputs
+                pooled_output = self.pooler(sequence_output)
+                return (sequence_output, pooled_output)
+            else:
+                sequence_output = encoder_outputs[0]
+                pooled_output = self.pooler(sequence_output)
+                return (sequence_output, pooled_output, *encoder_outputs[1:])
+
+
+class Bert(nn.Layer):
+    def __init__(self, to_static, enable_cinn):
+        super().__init__()
+        config = BertConfig()
+        self.bert = BertModel(config, to_static, enable_cinn)
+        self.cls = BertPretrainingHeads(
+            config,
+            embedding_weights=self.bert.embeddings.word_embeddings.weight,
+        )
+
+        # self.apply(self.init_weights)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        token_type_ids: Tensor | None = None,
+        position_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        masked_positions: Tensor | None = None,
+        labels: Tensor | None = None,
+        next_sentence_label: Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ):
+        with paddle.static.amp.fp16_guard():
+            outputs = self.bert(
+                input_ids,
+                token_type_ids=token_type_ids,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+            sequence_output, pooled_output = outputs[:2]
+            prediction_scores, seq_relationship_score = self.cls(
+                sequence_output, pooled_output, masked_positions
+            )
+
+            total_loss = None
+            if labels is not None and next_sentence_label is not None:
+                loss_fct = paddle.nn.CrossEntropyLoss()
+                masked_lm_loss = loss_fct(
+                    prediction_scores.reshape(
+                        (-1, prediction_scores.shape[-1])
+                    ),
+                    labels.reshape((-1,)),
+                )
+                next_sentence_loss = loss_fct(
+                    seq_relationship_score.reshape((-1, 2)),
+                    next_sentence_label.reshape((-1,)),
+                )
+                total_loss = masked_lm_loss + next_sentence_loss
+
+            output = (prediction_scores, seq_relationship_score, *outputs[2:])
+            return ((total_loss, *output)) if total_loss is not None else output
+
+
+class BertPretrainingCriterion(paddle.nn.Layer):
+    def __init__(self, vocab_size=VOCAB_SIZE):
+        super().__init__()
+        # CrossEntropyLoss is expensive since the inner reshape (copy)
+        self.loss_fn = paddle.nn.loss.CrossEntropyLoss(ignore_index=-1)
+        self.vocab_size = vocab_size
+
+    def forward(
+        self,
+        prediction_scores,
+        seq_relationship_score,
+        masked_lm_labels,
+        next_sentence_labels,
+        masked_lm_scale,
+    ):
+        with paddle.static.amp.fp16_guard():
+            masked_lm_loss = F.cross_entropy(
+                prediction_scores,
+                masked_lm_labels,
+                reduction="none",
+                ignore_index=-1,
+            )
+            masked_lm_loss = masked_lm_loss / masked_lm_scale
+            next_sentence_loss = F.cross_entropy(
+                seq_relationship_score, next_sentence_labels, reduction="none"
+            )
+        return paddle.sum(masked_lm_loss) + paddle.mean(next_sentence_loss)
+
+
+def layer_init_wrapper(func):
+    @functools.wraps(func)
+    def _impl(self, *args, **kwargs):
+        enable_recompute = kwargs.pop("enable_recompute", False)
+        func(self, *args, **kwargs)
+        if paddle.in_dynamic_mode():
+            self.enable_recompute = enable_recompute
+        else:
+            self.enable_recompute = False
+
+    return _impl
+
+
+def _convert_attention_mask(attn_mask, dtype):
+    if attn_mask is not None and attn_mask.dtype != dtype:
+        attn_mask_dtype = convert_dtype(attn_mask.dtype)
+        if attn_mask_dtype == 'bool' or 'int' in attn_mask_dtype:
+            attn_mask = (paddle.cast(attn_mask, dtype) - 1.0) * 1e9
+        else:
+            attn_mask = paddle.cast(attn_mask, dtype)
+    return attn_mask
+
+
+def _transformer_encoder_layer_fwd(
+    self, src, src_mask=None, cache=None, output_attentions=False
+):
+    self.self_attn.need_weights = output_attentions
+    src_mask = _convert_attention_mask(src_mask, src.dtype)
+
+    residual = src
+    if self.normalize_before:
+        src = self.norm1(src)
+
+    attn_outputs = self.self_attn(src, src, src, src_mask, cache)
+    if isinstance(attn_outputs, tuple):
+        src = attn_outputs[0]
+        outputs = attn_outputs[1:]
+    else:
+        src = attn_outputs
+        outputs = None
+
+    src = residual + self.dropout1(src)
+    if not self.normalize_before:
+        src = self.norm1(src)
+
+    residual = src
+    if self.normalize_before:
+        src = self.norm2(src)
+    src = self.linear2(self.dropout(self.activation(self.linear1(src))))
+    src = residual + self.dropout2(src)
+    if not self.normalize_before:
+        src = self.norm2(src)
+
+    return (
+        src if outputs is None else ((src, *outputs[::-1]))
+    )  # hidden_states, cache, attentions
+
+
+def _transformer_decoder_layer_fwd(
+    self,
+    tgt,
+    memory,
+    tgt_mask=None,
+    memory_mask=None,
+    cache=None,
+    output_attentions=False,
+):
+    residual = tgt
+
+    # self attention
+    self.self_attn.need_weights = output_attentions
+    tgt_mask = _convert_attention_mask(tgt_mask, tgt.dtype)
+
+    if self.normalize_before:
+        tgt = self.norm1(tgt)
+
+    self_attn_outputs = self.self_attn(
+        tgt, tgt, tgt, tgt_mask, cache[0] if cache else None
+    )
+    # self_attn_outputs = (tgt, attn_weights, incremental_cache) or only tgt
+    if isinstance(self_attn_outputs, type(tgt)):
+        tgt = self_attn_outputs
+    else:
+        tgt = self_attn_outputs[0]
+        if output_attentions:
+            self_attn_weights = self_attn_outputs[1]
+        if cache:
+            incremental_cache = self_attn_outputs[-1]
+
+    tgt = residual + self.dropout1(tgt)
+    if not self.normalize_before:
+        tgt = self.norm1(tgt)
+
+    residual = tgt
+
+    # cross attention
+    if memory is not None:
+        self.cross_attn.need_weights = output_attentions
+        memory_mask = _convert_attention_mask(memory_mask, memory.dtype)
+
+        if self.normalize_before:
+            tgt = self.norm2(tgt)
+
+        cross_attn_outputs = self.cross_attn(
+            tgt, memory, memory, memory_mask, cache[1] if cache else None
+        )
+        if isinstance(cross_attn_outputs, type(tgt)):
+            tgt = cross_attn_outputs
+        else:
+            tgt = cross_attn_outputs[0]
+            if output_attentions:
+                cross_attn_weights = cross_attn_outputs[1]
+            if cache:
+                static_cache = cross_attn_outputs[-1]
+
+        tgt = residual + self.dropout2(tgt)
+        if not self.normalize_before:
+            tgt = self.norm2(tgt)
+
+        residual = tgt
+
+    if self.normalize_before:
+        tgt = self.norm3(tgt)
+    tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+    tgt = residual + self.dropout3(tgt)
+    if not self.normalize_before:
+        tgt = self.norm3(tgt)
+
+    if not output_attentions and cache is None:
+        return tgt
+    else:
+        outputs = (tgt,)
+        if output_attentions:
+            outputs += (
+                self_attn_weights,
+                cross_attn_weights if memory is not None else None,
+            )
+        if cache:
+            outputs += (
+                (
+                    incremental_cache,
+                    static_cache if memory is not None else None,
+                ),
+            )
+        return outputs
+
+
+def _transformer_encoder_fwd(
+    self,
+    src,
+    src_mask=None,
+    cache=None,
+    output_attentions=False,
+    output_hidden_states=False,
+    return_dict=False,
+):
+    src_mask = _convert_attention_mask(src_mask, src.dtype)
+
+    output = src
+    # To get cache from None when use_cache is True, which is compatible with HF
+    # while HF requires decoder. The implementation here uses cache update in the
+    # MultiHeadAttention not so efficiently, and maybe optimize it later.
+    if cache is None and getattr(self, "_use_cache", False):
+        cache = [tuple(self.layers[0].gen_cache(src))] * len(self.layers)
+    # To be compatible with `TransformerEncoder.forward`, `_use_cache` defaults
+    # to True when cache is not None.
+    new_caches = (
+        [] if cache is not None and getattr(self, "_use_cache", True) else None
+    )
+    all_attentions = [] if output_attentions else None
+    # NOTE: Also includes embedding output which is same as HF.
+    all_hidden_states = [output] if output_hidden_states else None
+    for i, mod in enumerate(self.layers):
+        if self.enable_recompute:
+            # Note: recompute do not support pass as **kwargs yet.
+            layer_outputs = recompute(
+                mod,
+                output,
+                src_mask,
+                (
+                    None
+                    if cache is None
+                    else (
+                        cache[i]
+                        if isinstance(cache[i], MultiHeadAttention.Cache)
+                        else MultiHeadAttention.Cache(*cache[i])
+                    )
+                ),
+                output_attentions,
+            )
+        else:
+            layer_outputs = mod(
+                output,
+                src_mask=src_mask,
+                cache=(
+                    None
+                    if cache is None
+                    else (
+                        cache[i]
+                        if isinstance(cache[i], MultiHeadAttention.Cache)
+                        else MultiHeadAttention.Cache(*cache[i])
+                    )
+                ),
+                output_attentions=output_attentions,
+            )
+
+        if isinstance(layer_outputs, tuple):
+            output = layer_outputs[0]
+            outputs = layer_outputs[1:]
+        else:
+            output = layer_outputs
+            outputs = None
+
+        if output_hidden_states:
+            all_hidden_states.append(output)
+        if output_attentions:
+            all_attentions.append(outputs[-1])
+        if new_caches is not None:
+            new_caches.append(
+                outputs[0]
+                if isinstance(cache[i], MultiHeadAttention.Cache)
+                else (tuple(outputs[0]))
+            )
+
+    if self.norm is not None:
+        output = self.norm(output)
+
+        if output_hidden_states:
+            all_hidden_states[-1] = output
+
+    outputs = tuple(
+        tuple(v) if isinstance(v, list) else v
+        for v in [
+            output,
+            new_caches,
+            all_hidden_states,
+            all_attentions,
+        ]
+        if v is not None
+    )
+    if len(outputs) == 1:
+        return output
+    else:
+        return outputs
+
+
+def _transformer_decoder_fwd(
+    self,
+    tgt,
+    memory=None,
+    tgt_mask=None,
+    memory_mask=None,
+    cache=None,
+    output_attentions=False,
+    output_hidden_states=False,
+    return_dict=False,
+):
+    tgt_mask = _convert_attention_mask(tgt_mask, tgt.dtype)
+    if memory is not None:
+        memory_mask = _convert_attention_mask(memory_mask, memory.dtype)
+
+    new_caches = [] if cache else None
+    all_hidden_states = [tgt] if output_hidden_states else None
+    all_self_attns = [] if output_attentions else None
+    all_cross_attns = [] if output_attentions else None
+
+    for i, mod in enumerate(self.layers):
+        if cache is None:
+            if self.enable_recompute:
+                outputs = recompute(
+                    mod,
+                    tgt,
+                    memory,
+                    tgt_mask,
+                    memory_mask,
+                    None,
+                    output_attentions,
+                )
+            else:
+                outputs = mod(
+                    tgt,
+                    memory,
+                    tgt_mask=tgt_mask,
+                    memory_mask=memory_mask,
+                    cache=None,
+                    output_attentions=output_attentions,
+                )
+        else:
+            outputs = mod(
+                tgt,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                cache=cache[i] if cache else None,
+                output_attentions=output_attentions,
+            )
+        if isinstance(outputs, type(tgt)):
+            tgt = outputs
+        else:
+            tgt = outputs[0]
+        if cache:
+            new_caches.append(outputs[-1])
+        if output_attentions:
+            all_self_attns.append(outputs[1])
+            all_cross_attns.append(outputs[2])
+        if output_hidden_states:
+            all_hidden_states.append(tgt)
+
+    if self.norm is not None:
+        tgt = self.norm(tgt)
+        if output_hidden_states:
+            all_hidden_states[-1] = tgt
+
+    if isinstance(outputs, type(tgt)):
+        return tgt
+
+    temp_list = [
+        tgt,
+        new_caches if cache else None,
+        all_hidden_states,
+        all_self_attns,
+        all_cross_attns,
+    ]
+    return tuple(v for v in temp_list if v is not None)
+
+
+# patches of paddle.nn.Transformer to get all hidden_states and attentions
+paddle.nn.TransformerEncoderLayer.forward = _transformer_encoder_layer_fwd
+paddle.nn.TransformerDecoderLayer.forward = _transformer_decoder_layer_fwd
+paddle.nn.TransformerEncoder.forward = _transformer_encoder_fwd
+paddle.nn.TransformerDecoder.forward = _transformer_decoder_fwd
+
+_encoder_init = paddle.nn.TransformerEncoder.__init__
+_decoder_init = paddle.nn.TransformerDecoder.__init__
+paddle.nn.TransformerEncoder.__init__ = layer_init_wrapper(_encoder_init)
+paddle.nn.TransformerDecoder.__init__ = layer_init_wrapper(_decoder_init)
+
+
+class PretrainingDataset(Dataset):
+    def __init__(self, input_file, max_pred_length):
+        self.input_file = input_file
+        self.max_pred_length = max_pred_length
+        keys = [
+            "input_ids",
+            "input_mask",
+            "segment_ids",
+            "masked_lm_positions",
+            "masked_lm_ids",
+            "next_sentence_labels",
+        ]
+        self.inputs = np.load(input_file)
+        self.inputs = [self.inputs[key] for key in keys]
+
+    def __len__(self):
+        "Denotes the total number of samples"
+        return len(self.inputs[0])
+
+    def __getitem__(self, index):
+        [
+            input_ids,
+            input_mask,
+            segment_ids,
+            masked_lm_positions,
+            masked_lm_ids,
+            next_sentence_labels,
+        ] = [
+            (
+                input[index].astype(np.int64)
+                if indice < 5
+                else np.asarray(input[index].astype(np.int64))
+            )
+            for indice, input in enumerate(self.inputs)
+        ]
+        # TODO: whether to use reversed mask by changing 1s and 0s to be
+        # consistent with nv bert
+        input_mask = (
+            1
+            - np.reshape(
+                input_mask.astype(np.float32), [1, 1, input_mask.shape[0]]
+            )
+        ) * -1e9
+
+        index = self.max_pred_length
+        # store number of  masked tokens in index
+        # outputs of torch.nonzero diff with that of numpy.nonzero by zip
+        padded_mask_indices = (masked_lm_positions == 0).nonzero()[0]
+        if len(padded_mask_indices) != 0:
+            index = padded_mask_indices[0].item()
+        else:
+            index = self.max_pred_length
+        # masked_lm_labels = np.full(input_ids.shape, -1, dtype=np.int64)
+        # masked_lm_labels[masked_lm_positions[:index]] = masked_lm_ids[:index]
+        masked_lm_labels = masked_lm_ids[:index]
+        masked_lm_positions = masked_lm_positions[:index]
+        # softmax_with_cross_entropy enforce last dim size equal 1
+        masked_lm_labels = np.expand_dims(masked_lm_labels, axis=-1)
+        next_sentence_labels = np.expand_dims(next_sentence_labels, axis=-1)
+        return [
+            input_ids,
+            segment_ids,
+            input_mask,
+            masked_lm_positions,
+            masked_lm_labels,
+            next_sentence_labels,
+        ]
+
+
+def create_pretraining_dataset(
+    input_file, max_pred_length, shared_list, batch_size, worker_init
+):
+    train_data = PretrainingDataset(
+        input_file=input_file, max_pred_length=max_pred_length
+    )
+    # files have been sharded, no need to dispatch again
+    train_batch_sampler = paddle.io.BatchSampler(
+        train_data, batch_size=batch_size, shuffle=True
+    )
+
+    # DataLoader cannot be pickled because of its place.
+    # If it can be pickled, use global function instead of lambda and use
+    # ProcessPoolExecutor instead of ThreadPoolExecutor to prefetch.
+    def _collate_data(data, stack_fn=Stack()):
+        num_fields = len(data[0])
+        out = [None] * num_fields
+        # input_ids, segment_ids, input_mask, masked_lm_positions,
+        # masked_lm_labels, next_sentence_labels, mask_token_num
+        for i in (0, 1, 2, 5):
+            out[i] = stack_fn([x[i] for x in data])
+        _, seq_length = out[0].shape
+        size = sum(len(x[3]) for x in data)
+        # Padding for divisibility by 8 for fp16 or int8 usage
+        if size % 8 != 0:
+            size += 8 - (size % 8)
+        # masked_lm_positions
+        # Organize as a 1D tensor for gather or use gather_nd
+        out[3] = np.full(size, 0, dtype=np.int32)
+        # masked_lm_labels
+        out[4] = np.full([size, 1], -1, dtype=np.int64)
+        mask_token_num = 0
+        for i, x in enumerate(data):
+            for j, pos in enumerate(x[3]):
+                out[3][mask_token_num] = i * seq_length + pos
+                out[4][mask_token_num] = x[4][j]
+                mask_token_num += 1
+        # mask_token_num
+        out.append(np.asarray([mask_token_num], dtype=np.float32))
+        return out
+
+    train_data_loader = DataLoader(
+        dataset=train_data,
+        batch_sampler=train_batch_sampler,
+        collate_fn=_collate_data,
+        num_workers=0,
+        worker_init_fn=worker_init,
+        return_list=True,
+    )
+    return train_data_loader
+
+
+if __name__ == '__main__':
+    bert = Bert()

test/prim/model/test_bert_cinn.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+import unittest
+
+import numpy as np
+from bert import Bert, BertPretrainingCriterion, create_pretraining_dataset
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.dataset.common import DATA_HOME, download
+
+SEED = 2023
+BATCH_SIZE = 2
+
+URL = 'https://paddle-ci.gz.bcebos.com/prim_cinn/bert_training_data.npz'
+MODULE_NAME = 'test_bert_prim_cinn'
+MD5SUM = '71e730ee8d7aa77a215b7e898aa089af'
+SAVE_NAME = 'bert_training_data.npz'
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    base.core._set_prim_all_enabled(enable_prim)
+
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    # paddle.framework.random._manual_program_seed(SEED)
+
+    train_data_loader = create_pretraining_dataset(
+        os.path.join(DATA_HOME, MODULE_NAME, SAVE_NAME),
+        20,
+        {},
+        batch_size=BATCH_SIZE,
+        worker_init=None,
+    )
+
+    # Now only apply dy2st for encoder
+    bert = Bert(to_static, enable_cinn)
+    criterion = BertPretrainingCriterion()
+
+    optimizer = paddle.optimizer.Adam(parameters=bert.parameters())
+
+    losses = []
+    for step, batch in enumerate(train_data_loader):
+        start_time = time.time()
+        (
+            input_ids,
+            segment_ids,
+            input_mask,
+            masked_lm_positions,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        ) = batch
+
+        prediction_scores, seq_relationship_score = bert(
+            input_ids=input_ids,
+            token_type_ids=segment_ids,
+            attention_mask=input_mask,
+            masked_positions=masked_lm_positions,
+        )
+
+        loss = criterion(
+            prediction_scores,
+            seq_relationship_score,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        )
+
+        loss.backward()
+        optimizer.minimize(loss)
+        bert.clear_gradients()
+        losses.append(loss.numpy().item())
+
+        print(
+            f"step: {step}, loss: {loss.numpy()}, batch_cost: {time.time() - start_time:.5}"
+        )
+        if step >= 9:
+            break
+    print(losses)
+    return losses
+
+
+class TestBert(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        download(URL, MODULE_NAME, MD5SUM, SAVE_NAME)
+
+    def tearDown(self):
+        paddle.set_flags({'FLAGS_deny_cinn_ops': ''})
+
+    @unittest.skipIf(
+        not (paddle.is_compiled_with_cinn() and paddle.is_compiled_with_cuda()),
+        "paddle is not compiled with CINN and CUDA",
+    )
+    def test_cinn(self):
+        paddle.set_flags({'FLAGS_deny_cinn_ops': "dropout"})
+        dy2st_cinn = train(to_static=True, enable_prim=False, enable_cinn=True)
+        dy2st_pir = train(to_static=True, enable_prim=False, enable_cinn=False)
+        np.testing.assert_allclose(dy2st_cinn, dy2st_pir, rtol=1e-5)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_bert_prim.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+import unittest
+
+import numpy as np
+from bert import Bert, BertPretrainingCriterion, create_pretraining_dataset
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.dataset.common import DATA_HOME, download
+
+SEED = 2023
+BATCH_SIZE = 2
+
+URL = 'https://paddle-ci.gz.bcebos.com/prim_cinn/bert_training_data.npz'
+MODULE_NAME = 'test_bert_prim_cinn'
+MD5SUM = '71e730ee8d7aa77a215b7e898aa089af'
+SAVE_NAME = 'bert_training_data.npz'
+
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    base.core._set_prim_all_enabled(enable_prim)
+
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    # paddle.framework.random._manual_program_seed(SEED)
+
+    train_data_loader = create_pretraining_dataset(
+        os.path.join(DATA_HOME, MODULE_NAME, SAVE_NAME),
+        20,
+        {},
+        batch_size=BATCH_SIZE,
+        worker_init=None,
+    )
+
+    # Now only apply dy2st for encoder
+    bert = Bert(to_static, enable_cinn)
+    criterion = BertPretrainingCriterion()
+
+    optimizer = paddle.optimizer.Adam(parameters=bert.parameters())
+
+    losses = []
+    for step, batch in enumerate(train_data_loader):
+        start_time = time.time()
+        (
+            input_ids,
+            segment_ids,
+            input_mask,
+            masked_lm_positions,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        ) = batch
+
+        prediction_scores, seq_relationship_score = bert(
+            input_ids=input_ids,
+            token_type_ids=segment_ids,
+            attention_mask=input_mask,
+            masked_positions=masked_lm_positions,
+        )
+
+        loss = criterion(
+            prediction_scores,
+            seq_relationship_score,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        )
+
+        loss.backward()
+        optimizer.minimize(loss)
+        bert.clear_gradients()
+        losses.append(loss.numpy().item())
+
+        print(
+            f"step: {step}, loss: {loss.numpy()}, batch_cost: {time.time() - start_time:.5}"
+        )
+        if step >= 9:
+            break
+    print(losses)
+    return losses
+
+
+class TestBert(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        download(URL, MODULE_NAME, MD5SUM, SAVE_NAME)
+
+    def tearDown(self):
+        paddle.set_flags({'FLAGS_deny_cinn_ops': ''})
+
+    @unittest.skipIf(
+        not (paddle.is_compiled_with_cinn() and paddle.is_compiled_with_cuda()),
+        "paddle is not compiled with CINN and CUDA",
+    )
+    def test_prim(self):
+        if "H20" in paddle.cuda.get_device_name():
+            DY2ST_PRIM_GT = [
+                10.834290504455566,
+                10.328838348388672,
+                10.342059135437012,
+                10.281204223632812,
+                10.226964950561523,
+                10.220486640930176,
+                10.174433708190918,
+                10.127359390258789,
+                10.134778022766113,
+                10.03632926940918,
+            ]
+        else:
+            DY2ST_PRIM_GT = [
+                10.649632453918457,
+                10.333406448364258,
+                10.33541202545166,
+                10.260543823242188,
+                10.219606399536133,
+                10.176884651184082,
+                10.124699592590332,
+                10.072620391845703,
+                10.112163543701172,
+                9.969392776489258,
+            ]
+        dy2st_prim = train(to_static=True, enable_prim=True, enable_cinn=False)
+        np.testing.assert_allclose(dy2st_prim, DY2ST_PRIM_GT, rtol=1e-5)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_bert_prim_cinn.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+import unittest
+
+import numpy as np
+from bert import Bert, BertPretrainingCriterion, create_pretraining_dataset
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.dataset.common import DATA_HOME, download
+
+SEED = 2023
+BATCH_SIZE = 2
+
+URL = 'https://paddle-ci.gz.bcebos.com/prim_cinn/bert_training_data.npz'
+MODULE_NAME = 'test_bert_prim_cinn'
+MD5SUM = '71e730ee8d7aa77a215b7e898aa089af'
+SAVE_NAME = 'bert_training_data.npz'
+
+
+DY2ST_PRIM_CINN_GT = [
+    11.086677551269531,
+    10.342002868652344,
+    10.336370468139648,
+    10.272554397583008,
+    10.224645614624023,
+    10.184449195861816,
+    10.14853572845459,
+    10.096378326416016,
+    10.117865562438965,
+    9.99058723449707,
+]
+
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    base.core._set_prim_all_enabled(enable_prim)
+
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    # paddle.framework.random._manual_program_seed(SEED)
+
+    train_data_loader = create_pretraining_dataset(
+        os.path.join(DATA_HOME, MODULE_NAME, SAVE_NAME),
+        20,
+        {},
+        batch_size=BATCH_SIZE,
+        worker_init=None,
+    )
+
+    # Now only apply dy2st for encoder
+    bert = Bert(to_static, enable_cinn)
+    criterion = BertPretrainingCriterion()
+
+    optimizer = paddle.optimizer.Adam(parameters=bert.parameters())
+
+    losses = []
+    for step, batch in enumerate(train_data_loader):
+        start_time = time.time()
+        (
+            input_ids,
+            segment_ids,
+            input_mask,
+            masked_lm_positions,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        ) = batch
+
+        prediction_scores, seq_relationship_score = bert(
+            input_ids=input_ids,
+            token_type_ids=segment_ids,
+            attention_mask=input_mask,
+            masked_positions=masked_lm_positions,
+        )
+
+        loss = criterion(
+            prediction_scores,
+            seq_relationship_score,
+            masked_lm_labels,
+            next_sentence_labels,
+            masked_lm_scale,
+        )
+
+        loss.backward()
+        optimizer.minimize(loss)
+        bert.clear_gradients()
+        losses.append(loss.numpy().item())
+
+        print(
+            f"step: {step}, loss: {loss.numpy()}, batch_cost: {time.time() - start_time:.5}"
+        )
+        if step >= 9:
+            break
+    print(losses)
+    return losses
+
+
+class TestBert(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        download(URL, MODULE_NAME, MD5SUM, SAVE_NAME)
+
+    def tearDown(self):
+        paddle.set_flags({'FLAGS_deny_cinn_ops': ''})
+
+    @unittest.skipIf(
+        not (paddle.is_compiled_with_cinn() and paddle.is_compiled_with_cuda()),
+        "paddle is not compiled with CINN and CUDA",
+    )
+    def test_prim_cinn(self):
+        dy2st_prim_cinn = train(
+            to_static=True, enable_prim=True, enable_cinn=True
+        )
+        np.testing.assert_allclose(
+            dy2st_prim_cinn, DY2ST_PRIM_CINN_GT, rtol=1e-5
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_comp_model_simple_net.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core, framework
+
+
+@param.parameterized_class(
+    ('name', 'primals', 'stop_gradients', 'cotangents', 'dtype'),
+    (
+        (
+            'test_normal_case',
+            (np.random.rand(2, 3, 4), np.random.rand(2, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_broadcast_diff_rank',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(3, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_broadcast_same_rank',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(2, 1, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_stop_gradient',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(2, 1, 3, 4)),
+            (False, True),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+    ),
+)
+class TestMultiplyGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primals = tuple(primal.astype(cls.dtype) for primal in cls.primals)
+        cls.cotangents = tuple(co.astype(cls.dtype) for co in cls.cotangents)
+
+    def setUp(self):
+        paddle.enable_static()
+
+    def tearDown(self):
+        paddle.disable_static()
+
+    def as_tuple(self, x):
+        return (x,) if isinstance(x, framework.Variable) else x
+
+    def net(self):
+        primals, cotangents = self.primals, self.cotangents
+        mp, sp = paddle.static.Program(), paddle.static.Program()
+        with paddle.static.program_guard(mp, sp):
+            primals = tuple(
+                paddle.static.data(f'primal{i}', primal.shape, primal.dtype)
+                for i, primal in enumerate(primals)
+            )
+            for primal, flag in zip(primals, self.stop_gradients):
+                primal.stop_gradient = flag
+            cotangents = tuple(
+                paddle.static.data(f'cotangent{i}', co.shape, co.dtype)
+                for i, co in enumerate(cotangents)
+            )
+            out = self.as_tuple(paddle.tanh(paddle.multiply(*primals)))
+            grads = paddle.static.gradients(out, primals)
+        exe = paddle.static.Executor()
+        exe.run(sp)
+        return exe.run(
+            program=mp,
+            feed={f'primal{i}': primal for i, primal in enumerate(self.primals)}
+            | {f'cotangent{i}': co for i, co in enumerate(self.cotangents)},
+            fetch_list=[g for g in grads if g is not None],
+        )
+
+    def test_comp(self):
+        core._set_prim_backward_enabled(True)
+        actual = self.net()
+
+        core._set_prim_backward_enabled(False)
+        desired = self.net()
+
+        self.assertEqual(len(actual), len(desired))
+        for i, j in zip(actual, desired):
+            np.testing.assert_allclose(
+                i,
+                j,
+                rtol=1e-6,
+                atol=0,
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_resnet_cinn.py

@@ -0,0 +1,191 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.vision.models import resnet50
+
+SEED = 2020
+base_lr = 0.001
+momentum_rate = 0.9
+l2_decay = 1e-4
+batch_size = 2
+epoch_num = 1
+
+# In V100, 16G, CUDA 11.2, the results are as follows:
+# DY2ST_CINN_GT = [
+#     5.847336769104004,
+#     8.336246490478516,
+#     5.108744144439697,
+#     8.316713333129883,
+#     8.175262451171875,
+#     7.590441703796387,
+#     9.895681381225586,
+#     8.196207046508789,
+#     8.438933372497559,
+#     10.305074691772461,
+
+# note: Version 2.0 momentum is fused to OP when L2Decay is available, and the results are different from the base version.
+# The results in ci as as follows:
+DY2ST_CINN_GT = [
+    5.847333908081055,
+    8.347576141357422,
+    5.1415300369262695,
+    8.373777389526367,
+    8.05331802368164,
+    7.437496185302734,
+    9.630914688110352,
+    8.547889709472656,
+    8.343082427978516,
+    10.203847885131836,
+]
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def reader_decorator(reader):
+    def __reader__():
+        for item in reader():
+            img = np.array(item[0]).astype('float32').reshape(3, 224, 224)
+            label = np.array(item[1]).astype('int64').reshape(1)
+            yield img, label
+
+    return __reader__
+
+
+class TransedFlowerDataSet(paddle.io.Dataset):
+    def __init__(self, flower_data, length):
+        self.img = []
+        self.label = []
+        self.flower_data = flower_data()
+        self._generate(length)
+
+    def _generate(self, length):
+        for i, data in enumerate(self.flower_data):
+            if i >= length:
+                break
+            self.img.append(data[0])
+            self.label.append(data[1])
+
+    def __getitem__(self, idx):
+        return self.img[idx], self.label[idx]
+
+    def __len__(self):
+        return len(self.img)
+
+
+def optimizer_setting(parameter_list=None):
+    optimizer = paddle.optimizer.Momentum(
+        learning_rate=base_lr,
+        momentum=momentum_rate,
+        weight_decay=paddle.regularizer.L2Decay(l2_decay),
+        parameters=parameter_list,
+    )
+
+    return optimizer
+
+
+def run(model, data_loader, optimizer, mode):
+    if mode == 'train':
+        model.train()
+        end_step = 9
+    elif mode == 'eval':
+        model.eval()
+        end_step = 1
+
+    for epoch in range(epoch_num):
+        total_acc1 = 0.0
+        total_acc5 = 0.0
+        total_sample = 0
+        losses = []
+
+        for batch_id, data in enumerate(data_loader()):
+            start_time = time.time()
+            img, label = data
+
+            pred = model(img)
+            avg_loss = paddle.nn.functional.cross_entropy(
+                input=pred,
+                label=label,
+                soft_label=False,
+                reduction='mean',
+                use_softmax=True,
+            )
+
+            acc_top1 = paddle.static.accuracy(input=pred, label=label, k=1)
+            acc_top5 = paddle.static.accuracy(input=pred, label=label, k=5)
+
+            if mode == 'train':
+                avg_loss.backward()
+                optimizer.minimize(avg_loss)
+                model.clear_gradients()
+
+            total_acc1 += acc_top1
+            total_acc5 += acc_top5
+            total_sample += 1
+            losses.append(avg_loss.numpy().item())
+
+            end_time = time.time()
+            print(
+                f"[{mode}]epoch {epoch} | batch step {batch_id}, "
+                f"loss {avg_loss:0.8f}, "
+                f"acc1 {total_acc1.numpy() / total_sample:0.3f}, "
+                f"acc5 {total_acc5.numpy() / total_sample:0.3f}, "
+                f"time {end_time - start_time:f}"
+            )
+            if batch_id >= end_step:
+                break
+    print(losses)
+    return losses
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    paddle.framework.random._manual_program_seed(SEED)
+    base.core._set_prim_all_enabled(enable_prim)
+
+    dataset = TransedFlowerDataSet(
+        reader_decorator(paddle.dataset.flowers.train(use_xmap=False)),
+        batch_size * (10 + 1),
+    )
+    data_loader = paddle.io.DataLoader(
+        dataset, batch_size=batch_size, drop_last=True
+    )
+
+    resnet = resnet50(False)
+    if to_static:
+        backend = "CINN" if enable_cinn else None
+        resnet = paddle.jit.to_static(resnet, backend=backend, full_graph=True)
+    optimizer = optimizer_setting(parameter_list=resnet.parameters())
+
+    train_losses = run(resnet, data_loader, optimizer, 'train')
+    if to_static and enable_prim and enable_cinn:
+        eval_losses = run(resnet, data_loader, optimizer, 'eval')
+    return train_losses
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_resnet_prim.py

@@ -0,0 +1,219 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.vision.models import resnet50
+
+SEED = 2020
+base_lr = 0.001
+momentum_rate = 0.9
+l2_decay = 1e-4
+batch_size = 2
+epoch_num = 1
+
+# In V100, 16G, CUDA 11.2, the results are as follows:
+# DY2ST_PRIM_GT = [
+#     5.8473358154296875,
+#     8.354944229125977,
+#     5.098367691040039,
+#     8.533346176147461,
+#     8.179085731506348,
+#     7.285282135009766,
+#     9.824585914611816,
+#     8.56928825378418,
+#     8.539499282836914,
+#     10.256929397583008,
+# ]
+
+# note: Version 2.0 momentum is fused to OP when L2Decay is available, and the results are different from the base version.
+# The results in ci as as follows:
+DY2ST_PRIM_GT = [
+    8.852441787719727,
+    8.403528213500977,
+    7.157894134521484,
+    8.536691665649414,
+    7.061797142028809,
+    7.612838268280029,
+    7.831070423126221,
+    8.562232971191406,
+    8.49091911315918,
+    7.967043876647949,
+]
+
+# IN V100, 16G, CUDA 12.0, the results are as follows:
+DY2ST_PRIM_GT_CUDA12 = [
+    8.852443695068359,
+    8.403528213500977,
+    7.158357620239258,
+    8.539973258972168,
+    7.070737838745117,
+    7.629122734069824,
+    7.809022426605225,
+    8.636153221130371,
+    8.410324096679688,
+    7.992737293243408,
+]
+
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def reader_decorator(reader):
+    def __reader__():
+        for item in reader():
+            img = np.array(item[0]).astype('float32').reshape(3, 224, 224)
+            label = np.array(item[1]).astype('int64').reshape(1)
+            yield img, label
+
+    return __reader__
+
+
+class TransedFlowerDataSet(paddle.io.Dataset):
+    def __init__(self, flower_data, length):
+        self.img = []
+        self.label = []
+        self.flower_data = flower_data()
+        self._generate(length)
+
+    def _generate(self, length):
+        for i, data in enumerate(self.flower_data):
+            if i >= length:
+                break
+            self.img.append(data[0])
+            self.label.append(data[1])
+
+    def __getitem__(self, idx):
+        return self.img[idx], self.label[idx]
+
+    def __len__(self):
+        return len(self.img)
+
+
+def optimizer_setting(parameter_list=None):
+    optimizer = paddle.optimizer.Momentum(
+        learning_rate=base_lr,
+        momentum=momentum_rate,
+        weight_decay=paddle.regularizer.L2Decay(l2_decay),
+        parameters=parameter_list,
+    )
+
+    return optimizer
+
+
+def run(model, data_loader, optimizer, mode):
+    if mode == 'train':
+        model.train()
+        end_step = 9
+    elif mode == 'eval':
+        model.eval()
+        end_step = 1
+
+    for epoch in range(epoch_num):
+        total_acc1 = 0.0
+        total_acc5 = 0.0
+        total_sample = 0
+        losses = []
+
+        for batch_id, data in enumerate(data_loader()):
+            start_time = time.time()
+            img, label = data
+            pred = model(img)
+            avg_loss = paddle.nn.functional.cross_entropy(
+                input=pred,
+                label=label,
+                soft_label=False,
+                reduction='mean',
+                use_softmax=True,
+            )
+            acc_top1 = paddle.static.accuracy(input=pred, label=label, k=1)
+            acc_top5 = paddle.static.accuracy(input=pred, label=label, k=5)
+
+            if mode == 'train':
+                avg_loss.backward()
+                optimizer.minimize(avg_loss)
+                model.clear_gradients()
+
+            total_acc1 += acc_top1
+            total_acc5 += acc_top5
+            total_sample += 1
+            losses.append(avg_loss.numpy().item())
+
+            end_time = time.time()
+            print(
+                f"[{mode}]epoch {epoch} | batch step {batch_id}, "
+                f"loss {avg_loss:0.15f}, "
+                f"acc1 {total_acc1.numpy() / total_sample:0.3f}, "
+                f"acc5 {total_acc5.numpy() / total_sample:0.3f}, "
+                f"time {end_time - start_time:f}"
+            )
+            if batch_id >= end_step:
+                break
+    return losses
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    paddle.framework.random._manual_program_seed(SEED)
+    base.core._set_prim_all_enabled(enable_prim)
+
+    dataset = TransedFlowerDataSet(
+        reader_decorator(paddle.dataset.flowers.train(use_xmap=False)),
+        batch_size * (10 + 1),
+    )
+    data_loader = paddle.io.DataLoader(
+        dataset, batch_size=batch_size, drop_last=True
+    )
+    resnet = resnet50(True)
+    if to_static:
+        backend = "CINN" if enable_cinn else None
+        resnet = paddle.jit.to_static(resnet, backend=backend, full_graph=True)
+    optimizer = optimizer_setting(parameter_list=resnet.parameters())
+
+    train_losses = run(resnet, data_loader, optimizer, 'train')
+    if to_static and enable_prim and enable_cinn:
+        eval_losses = run(resnet, data_loader, optimizer, 'eval')
+    return train_losses
+
+
+class TestResnet(unittest.TestCase):
+    @unittest.skipIf(
+        not (paddle.is_compiled_with_cinn() and paddle.is_compiled_with_cuda()),
+        "paddle is not compiled with CINN and CUDA",
+    )
+    def test_prim(self):
+        dy2st_prim = train(to_static=True, enable_prim=True, enable_cinn=False)
+        standard_prim = DY2ST_PRIM_GT
+
+        if paddle.version.cuda() == "12.0":
+            standard_prim = DY2ST_PRIM_GT_CUDA12
+        np.testing.assert_allclose(
+            dy2st_prim, standard_prim, rtol=2e-2, atol=1e-2
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/model/test_resnet_prim_cinn.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.vision.models import resnet50
+
+SEED = 2020
+base_lr = 0.001
+momentum_rate = 0.9
+l2_decay = 1e-4
+batch_size = 2
+epoch_num = 1
+
+# In V100, 16G, CUDA 11.2, the results are as follows:
+# DY2ST_PRIM_CINN_GT = [
+#     5.8473358154296875,
+#     8.322463989257812,
+#     5.169863700866699,
+#     8.399882316589355,
+#     7.859550476074219,
+#     7.4672698974609375,
+#     9.828727722167969,
+#     8.270355224609375,
+#     8.456792831420898,
+#     9.919631958007812,
+# ]
+
+# note: Version 2.0 momentum is fused to OP when L2Decay is available, and the results are different from the base version.
+# The results in ci as as follows:
+DY2ST_PRIM_CINN_GT = [
+    5.847333908081055,
+    8.342670440673828,
+    5.130363941192627,
+    8.511886596679688,
+    8.13458251953125,
+    7.35969352722168,
+    9.874241828918457,
+    8.126291275024414,
+    8.637175559997559,
+    10.385666847229004,
+]
+
+if core.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
+
+
+def reader_decorator(reader):
+    def __reader__():
+        for item in reader():
+            img = np.array(item[0]).astype('float32').reshape(3, 224, 224)
+            label = np.array(item[1]).astype('int64').reshape(1)
+            yield img, label
+
+    return __reader__
+
+
+class TransedFlowerDataSet(paddle.io.Dataset):
+    def __init__(self, flower_data, length):
+        self.img = []
+        self.label = []
+        self.flower_data = flower_data()
+        self._generate(length)
+
+    def _generate(self, length):
+        for i, data in enumerate(self.flower_data):
+            if i >= length:
+                break
+            self.img.append(data[0])
+            self.label.append(data[1])
+
+    def __getitem__(self, idx):
+        return self.img[idx], self.label[idx]
+
+    def __len__(self):
+        return len(self.img)
+
+
+def optimizer_setting(parameter_list=None):
+    optimizer = paddle.optimizer.Momentum(
+        learning_rate=base_lr,
+        momentum=momentum_rate,
+        weight_decay=paddle.regularizer.L2Decay(l2_decay),
+        parameters=parameter_list,
+    )
+
+    return optimizer
+
+
+def run(model, data_loader, optimizer, mode):
+    if mode == 'train':
+        model.train()
+        end_step = 9
+    elif mode == 'eval':
+        model.eval()
+        end_step = 1
+
+    for epoch in range(epoch_num):
+        total_acc1 = 0.0
+        total_acc5 = 0.0
+        total_sample = 0
+        losses = []
+
+        for batch_id, data in enumerate(data_loader()):
+            start_time = time.time()
+            img, label = data
+
+            pred = model(img)
+            avg_loss = paddle.nn.functional.cross_entropy(
+                input=pred,
+                label=label,
+                soft_label=False,
+                reduction='mean',
+                use_softmax=True,
+            )
+
+            acc_top1 = paddle.static.accuracy(input=pred, label=label, k=1)
+            acc_top5 = paddle.static.accuracy(input=pred, label=label, k=5)
+
+            if mode == 'train':
+                avg_loss.backward()
+                optimizer.minimize(avg_loss)
+                model.clear_gradients()
+
+            total_acc1 += acc_top1
+            total_acc5 += acc_top5
+            total_sample += 1
+            losses.append(avg_loss.numpy().item())
+
+            end_time = time.time()
+            print(
+                f"[{mode}]epoch {epoch} | batch step {batch_id}, "
+                f"loss {avg_loss:0.8f}, "
+                f"acc1 {total_acc1.numpy() / total_sample:0.3f}, "
+                f"acc5 {total_acc5.numpy() / total_sample:0.3f}, "
+                f"time {end_time - start_time:f}"
+            )
+            if batch_id >= end_step:
+                break
+    print(losses)
+    return losses
+
+
+def train(to_static, enable_prim, enable_cinn):
+    if core.is_compiled_with_cuda():
+        paddle.set_device('gpu')
+    else:
+        paddle.set_device('cpu')
+    np.random.seed(SEED)
+    paddle.seed(SEED)
+    paddle.framework.random._manual_program_seed(SEED)
+    base.core._set_prim_all_enabled(enable_prim)
+
+    dataset = TransedFlowerDataSet(
+        reader_decorator(paddle.dataset.flowers.train(use_xmap=False)),
+        batch_size * (10 + 1),
+    )
+    data_loader = paddle.io.DataLoader(
+        dataset, batch_size=batch_size, drop_last=True
+    )
+
+    resnet = resnet50(False)
+    if to_static:
+        backend = "CINN" if enable_cinn else None
+        resnet = paddle.jit.to_static(resnet, backend=backend, full_graph=True)
+    optimizer = optimizer_setting(parameter_list=resnet.parameters())
+
+    train_losses = run(resnet, data_loader, optimizer, 'train')
+    if to_static and enable_prim and enable_cinn:
+        eval_losses = run(resnet, data_loader, optimizer, 'eval')
+    return train_losses
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/CMakeLists.txt

@@ -0,0 +1,98 @@
+set(TEST_PRIM_PURE_PIR_CASES
+    test_custom_vjp_trait
+    test_decomp_op
+    test_decompose_op
+    test_vjp_prim
+    test_batch_norm_shape_check
+    test_builtin_slice
+    test_prim_program
+    test_prim_simpnet
+    test_prim_custom_vjp
+    test_prim_jit
+    test_pir_prim_flags
+    test_pir_prim_flags_v2
+    test_sink_decomp
+    test_prim_skip_dynamic
+    test_prim_dynamic
+    test_prim_jit_dynamic
+    test_auto_recompute
+    test_auto_recompute_dy2static
+    test_prim_sub_graph_dynamic_shape
+    test_decompose_control_flow
+    test_decomp_whole_program
+    test_dynamic_combine1
+    test_dynamic_combine2
+    test_decomp_fallback
+    test_prim_amax_amin_op)
+
+foreach(target ${TEST_PRIM_PURE_PIR_CASES})
+  py_test_modules(
+    ${target}
+    MODULES
+    ${target}
+    ENVS
+    FLAGS_enable_pir_api=true
+    FLAGS_comp_skip_default_ops=0
+    FLAGS_prim_enable_dynamic=true)
+endforeach()
+
+set(TEST_PRIM_DYNAMIC_SHAPE_BACKWARD_PIR_CASES
+    test_prim_sub_graph_backward_dynamic_shape
+    test_prim_sub_graph_abcde_backward_dynamic_shape
+    test_prim_sub_graph_fghij_backward_dynamic_shape
+    test_prim_sub_graph_klmno_backward_dynamic_shape
+    test_prim_sub_graph_pqrst_backward_dynamic_shape
+    test_prim_sub_graph_uvwxyz_backward_dynamic_shape)
+
+foreach(target ${TEST_PRIM_DYNAMIC_SHAPE_BACKWARD_PIR_CASES})
+  py_test_modules(
+    ${target}
+    MODULES
+    ${target}
+    ENVS
+    FLAGS_prim_vjp_skip_default_ops=0
+    FLAGS_enable_pir_api=true
+    FLAGS_comp_skip_default_ops=0
+    FLAGS_prim_enable_dynamic=true)
+endforeach()
+
+py_test_modules(test_pir_prim_flags_v3 MODULES test_pir_prim_flags_v3 ENVS
+                FLAGS_enable_pir_api=true FLAGS_prim_vjp_skip_default_ops=0)
+
+set_tests_properties(test_auto_recompute PROPERTIES TIMEOUT 40)
+set_tests_properties(test_auto_recompute_dy2static PROPERTIES TIMEOUT 40)
+set_tests_properties(test_pir_prim_flags PROPERTIES TIMEOUT 150)
+
+set(TEST_PRIM_PURE_PIR_CINN test_prim_rms_norm_st_shape
+                            test_prim_flags_check_ops)
+
+if(WITH_CINN)
+  foreach(target ${TEST_PRIM_PURE_PIR_CINN})
+    py_test_modules(
+      ${target}
+      MODULES
+      ${target}
+      ENVS
+      FLAGS_prim_check_ops=true
+      FLAGS_enable_pir_api=true
+      FLAGS_prim_enable_dynamic=true
+      FLAGS_prim_vjp_skip_default_ops=false)
+    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=CINN")
+  endforeach()
+endif()
+
+foreach(target ${TEST_PRIM_TRANS_PIR_CASES})
+  py_test_modules(${target} MODULES ${target} ENVS
+                  FLAGS_enable_pir_in_executor=true)
+endforeach()
+
+set(TEST_PRIM_BACKWARD_BLACKLIST "test_pir_prim_flags_v4")
+
+py_test_modules(
+  ${TEST_PRIM_BACKWARD_BLACKLIST}
+  MODULES
+  ${TEST_PRIM_BACKWARD_BLACKLIST}
+  ENVS
+  FLAGS_enable_pir_api=true
+  FLAGS_prim_enable_dynamic=true
+  FLAGS_prim_backward_blacklist=pd_op.tanh_grad)

test/prim/pir_prim/test_auto_recompute.py

@@ -0,0 +1,182 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.autograd.ir_backward import grad as ir_grad
+from paddle.base import core
+from paddle.decomposition import decompose
+
+TOLERANCE = {
+    "float64": {"rtol": 1e-15, "atol": 1e-15},
+    "float32": {"rtol": 1e-6, "atol": 1e-6},
+    "float16": {"rtol": 1e-3, "atol": 1e-3},
+    "bfloat16": {"rtol": 1e-2, "atol": 1e-2},
+}
+
+
+def rms_norm(weight, hidden):
+    variance = paddle.mean(paddle.pow(hidden, 2), axis=-1, keepdim=True)
+    hidden = paddle.rsqrt(variance + 0.00001) * hidden
+    return hidden * weight
+
+
+places = []
+if (
+    os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+    in ['1', 'true', 'on']
+    or not paddle.is_compiled_with_cuda()
+):
+    places.append(paddle.CPUPlace())
+if paddle.is_compiled_with_cuda():
+    places.append(paddle.CUDAPlace(0))
+
+
+@param.parameterized_class(
+    ('name', 'inputs', 'dtype', 'places'),
+    (
+        (
+            "auto_recompute_rms_norm_test1",
+            [
+                np.random.random(size=[4096, 4096]),
+                np.random.random(size=[4096, 4096]),
+            ],
+            "float32",
+            places,
+        ),
+        (
+            "auto_recompute_rms_norm_test2",
+            [
+                np.random.random(size=[128, 256]),
+                np.random.random(size=[128, 256]),
+            ],
+            "float32",
+            places,
+        ),
+    ),
+)
+class TestAutoRecomputeRmsNorm(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.inputs = [
+            (
+                x.astype(cls.dtype)
+                if cls.dtype != "bfloat16"
+                else x.astype("float32")
+            )
+            for x in cls.inputs
+        ]
+        core._set_prim_all_enabled(True)
+        paddle.enable_static()
+
+    @classmethod
+    def tearDownClass(cls):
+        core._set_prim_all_enabled(False)
+        paddle.disable_static()
+
+    def product_rms_norm_inputs(self):
+        weight = paddle.static.data(
+            name="weight", shape=self.inputs[0].shape, dtype=self.dtype
+        )
+        hidden = paddle.static.data(
+            name="hidden", shape=self.inputs[1].shape, dtype=self.dtype
+        )
+        weight.stop_gradient = False
+        hidden.stop_gradient = False
+        return [weight, hidden]
+
+    def cal_rms_norm_decomp_res(self, place):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            weight, hidden = self.product_rms_norm_inputs()
+            out = rms_norm(weight, hidden)
+            out_grad = paddle.full(
+                shape=out.shape, fill_value=3, dtype="float32"
+            )
+            [out] = decompose(main_program, [out])
+            [dweight, dhidden] = ir_grad(out, [weight, hidden], out_grad)
+            exe = paddle.static.Executor(place)
+            res = exe.run(
+                feed={'weight': self.inputs[0], 'hidden': self.inputs[1]},
+                fetch_list=[dweight, dhidden],
+            )
+        return res, main_program
+
+    def cal_rms_norm_auto_recompute_decomp_res(self, place):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            weight, hidden = self.product_rms_norm_inputs()
+            out = rms_norm(weight, hidden)
+            out_grad = paddle.full(
+                shape=out.shape, fill_value=3, dtype="float32"
+            )
+            [out] = decompose(main_program, [out])
+            [dweight, dhidden] = ir_grad(out, [weight, hidden], out_grad)
+            main_program, _ = paddle.decomposition.auto_recompute(
+                main_program,
+                [weight, hidden],
+                [out],
+                grad_outputs=[out_grad],
+                fwd_op_end_idx=13,
+                backward_op_start_idx=15,
+            )
+            exe = paddle.static.Executor(place)
+            res = exe.run(
+                feed={'weight': self.inputs[0], 'hidden': self.inputs[1]},
+                fetch_list=[dweight, dhidden],
+            )
+        return res, main_program
+
+    def test_auto_recompute(self):
+        for place in places:
+            res_desire, orig_program = self.cal_rms_norm_decomp_res(place)
+            (
+                res_recompute,
+                recompute_program,
+            ) = self.cal_rms_norm_auto_recompute_decomp_res(place)
+            np.testing.assert_allclose(
+                res_desire[0],
+                res_recompute[0],
+                atol=TOLERANCE[self.dtype]["atol"],
+                rtol=TOLERANCE[self.dtype]["rtol"],
+            )
+            np.testing.assert_allclose(
+                res_desire[1],
+                res_recompute[1],
+                atol=TOLERANCE[self.dtype]["atol"],
+                rtol=TOLERANCE[self.dtype]["rtol"],
+            )
+            # The following code is related to coverage, although backward_ops,define_op,all_used_ops is not used, it needs to be retained
+            forward_ops = recompute_program.global_block().ops[:13]
+            backward_ops = recompute_program.global_block().ops[13:]
+            saved_values = forward_ops[10].results()[0]
+            define_op = saved_values.get_defining_op()
+            for op in forward_ops:
+                if op.name() == "pd_op.data":
+                    continue
+                op_results = op.results()
+                for op_result in op_results:
+                    if op_result.is_same(saved_values):
+                        continue
+                    else:
+                        all_used_ops = op_result.all_used_ops()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_auto_recompute_dy2static.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+TOLERANCE = {
+    "float64": {"rtol": 1e-15, "atol": 1e-15},
+    "float32": {"rtol": 1e-6, "atol": 1e-6},
+    "float16": {"rtol": 1e-3, "atol": 1e-3},
+    "bfloat16": {"rtol": 1e-2, "atol": 1e-2},
+}
+
+
+def rms_norm(weight, hidden):
+    variance = paddle.mean(paddle.pow(hidden, 2), axis=-1, keepdim=True)
+    hidden = paddle.rsqrt(variance + 0.00001) * hidden
+    return hidden * weight
+
+
+class PrimNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, weight, hidden):
+        out = rms_norm(weight, hidden)
+        return out
+
+
+places = ["cpu"]
+if paddle.is_compiled_with_cuda():
+    places.append("gpu")
+
+
+@param.parameterized_class(
+    ('name', 'inputs', 'dtype', 'places'),
+    (
+        (
+            "auto_recompute_rms_norm_test1",
+            [
+                np.random.random(size=[4096, 4096]),
+                np.random.random(size=[4096, 4096]),
+            ],
+            "float32",
+            places,
+        ),
+        (
+            "auto_recompute_rms_norm_test2",
+            [
+                np.random.random(size=[128, 256]),
+                np.random.random(size=[128, 256]),
+            ],
+            "float32",
+            places,
+        ),
+    ),
+)
+class TestDy2StaticAutoRecomputeRmsNorm(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.inputs = [
+            (
+                x.astype(cls.dtype)
+                if cls.dtype != "bfloat16"
+                else x.astype("float32")
+            )
+            for x in cls.inputs
+        ]
+
+    def product_rms_norm_inputs(self, place):
+        weight = paddle.to_tensor(self.inputs[0], dtype=self.dtype, place=place)
+        hidden = paddle.to_tensor(self.inputs[1], dtype=self.dtype, place=place)
+        weight.stop_gradient = False
+        hidden.stop_gradient = False
+        return [weight, hidden]
+
+    def cal_rms_norm_res(self, place):
+        weight, hidden = self.product_rms_norm_inputs(place)
+        net = PrimNet()
+        net = paddle.jit.to_static(net, full_graph=True, backend=None)
+        program = net.forward.get_concrete_program(weight, hidden)[
+            -1
+        ].program.program
+        out = net(weight, hidden)
+        [dweight, dhidden] = paddle.grad(out, [weight, hidden])
+        return program, out, dweight, dhidden
+
+    def prepare_run_desire_res(self):
+        if os.environ.get('FLAGS_enable_auto_recompute'):
+            del os.environ['FLAGS_enable_auto_recompute']
+        core._set_prim_all_enabled(False)
+
+    def prepare_run_actual_res(self):
+        os.environ['FLAGS_enable_auto_recompute'] = "1"
+        core._set_prim_all_enabled(True)
+
+    def test_auto_recompute(self):
+        for place in places:
+            self.prepare_run_desire_res()
+            res_desire = self.cal_rms_norm_res(place)
+
+            self.prepare_run_actual_res()
+            res_actual = self.cal_rms_norm_res(place)
+            for desire, actual in zip(res_desire[1:], res_actual[1:]):
+                np.testing.assert_allclose(
+                    desire,
+                    actual,
+                    atol=TOLERANCE[self.dtype]["atol"],
+                    rtol=TOLERANCE[self.dtype]["rtol"],
+                )
+            actual_program = res_actual[0]
+            forward_ops = actual_program.global_block().ops[:14]
+            mid_ops = actual_program.global_block().ops[14:17]
+            backward_ops = actual_program.global_block().ops[17:]
+            saved_values = forward_ops[10].results()[0]
+            define_op = saved_values.get_defining_op()
+            self.assertTrue(define_op.name() == "pd_op.rsqrt")
+            for op in forward_ops:
+                if op.name() == "pd_op.data":
+                    continue
+                op_results = op.results()
+                for op_result in op_results:
+                    if op_result.is_same(saved_values):
+                        continue
+                    else:
+                        all_used_ops = op_result.all_used_ops()
+                        for used_op in all_used_ops:
+                            self.assertTrue(used_op in forward_ops + mid_ops)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_batch_norm_shape_check.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.decomposition import decompose
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+def batch_norm_net1(x, r_m, r_v, w, b):
+    return paddle.nn.functional.batch_norm(x, r_m, r_v, w, b, training=False)
+
+
+class TestBuildOp(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [1, 64, 512, 1024]
+        self.c_shape = [64]
+        self.dtype_x = "float32"
+        self.init_x_shape = [1, 64, 512, 1024]
+        self.x = np.random.random(self.x_shape).astype(self.dtype_x)
+        self.r_m = np.random.random(self.x_shape[1]).astype(self.dtype)
+        self.r_v = np.random.random(self.x_shape[1]).astype(self.dtype)
+        self.w = np.random.random(self.x_shape[1]).astype(self.dtype)
+        self.b = np.random.random(self.x_shape[1]).astype(self.dtype)
+        self.net = batch_norm_net1
+        self.necessary_ops = "pd_op.batch_norm"
+        self.enable_cinn = False
+        self.tol = 5e-6
+
+    def get_ir_program(self):
+        paddle.enable_static()
+        x = paddle.randn([4, 4])
+        main_program, start_program = (
+            paddle.static.Program(),
+            paddle.static.Program(),
+        )
+        with paddle.static.program_guard(main_program, start_program):
+            x = paddle.static.data('x', self.x_shape, x.dtype)
+            x.stop_gradient = False
+            r_m = paddle.static.data('r_m', self.c_shape, x.dtype)
+            r_v = paddle.static.data('r_v', self.c_shape, x.dtype)
+            w = paddle.static.data('w', self.c_shape, x.dtype)
+            b = paddle.static.data('b', self.c_shape, x.dtype)
+            y = batch_norm_net1(x, r_m, r_v, w, b)
+            _ = paddle.tanh(y)
+        return main_program
+
+    def test_build_op(self):
+        pir_program = self.get_ir_program()
+        y = pir_program.global_block().ops[-2].results()
+        orig_shape = y[0].shape
+        with paddle.pir_utils.IrGuard():
+            core._set_prim_forward_enabled(True)
+            y_new = decompose(pir_program, y)
+            core._set_prim_forward_enabled(False)
+            new_shape = y_new[0].shape
+            assert orig_shape == new_shape, (
+                f"Original shape {orig_shape} is not equal to new shape {new_shape}"
+            )
+            op_name_list = [op.name() for op in pir_program.global_block().ops]
+            assert "pd_op.batch_norm_" not in op_name_list
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_builtin_slice.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.decomposition import decompose
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+def meshgrid_net(x1, x2, x3, x4):
+    return paddle.meshgrid(x1, x2, x3, x4)
+
+
+class TestBuildOp(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.c_shape = [64]
+        self.init_x_shape = [1, 64, 512, 1024]
+        self.x1 = np.random.random(self.c_shape).astype(self.dtype)
+        self.x2 = np.random.random(self.c_shape).astype(self.dtype)
+        self.x3 = np.random.random(self.c_shape).astype(self.dtype)
+        self.x4 = np.random.random(self.c_shape).astype(self.dtype)
+        self.net = meshgrid_net
+
+    def get_ir_program(self):
+        paddle.enable_static()
+        main_program, start_program = (
+            paddle.static.Program(),
+            paddle.static.Program(),
+        )
+        with paddle.static.program_guard(main_program, start_program):
+            x1 = paddle.static.data('x1', self.c_shape, self.dtype)
+            x2 = paddle.static.data('x2', self.c_shape, self.dtype)
+            x3 = paddle.static.data('x3', self.c_shape, self.dtype)
+            x4 = paddle.static.data('x4', self.c_shape, self.dtype)
+            y = meshgrid_net(x1, x2, x3, x4)
+            paddle.tanh(y[0])
+            paddle.sin(y[1])
+            paddle.cos(y[2])
+        return main_program
+
+    def test_build_op(self):
+        pir_program = self.get_ir_program()
+        y = pir_program.global_block().ops[-1].results()
+        orig_shape = y[0].shape
+        with paddle.pir_utils.IrGuard():
+            core._set_prim_forward_enabled(True)
+            y_new = decompose(pir_program, y)
+            core._set_prim_forward_enabled(False)
+            new_shape = y_new[0].shape
+            assert orig_shape == new_shape, (
+                f"Original shape {orig_shape} is not equal to new shape {new_shape}"
+            )
+            op_name_list = [op.name() for op in pir_program.global_block().ops]
+            assert "pd_op.meshgrid" not in op_name_list
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_custom_vjp_trait.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle import nn
+from paddle.base.core import has_custom_vjp
+
+paddle.enable_static()
+
+
+def get_gelu_program_pir():
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.static.data('x', [2, 3, 3], dtype='float32')
+        net = nn.GELU()
+        net(x)
+    return main_program
+
+
+def get_multiply_program_pir():
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.static.data('x', [2, 3, 3], dtype='float32')
+        y = paddle.static.data('y', [2, 3, 3], dtype='float32')
+        _ = x * y
+    return main_program
+
+
+class TestCustomVjpTrait(unittest.TestCase):
+    def test_gelu_op_custom_vjp_trait(self):
+        pir_program = get_gelu_program_pir()
+        op = pir_program.global_block().ops[-1]
+        self.assertEqual(op.name(), "pd_op.gelu")
+        self.assertEqual(has_custom_vjp(op), True)
+
+    def test_multiply_op_custom_vjp_trait(self):
+        pir_program = get_multiply_program_pir()
+        op = pir_program.global_block().ops[-1]
+        self.assertEqual(op.name(), "pd_op.multiply")
+        self.assertEqual(has_custom_vjp(op), False)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_decomp_fallback.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+class TestFallBackBase(unittest.TestCase):
+    def setUp(self):
+        self.func_api = None
+        self.dtype = np.float32
+        self.tol = 1e-6
+
+
+def custom_dropout(x, p):
+    return paddle.nn.functional.dropout(x, p) + 2.0
+
+
+class TestDropOutFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_dropout
+        self.x = paddle.to_tensor([[1.0, -2], [3.0, 4]], dtype=self.dtype)
+        self.p = paddle.to_tensor(0.0, dtype=self.dtype)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(
+            self.func_api, full_graph=True, backend=None
+        )
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.p)
+        ref_out = dynamic_func(self.x, self.p)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_full(shape, value):
+    return paddle.full_like(shape, value) + 2.0
+
+
+class TestFullLikeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_full
+        self.x = paddle.to_tensor([[1.0, -2], [3.0, 4]], dtype=self.dtype)
+        self.value = paddle.to_tensor(2, dtype=self.dtype)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(
+            self.func_api, full_graph=True, backend=None
+        )
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.value)
+        ref_out = dynamic_func(self.x, self.value)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_squeeze(x, axis):
+    return paddle.squeeze(x, axis) + 2.0
+
+
+class TestSqueezeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_squeeze
+        self.x = paddle.rand([5, 1, 10], dtype=self.dtype)
+        self.axis = paddle.to_tensor(1, dtype=paddle.int64)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(
+            self.func_api, full_graph=True, backend=None
+        )
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_unsqueeze(x, axis):
+    return paddle.unsqueeze(x, axis) + 2.0
+
+
+class TestUnsqueezeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_unsqueeze
+        self.x = paddle.rand([5, 10], dtype=self.dtype)
+        self.axis = paddle.to_tensor([0, 2], dtype=paddle.int64)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(
+            self.func_api, full_graph=True, backend=None
+        )
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_any(x, axis):
+    return paddle.any(x, axis)
+
+
+class TestAnyFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_any
+        self.x = paddle.to_tensor([[1, 0], [1, 1]], dtype='int32').cast('bool')
+        # Axis cannot accept a list of tensors,
+        # the framework will check the argument type before decomposition.
+        self.axis = [0]
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(
+            self.func_api, full_graph=True, backend=None
+        )
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_decomp_op.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.decomposition import decompose
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+def get_ir_program():
+    paddle.enable_static()
+    x = paddle.randn([4, 4])
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x_s = paddle.static.data('x', [4, 4], x.dtype)
+        x_s.stop_gradient = False
+        y_s = paddle.divide(x_s, x_s)
+        y_s = paddle.add(x_s, y_s)
+        y_s = paddle.mean(y_s)
+        y_s = paddle.tanh(y_s)
+    pir_program = main_program
+
+    all_ops = pir_program.global_block().ops
+    for op in all_ops:
+        op.op_role = 1
+
+    return pir_program
+
+
+class TestBuildOp(unittest.TestCase):
+    def test_build_op(self):
+        pir_program = get_ir_program()
+        y = pir_program.global_block().ops[-2].results()
+        orig_shape = y[0].shape
+        with paddle.pir_utils.IrGuard():
+            core._set_prim_forward_enabled(True)
+            y_new = decompose(pir_program, y)
+            core._set_prim_forward_enabled(False)
+            new_shape = y_new[0].shape
+            assert orig_shape == new_shape, (
+                f"Original shape {orig_shape} is not equal to new shape {new_shape}"
+            )
+            op_name_list = [op.name() for op in pir_program.global_block().ops]
+            self.assertEqual(
+                op_name_list,
+                [
+                    'pd_op.data',
+                    'pd_op.divide',
+                    'pd_op.add',
+                    'pd_op.full_int_array',
+                    'pd_op.full_int_array',
+                    'pd_op.sum',
+                    'pd_op.full',
+                    'pd_op.divide',
+                    'pd_op.tanh',
+                ],
+            )
+            op_role_list = [op.op_role for op in pir_program.global_block().ops]
+            self.assertEqual(
+                all(op_role == 1 for op_role in op_role_list), True
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_decomp_whole_program.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.autograd.ir_backward import grad
+from paddle.decomposition import decomp
+
+paddle.enable_static()
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [32, 32]
+        self.shape_y = [32, 32]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            x1 = paddle.sin(x)
+            y1 = paddle.cos(y)
+            y3 = paddle.matmul(x1, y1)
+            tmp1 = paddle.concat((x1, y1, y3))
+            tmp1 = paddle.slice(tmp1, axes=[1], starts=[0], ends=[2])
+            tmp2 = paddle.mean(tmp1)
+            sum_out = paddle.sin(tmp2)
+            gradients = grad(sum_out, (x, y))
+            if flag == "prim":
+                with decomp.prim_guard():
+                    decomp.decompose_dist_program(main_program)
+            exe = paddle.static.Executor()
+            [fwd, dx, dy] = exe.run(
+                feed={'x': self.x, 'y': self.y}, fetch_list=[sum_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "prim":
+            assert 'pd_op.concat_grad' not in whole_ops
+        else:
+            assert 'pd_op.concat_grad' in whole_ops
+
+        return fwd, dx, dy
+
+    def test_prim_all(self):
+        paddle.base.core._set_prim_backward_blacklist("sin_grad", "cos_grad")
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6, atol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_decompose_control_flow.py

@@ -0,0 +1,333 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pathlib
+import sys
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import base
+from paddle.autograd.ir_backward import grad as ir_grad
+from paddle.framework import core
+
+sys.path.append(
+    str(pathlib.Path(__file__).resolve().parents[2] / 'legacy_test')
+)
+from utils import static_guard
+
+
+class IfNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y, cond):
+        if cond:
+            x = x + 1
+            out1 = paddle.mean(x)
+            out2 = paddle.mean(y)
+        else:
+            y = y + 1
+            out1 = paddle.mean(x)
+            out2 = paddle.mean(y)
+        return out1, out2
+
+
+class WhileNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        while paddle.all(x < y):
+            x = x + 1
+            out = paddle.mean(y**2)
+        return out
+
+
+class WhileAndIfNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        while paddle.all(x < y):
+            if paddle.all(x + 1 < y):
+                x = x + 0.5
+                out = paddle.mean(y**2)
+            else:
+                x = x + 0.6
+                out = paddle.mean(paddle.nn.functional.softmax(y))
+        return out
+
+
+class TestPrimControlFlowIf(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core._set_prim_forward_enabled(False)
+
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32]
+        self.shape_y = [8, 16, 32]
+        self.x_np = np.random.random(self.shape_x).astype("float32")
+        self.y_np = np.random.random(self.shape_y).astype("float32")
+        self.places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            self.places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+    def get_control_if_res(self, x, y, cond, prim_forward=False):
+        if prim_forward:
+            core._set_prim_forward_enabled(True)
+        net = IfNet()
+        net = paddle.jit.to_static(net, full_graph=True)
+        out1, out2 = net(x, y, cond)
+        core._set_prim_forward_enabled(False)
+        return out1, out2
+
+    def test_decompose_if_true(self):
+        for place in self.places:
+            if isinstance(place, paddle.base.CPUPlace):
+                paddle.set_device("cpu")
+            elif isinstance(place, paddle.base.CUDAPlace):
+                paddle.set_device("gpu")
+            x = paddle.to_tensor(self.x_np, dtype="float32")
+            y = paddle.to_tensor(self.y_np, dtype="float32")
+            cond = paddle.full(shape=[1], fill_value=1)
+            out1_baseline, out2_baseline = self.get_control_if_res(
+                x, y, cond, prim_forward=False
+            )
+            out1, out2 = self.get_control_if_res(x, y, cond, prim_forward=True)
+            np.testing.assert_allclose(out1_baseline, out1, rtol=1e-6, atol=0)
+            np.testing.assert_allclose(out2_baseline, out2, rtol=1e-6, atol=0)
+
+    def test_decompose_if_false(self):
+        for place in self.places:
+            if isinstance(place, paddle.base.CPUPlace):
+                paddle.set_device("cpu")
+            elif isinstance(place, paddle.base.CUDAPlace):
+                paddle.set_device("gpu")
+            x = paddle.to_tensor(self.x_np, dtype="float32")
+            y = paddle.to_tensor(self.y_np, dtype="float32")
+            cond = paddle.full(shape=[1], fill_value=0)
+            out1_baseline, out2_baseline = self.get_control_if_res(
+                x, y, cond, prim_forward=False
+            )
+            out1, out2 = self.get_control_if_res(x, y, cond, prim_forward=True)
+            np.testing.assert_allclose(out1_baseline, out1, rtol=1e-6, atol=0)
+            np.testing.assert_allclose(out2_baseline, out2, rtol=1e-6, atol=0)
+
+    @classmethod
+    def tearDownClass(cls):
+        core._set_prim_forward_enabled(False)
+
+
+class TestPrimControlFlowWhile(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core._set_prim_forward_enabled(False)
+
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32]
+        self.shape_y = [8, 16, 32]
+        self.x_np = np.random.random(self.shape_x).astype("float32")
+        self.y_np = np.random.random(self.shape_y).astype("float32") + 3
+        self.places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            self.places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+    def get_control_while_res(self, x, y, prim_forward=False):
+        if prim_forward:
+            core._set_prim_forward_enabled(True)
+        net = WhileNet()
+        net = paddle.jit.to_static(net, full_graph=True)
+        out = net(x, y)
+        core._set_prim_forward_enabled(False)
+        return out
+
+    def test_decompose_while(self):
+        for place in self.places:
+            if isinstance(place, paddle.base.CPUPlace):
+                paddle.set_device("cpu")
+            elif isinstance(place, paddle.base.CUDAPlace):
+                paddle.set_device("gpu")
+            x = paddle.to_tensor(self.x_np, dtype="float32")
+            y = paddle.to_tensor(self.y_np, dtype="float32")
+            out_baseline = self.get_control_while_res(x, y, prim_forward=False)
+            out = self.get_control_while_res(x, y, prim_forward=True)
+            np.testing.assert_allclose(out_baseline, out, rtol=1e-6, atol=0)
+
+    @classmethod
+    def tearDownClass(cls):
+        core._set_prim_forward_enabled(False)
+
+
+class TestPrimControlFlowWhileAndIf(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core._set_prim_forward_enabled(False)
+
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32]
+        self.shape_y = [8, 16, 32]
+        self.x_np = np.random.random(self.shape_x).astype("float32")
+        self.y_np = np.random.random(self.shape_y).astype("float32") + 3
+        self.places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            self.places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+    def get_control_flow_res(self, x, y, prim_forward=False):
+        if prim_forward:
+            core._set_prim_forward_enabled(True)
+        net = WhileAndIfNet()
+        net = paddle.jit.to_static(net, full_graph=True)
+        out = net(x, y)
+        core._set_prim_forward_enabled(False)
+        return out
+
+    def test_decompose_while_and_if(self):
+        for place in self.places:
+            if isinstance(place, paddle.base.CPUPlace):
+                paddle.set_device("cpu")
+            elif isinstance(place, paddle.base.CUDAPlace):
+                paddle.set_device("gpu")
+            x = paddle.to_tensor(self.x_np, dtype="float32")
+            y = paddle.to_tensor(self.y_np, dtype="float32")
+            out_baseline = self.get_control_flow_res(x, y, prim_forward=False)
+            out = self.get_control_flow_res(x, y, prim_forward=True)
+            np.testing.assert_allclose(out_baseline, out, rtol=1e-6, atol=0)
+
+    @classmethod
+    def tearDownClass(cls):
+        core._set_prim_forward_enabled(False)
+
+
+class TestPrimControlFlowWhileBackward(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core._set_prim_all_enabled(False)
+
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_i = [1]
+        self.shape_x = [1]
+        self.i_np = np.random.random(self.shape_i).astype("float32")
+        self.x_np = np.random.random(self.shape_x).astype("float32")
+
+    def cond(self, i, x):
+        return i < 3
+
+    def body(self, i, x):
+        x = paddle.pow(x, i)
+        i = i + 1
+        return [i, x]
+
+    def get_while_prim_grad_res(self):
+        core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            i = paddle.static.data(name='i', shape=[1], dtype='float32')
+            i.stop_gradient = False
+            i.persistable = True
+            x = paddle.static.data(name='x', shape=[1], dtype='float32')
+            x.stop_gradient = False
+            x.persistable = True
+
+            out = paddle.static.nn.while_loop(self.cond, self.body, [i, x])
+            [new_out] = paddle.decomposition.decomp.decompose(
+                main_program, [out[1]]
+            )
+            out_grad = ir_grad(new_out, [x])
+
+        place = (
+            base.CUDAPlace(0)
+            if core.is_compiled_with_cuda()
+            else base.CPUPlace()
+        )
+        exe = base.Executor(place)
+        out_grad = exe.run(
+            main_program,
+            feed={'i': self.i_np, 'x': self.x_np},
+            fetch_list=[out_grad],
+        )
+        core._set_prim_all_enabled(False)
+        return main_program, out_grad[0]
+
+    def get_while_grad_res(self):
+        core._set_prim_all_enabled(False)
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            i = paddle.static.data(name='i', shape=[1], dtype='float32')
+            i.stop_gradient = False
+            i.persistable = True
+            x = paddle.static.data(name='x', shape=[1], dtype='float32')
+            x.stop_gradient = False
+            x.persistable = True
+
+            out = paddle.static.nn.while_loop(self.cond, self.body, [i, x])
+            out_grad = ir_grad(out, [x])
+
+        place = (
+            base.CUDAPlace(0)
+            if core.is_compiled_with_cuda()
+            else base.CPUPlace()
+        )
+        exe = base.Executor(place)
+        out_grad = exe.run(
+            main_program,
+            feed={'i': self.i_np, 'x': self.x_np},
+            fetch_list=[out_grad],
+        )
+        return main_program, out_grad[0]
+
+    def test_while_loop_backward2(self):
+        with static_guard():
+            program_origin, out_grad_baseline = self.get_while_grad_res()
+            program_prim, out_grad = self.get_while_prim_grad_res()
+        np.testing.assert_allclose(
+            out_grad_baseline, out_grad, rtol=1e-6, atol=0
+        )
+        assert len(
+            program_origin.global_block().ops[-1].as_while_op().body().ops
+        ) != len(program_prim.global_block().ops[-1].as_while_op().body().ops)
+
+    @classmethod
+    def tearDownClass(cls):
+        core._set_prim_all_enabled(False)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_decompose_op.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+
+import paddle
+from paddle import pir
+
+paddle.enable_static()
+
+
+def get_pir_program_and_param_map():
+    with paddle.pir_utils.OldIrGuard():
+        shape = [3, 3]
+        mp = paddle.static.Program()
+        with paddle.static.program_guard(mp):
+            # construct graph
+            x = paddle.static.data('x', shape, dtype='float32')
+            x.stop_gradient = False
+            y = paddle.static.data('y', shape, dtype='float32')
+            y.stop_gradient = False
+            z = paddle.static.data('z', shape, dtype='float32')
+            z.stop_gradient = False
+            tmp1 = paddle.add(x, y)
+            tmp2 = paddle.multiply(tmp1, z)
+            tmp3 = paddle.matmul(tmp2, z)
+            tmp4 = paddle.mean(tmp3, axis=-1, keepdim=True)
+            tmp5 = paddle.rsqrt(tmp4)
+            scale = paddle.tensor.fill_constant(
+                shape=tmp5.shape[1:],
+                dtype=tmp5.dtype,
+                value=1.0,
+            )
+            scale.stop_gradient = True
+            tmp6 = paddle.nn.functional.layer_norm(
+                tmp5, tmp5.shape[1:], scale, None, 1e-5
+            )
+            tmp7 = paddle.nn.functional.dropout(tmp6, p=0.5)
+            tmp8 = paddle.add(x, tmp7)
+            tmp9 = paddle.concat(tmp8)
+
+            test = paddle.rand([5, 1, 10])
+            _ = paddle.squeeze(test, axis=1)
+            out = paddle.mean(tmp9)
+            # construct backward graph
+            _ = paddle.static.gradients(out, [x, y, z])
+
+        pir_program, param_mapping = pir.translate_to_pir_with_param_map(
+            mp.desc
+        )
+        return pir_program, param_mapping
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_dynamic_combine1.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+def meshgrid_net(x, y, z):
+    temp = paddle.meshgrid(x, y)[0]
+    return paddle.concat([temp, z])
+
+
+class TestPrimMode1(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [4098]
+        self.shape_y = [4098]
+        self.shape_z = [4098, 4098]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.z = np.random.random(self.shape_z).astype("float32")
+        self.net = meshgrid_net
+        self.enable_cinn = False
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        z = paddle.to_tensor(self.z)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=[None], dtype='float32'),
+                    InputSpec(shape=[4098], dtype='float32'),
+                    InputSpec(shape=[4098, 4098], dtype='float32'),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y, z)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert "pd_op.meshgrid" not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        np.testing.assert_allclose(res_ref.numpy(), res.numpy(), rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_dynamic_combine2.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+def stack_net(x, y, z):
+    temp = paddle.stack([x, y], axis=-1)
+    return paddle.concat([temp, z])
+
+
+class TestPrimMode1(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [128, 128]
+        self.shape_y = [128, 128]
+        self.shape_z = [128, 128, 2]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.z = np.random.random(self.shape_z).astype("float32")
+        self.net = stack_net
+        self.enable_cinn = False
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        z = paddle.to_tensor(self.z)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=[None, None], dtype='float32'),
+                    InputSpec(shape=[None, None], dtype='float32'),
+                    InputSpec(shape=[None, 128, None], dtype='float32'),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y, z)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert "pd_op.stack" not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        np.testing.assert_allclose(res_ref.numpy(), res.numpy(), rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_pir_prim_flags.py

@@ -0,0 +1,270 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ast
+import os
+import re
+import subprocess
+import sys
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle.base import core
+from paddle.decomposition import decomp
+
+REGEX_FLAGS = re.compile(
+    r"Result\(prim_all=(?P<prim_all>.*), prim_fwd=(?P<prim_fwd>.*), prim_bwd=(?P<prim_bwd>.*)\)"
+)
+
+
+class TestPrimFlags(unittest.TestCase):
+    def test_prim_flags_default(self):
+        self.assertFalse(core._is_bwd_prim_enabled())
+        self.assertFalse(core._is_fwd_prim_enabled())
+        self.assertFalse(core._is_all_prim_enabled())
+
+    def check_prim_flags_under_subprocess(
+        self, instructions, env, expected_flags
+    ):
+        all_instrs = [
+            "import paddle",
+            *instructions,
+            "prim_all = paddle.base.core._is_all_prim_enabled()",
+            "prim_fwd = paddle.base.core._is_fwd_prim_enabled()",
+            "prim_bwd = paddle.base.core._is_bwd_prim_enabled()",
+            "print(f'Result(prim_all={prim_all}, prim_fwd={prim_fwd}, prim_bwd={prim_bwd})', end='')",
+        ]
+        inherited_env = os.environ.copy()
+        inherited_env.update(env)
+        result = subprocess.run(
+            [sys.executable, '-c', '; '.join(all_instrs)],
+            capture_output=True,
+            env=inherited_env,
+        )
+        if result.returncode != 0:
+            self.fail(f"Failed to run subprocess: {result.stderr}")
+        matched_flags = REGEX_FLAGS.search(result.stdout.decode())
+        self.assertIsNotNone(
+            matched_flags, f"Failed to parse flags: {result.stdout}"
+        )
+        flags = (
+            ast.literal_eval(matched_flags.group("prim_all")),
+            ast.literal_eval(matched_flags.group("prim_fwd")),
+            ast.literal_eval(matched_flags.group("prim_bwd")),
+        )
+        self.assertEqual(
+            flags, expected_flags, f"Expected: {expected_flags}, got: {flags}"
+        )
+
+    def test_prim_flags_under_subprocess(self):
+        # Check envs
+        self.check_prim_flags_under_subprocess(
+            [],
+            {},
+            (False, False, False),  # (prim_all, prim_fwd, prim_bwd)
+        )
+        self.check_prim_flags_under_subprocess(
+            [],
+            {"FLAGS_prim_backward": "True"},
+            (False, False, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            [],
+            {"FLAGS_prim_forward": "True"},
+            (False, True, False),
+        )
+        self.check_prim_flags_under_subprocess(
+            [],
+            {"FLAGS_prim_all": "True"},
+            (True, True, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            [],
+            {"FLAGS_prim_all": "True", "FLAGS_prim_forward": "False"},
+            (False, False, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            [],
+            {"FLAGS_prim_all": "True", "FLAGS_prim_backward": "False"},
+            (False, True, False),
+        )
+        # Check apis
+        self.check_prim_flags_under_subprocess(
+            ["paddle.base.core._set_prim_all_enabled(True)"],
+            {},
+            (True, True, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            ["paddle.base.core._set_prim_forward_enabled(True)"],
+            {},
+            (False, True, False),
+        )
+        self.check_prim_flags_under_subprocess(
+            ["paddle.base.core._set_prim_backward_enabled(True)"],
+            {},
+            (False, False, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            [
+                "paddle.base.core._set_prim_all_enabled(True)",
+                "paddle.base.core._set_prim_forward_enabled(False)",
+            ],
+            {},
+            (False, False, True),
+        )
+        self.check_prim_flags_under_subprocess(
+            [
+                "paddle.base.core._set_prim_all_enabled(True)",
+                "paddle.base.core._set_prim_backward_enabled(False)",
+            ],
+            {},
+            (False, True, False),
+        )
+        self.check_prim_flags_under_subprocess(
+            [
+                "paddle.base.core._set_prim_forward_enabled(True)",
+                "paddle.base.core._set_prim_backward_enabled(True)",
+            ],
+            {},
+            (True, True, True),
+        )
+        # Check envs and apis
+        self.check_prim_flags_under_subprocess(
+            [
+                "paddle.base.core._set_prim_all_enabled(False)",
+            ],
+            {
+                "FLAGS_prim_all": "True",
+            },
+            (False, False, False),
+        )
+
+
+class TestPrimBlacklistFlags(unittest.TestCase):
+    def not_in_blacklist(self, op_name):
+        inputs = np.random.random([2, 3, 4]).astype("float32")
+        paddle.enable_static()
+        core._set_prim_forward_enabled(True)
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            x = paddle.static.data(
+                'x', shape=inputs.shape, dtype=str(inputs.dtype)
+            )
+            y = F.gelu(x)
+            z = F.silu(y)
+
+            fwd_ops = [op.name() for op in main_program.global_block().ops]
+            # Ensure that tanh in original block
+            self.assertTrue(op_name in fwd_ops)
+
+            z = decomp.decompose(main_program, [z])
+
+            fwd_ops_new = [op.name() for op in main_program.global_block().ops]
+            # Ensure that tanh is split into small ops
+            self.assertTrue(op_name not in fwd_ops_new)
+
+        exe = paddle.static.Executor()
+        exe.run(startup_program)
+        _ = exe.run(main_program, feed={'x': inputs}, fetch_list=[z])
+        paddle.disable_static()
+        core._set_prim_forward_enabled(False)
+
+    def in_blacklist(self, op_name):
+        inputs = np.random.random([2, 3, 4]).astype("float32")
+        paddle.enable_static()
+        core._set_prim_forward_enabled(True)
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            x = paddle.static.data(
+                'x', shape=inputs.shape, dtype=str(inputs.dtype)
+            )
+            y = F.gelu(x)
+            z = F.silu(y)
+
+            fwd_ops = [op.name() for op in main_program.global_block().ops]
+            # Ensure that tanh in original block
+            self.assertTrue(op_name in fwd_ops)
+
+            z = decomp.decompose(main_program, [z])
+
+            fwd_ops_new = [op.name() for op in main_program.global_block().ops]
+            # Ensure that tanh is split into small ops
+            self.assertTrue(op_name in fwd_ops_new)
+
+        exe = paddle.static.Executor()
+        exe.run(startup_program)
+        _ = exe.run(main_program, feed={'x': inputs}, fetch_list=[z])
+        paddle.disable_static()
+        core._set_prim_forward_enabled(False)
+
+    def test_prim_forward_blacklist(self):
+        self.not_in_blacklist("pd_op.gelu")
+        core._set_prim_forward_blacklist("pd_op.gelu")
+        self.in_blacklist("pd_op.gelu")
+
+    def test_prim_forward_blacklist_flag(self):
+        self.not_in_blacklist("pd_op.silu")
+        paddle.set_flags({"FLAGS_prim_forward_blacklist": "pd_op.silu"})
+        self.in_blacklist("pd_op.silu")
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = paddle.tanh(x)
+        x2 = paddle.exp(x)
+        x3 = x1 + x2
+        res = paddle.nn.functional.gelu(x3)
+        return res
+
+
+class TestPrimBackwardBlacklistFlags(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([2, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.tanh_grad' in ops)
+        self.assertTrue('pd_op.exp_grad' in ops)
+        self.assertTrue('pd_op.gelu_grad' not in ops)
+
+    def test_prim_backward_blacklist(self):
+        core._set_prim_all_enabled(True)
+        core._set_prim_backward_blacklist("pd_op.tanh_grad", "pd_op.exp_grad")
+        self.train()
+        core._set_prim_all_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_pir_prim_flags_v2.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.base import core
+from paddle.decomposition import decomp
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = paddle.tanh(x)
+        x2 = paddle.exp(x)
+        res = paddle.matmul(x1, x2)
+        return res
+
+
+class TestPrimMatmulDefault(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([4, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.matmul_grad' in ops)
+
+    def test_prim_matmul_default(self):
+        with decomp.prim_guard():
+            self.train()
+
+
+class TestPrimMatmulDefaultRevert(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([4, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.matmul_grad' not in ops)
+
+    def test_prim_matmul_default(self):
+        core._clear_prim_vjp_skip_default_ops()
+        with decomp.prim_guard():
+            self.train()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_pir_prim_flags_v3.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.decomposition import decomp
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = paddle.tanh(x)
+        x2 = paddle.exp(x)
+        res = paddle.matmul(x1, x2)
+        return res
+
+
+class TestPrimMatmulDefaultRevert(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([4, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.matmul_grad' not in ops)
+
+    def test_prim_matmul_default(self):
+        with decomp.prim_guard():
+            self.train()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_pir_prim_flags_v4.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.base import core
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = paddle.tanh(x)
+        x2 = paddle.exp(x)
+        x3 = x1 + x2
+        res = paddle.nn.functional.gelu(x3)
+        return res
+
+
+class TestPrimBackwardBlacklistFlags(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([2, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.tanh_grad' in ops)
+        self.assertTrue('pd_op.exp_grad' not in ops)
+        self.assertTrue('pd_op.gelu_grad' not in ops)
+
+    def test_prim_backward_blacklist_flag(self):
+        core._set_prim_all_enabled(True)
+        self.train()
+        core._set_prim_all_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/pir_prim/test_prim_amax_amin_op.py

@@ -0,0 +1,289 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+class TestPrimBaseWithGrad(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = None
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = None
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x)
+        res.backward()
+        x_grad = x.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_all_enabled(False)
+        return res, x_grad
+
+    def test_prim_all(self):
+        if self.net is None:
+            return
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+def amax_net1(x):
+    return paddle.amax(x, keepdim=True)
+
+
+def amax_net2(x):
+    return paddle.amax(x, keepdim=False)
+
+
+def amax_net3(x):
+    return paddle.amax(x, axis=[0, 1], keepdim=False)
+
+
+def amax_net4(x):
+    return paddle.amax(x, axis=[-1, -2], keepdim=False)
+
+
+def amax_net5(x):
+    return paddle.amax(x, axis=[-1, 0], keepdim=False)
+
+
+def amin_net1(x):
+    return paddle.amin(x, keepdim=True)
+
+
+def amin_net2(x):
+    return paddle.amin(x, keepdim=False)
+
+
+def amin_net3(x):
+    return paddle.amin(x, axis=[0, 1], keepdim=False)
+
+
+def amin_net4(x):
+    return paddle.amin(x, axis=[-1, -2], keepdim=False)
+
+
+def amin_net5(x):
+    return paddle.amin(x, axis=[-1, 0], keepdim=False)
+
+
+class TestPrimAmaxWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.ones(self.x_shape).astype(self.dtype)
+        self.net = amax_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAmaxWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [30]
+        self.init_x_shape = [30]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amax_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAmaxWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amax_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAmaxWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amax_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAmaxWithGrad5(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x[2] = self.x[4]
+        self.net = amax_net4
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAmaxWithGrad6(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amax_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amax_net5
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.ones(self.x_shape).astype(self.dtype)
+        self.net = amin_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [30]
+        self.init_x_shape = [30]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amin_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amin_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amin_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad5(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = amin_net4
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimAminWithGrad6(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.amin_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [10, 10, 10]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x[4] = self.x[7]
+        self.net = amin_net5
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_custom_vjp.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import _pir_ops, nn
+from paddle.autograd.ir_backward import grad
+from paddle.decomposition import decompose
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+class SimpNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, linear1_weight, linear2_weight):
+        x2 = _pir_ops.matmul(x, linear1_weight, False, False)
+        x3 = _pir_ops.gelu(x2, False)
+        res = _pir_ops.matmul(x3, linear2_weight, False, False)
+        return res
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2, 1024, 1024]
+        self.shape_y = [2, 1024, 1024]
+        self.shape_l1_w = [2, 1024, 4096]
+        self.shape_l2_w = [2, 4096, 1024]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.l1_w = np.random.random(self.shape_l1_w).astype("float32")
+        self.l2_w = np.random.random(self.shape_l2_w).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "backward":
+            core._set_prim_backward_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            net = SimpNet()
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            l1_w = paddle.static.data('l1_w', self.shape_l1_w, dtype='float32')
+            l2_w = paddle.static.data('l2_w', self.shape_l2_w, dtype='float32')
+            divide_out = paddle.divide(x, y)
+            res = net(divide_out, l1_w, l2_w)
+            [res2] = decompose(
+                main_program,
+                [res],
+            )
+            gradients = grad(res2, (x, y))
+
+            if flag == "backward":
+                whole_ops_before = [
+                    op.name() for op in main_program.global_block().ops
+                ]
+                assert (
+                    "pd_op.gelu" in whole_ops_before
+                    and "pd_op.gelu_grad" not in whole_ops_before
+                )
+                core._set_prim_forward_enabled(True)
+                [res2] = decompose(
+                    main_program, [res2], whitelist={"pd_op.gelu"}
+                )
+                whole_ops_after = [
+                    op.name() for op in main_program.global_block().ops
+                ]
+                assert "pd_op.gelu" not in whole_ops_after
+                core._set_prim_forward_enabled(False)
+
+            exe = paddle.static.Executor()
+            outs = exe.run(
+                feed={
+                    'x': self.x,
+                    'y': self.y,
+                    'l1_w': self.l1_w,
+                    'l2_w': self.l2_w,
+                },
+                fetch_list=[res2, gradients[0], gradients[1]],
+            )
+
+        if flag == "backward":
+            core._set_prim_backward_enabled(False)
+        return outs
+
+    def test_prim_custom_vjp(self):
+        res_ref = self.base_net()
+        res = self.base_net("backward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_dynamic.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.decomposition import decomp
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+def rms_norm(hidden_states, weight):
+    variance = hidden_states.pow(2).mean((0, 1), keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', [-1, -1, 4096], dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            res = rms_norm(x, y)
+            [res2] = decomp.decompose(main_program, [res])
+            if flag == "all":
+                # Todo(CZ): when symbolic shape rules of all op are ready, set flag to make this branch effective
+                pm = paddle.base.libpaddle.pir.PassManager()
+                paddle.base.libpaddle.pir.infer_symbolic_shape_pass(
+                    pm, main_program
+                )
+                pm.run(main_program)
+
+            exe = paddle.static.Executor()
+            outs = exe.run(
+                feed={
+                    'x': self.x,
+                    'y': self.y,
+                },
+                fetch_list=[res2],
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if not flag:
+            assert 'pd_op.mean' in whole_ops
+        if flag == "all":
+            core._set_prim_all_enabled(False)
+            assert 'pd_op.mean' not in whole_ops
+        return outs
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_flags_check_ops.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+def rms_norm(hidden_states, weight):
+    # From llama2, reduce dim is not equal to dynamic shape dim
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+class TestPrimMode1(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.net = rms_norm
+        self.enable_cinn = False
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=[1, 300, 4096], dtype='float32'),
+                    InputSpec(shape=[4096], dtype='float32'),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert "pd_op.mean" not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_jit.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+
+def func(x):
+    x1 = paddle.mean(x)
+    out = paddle.nn.functional.gelu(x1, False)
+    return out
+
+
+class TestDy2staticPir(unittest.TestCase):
+    def test_basic_network_backward(self):
+        core._set_prim_all_enabled(True)
+
+        # ==== dygraph computation ====
+        static_func = paddle.jit.to_static(func, full_graph=True)
+        x = paddle.randn((8, 16, 64))
+        x.stop_gradient = False
+        ref_out = func(x) * 2
+        ref_out.backward()
+        ref_grad = x.grad.numpy()
+        x.clear_gradient()
+
+        # ==== to static computation ====
+        out = static_func(x)
+        actual_out = out * 2
+        actual_out.backward()
+        actual_grad = x.grad
+
+        core._set_prim_all_enabled(False)
+        ops = [
+            op.name()
+            for op in static_func.program_cache.last()[-1][-1]
+            .train_program.program.global_block()
+            .ops
+        ]
+        assert "pd_op.erf" in ops
+        assert "pd_op.gelu" not in ops
+
+        np.testing.assert_allclose(
+            ref_out, actual_out.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+        np.testing.assert_allclose(
+            ref_grad, actual_grad.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+
+class TestDy2staticPirEval(unittest.TestCase):
+    def test_basic_network_backward_(self):
+        core._set_prim_all_enabled(True)
+
+        # ==== dygraph computation ====
+        static_func = paddle.jit.to_static(func, full_graph=True)
+        static_func.eval()
+        x = paddle.randn((8, 16, 64))
+        x.stop_gradient = False
+        ref_out = func(x) * 2
+
+        # ==== to static computation ====
+        out = static_func(x)
+        actual_out = out * 2
+
+        ops = [
+            op.name()
+            for op in static_func.program_cache.last()[-1][-1]
+            .infer_program.program.global_block()
+            .ops
+        ]
+        core._set_prim_all_enabled(False)
+
+        assert "pd_op.erf" in ops
+        assert "pd_op.gelu" not in ops
+
+        np.testing.assert_allclose(
+            ref_out, actual_out.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_jit_dynamic.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+def rms_norm1(hidden_states, weight):
+    # From llama2, reduce dim is not equal to dynamic shape dim
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+def rms_norm2(hidden_states, weight):
+    # reduce dim is not equal to dynamic shape dim
+    variance = hidden_states.pow(2).mean((0, 1), keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+class TestPrimMode1(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.net = rms_norm1
+        self.enable_cinn = False
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=[None, None, 4096], dtype='float32'),
+                    InputSpec(shape=[4096], dtype='float32'),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert "pd_op.mean" not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+class TestPrimMode2(TestPrimMode1):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.net = rms_norm2
+        self.enable_cinn = False
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_program.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.autograd.ir_backward import grad
+from paddle.decomposition import decomp
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.shape_y = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        elif flag == "backward":
+            core._set_prim_backward_enabled(True)
+        elif flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            divide_out = paddle.divide(x, y)
+            sum_out = paddle.mean(divide_out, axis=0)
+            [new_out] = decomp.decompose(main_program, [sum_out])
+            gradients = grad(new_out, (x, y))
+
+            exe = paddle.static.Executor()
+            [fwd, dx, dy] = exe.run(
+                feed={'x': self.x, 'y': self.y}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' in whole_ops
+            )
+        elif flag == "backward":
+            core._set_prim_backward_enabled(False)
+            assert (
+                'pd_op.mean' in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        elif flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        else:
+            assert (
+                'pd_op.mean' in whole_ops and 'pd_op.divide_grad' in whole_ops
+            )
+        return fwd, dx, dy
+
+    def test_prim_forward(self):
+        res_ref = self.base_net()
+        res = self.base_net("forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_equal(ref, actual)
+
+    def test_prim_backward(self):
+        res_ref = self.base_net()
+        res = self.base_net("backward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+    def test_prim_all(self):
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+class TestCompOpName(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.shape_y = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            divide_out = paddle.divide(x, y)
+            sum_out = paddle.mean(divide_out, axis=0)
+            [new_out] = decomp.decompose(main_program, [sum_out])
+            gradients = grad(new_out, (x, y))
+
+            exe = paddle.static.Executor()
+            [fwd, dx, dy] = exe.run(
+                feed={'x': self.x, 'y': self.y}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        else:
+            assert (
+                'pd_op.mean' in whole_ops and 'pd_op.divide_grad' in whole_ops
+            )
+        return main_program
+
+    def test_set_comp_op_name(self):
+        decomp_program = self.base_net("all")
+        for op in decomp_program.global_block().ops:
+            if op.name() == 'pd_op.sum':
+                assert op.attrs()['comp_op_name'] == 'pd_op.mean'
+            if op.name() == 'pd_op.expand':
+                assert op.attrs()['comp_op_name'] == 'pd_op.sum_grad'
+        origin_program = self.base_net()
+        for op in decomp_program.global_block().ops:
+            if op.name() == 'pd_op.mean':
+                assert not op.has_attr('comp_op_name')
+            if op.name() == 'pd_op.sum_grad':
+                assert not op.has_attr('comp_op_name')
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_rms_norm_st_shape.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+def rms_norm1(hidden_states, weight):
+    # From llama2, reduce dim is not equal to dynamic shape dim
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+def rms_norm2(hidden_states, weight):
+    # reduce dim is not equal to dynamic shape dim
+    variance = hidden_states.pow(2).mean((0, 1), keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+class TestPrimMode1(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.net = rms_norm1
+        self.enable_cinn = True
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=[1, 300, 4096], dtype='float32'),
+                    InputSpec(shape=[4096], dtype='float32'),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert "pd_op.mean" not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+class TestPrimMode2(TestPrimMode1):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.net = rms_norm2
+        self.enable_cinn = True
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_simpnet.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import _pir_ops, nn
+from paddle.autograd.ir_backward import grad
+from paddle.decomposition import decomp
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+class SimpNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, linear1_weight, linear2_weight):
+        x2 = _pir_ops.matmul(x, linear1_weight, False, False)
+        x3 = _pir_ops.gelu(x2, False)
+        res = _pir_ops.matmul(x3, linear2_weight, False, False)
+        return res
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2, 1024, 1024]
+        self.shape_y = [2, 1024, 1024]
+        self.shape_l1_w = [2, 1024, 4096]
+        self.shape_l2_w = [2, 4096, 1024]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.l1_w = np.random.random(self.shape_l1_w).astype("float32")
+        self.l2_w = np.random.random(self.shape_l2_w).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            net = SimpNet()
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            l1_w = paddle.static.data('l1_w', self.shape_l1_w, dtype='float32')
+            l2_w = paddle.static.data('l2_w', self.shape_l2_w, dtype='float32')
+            divide_out = paddle.divide(x, y)
+            res = net(divide_out, l1_w, l2_w)
+            [res2] = decomp.decompose(main_program, [res])
+            gradients = grad(res2, (x, y))
+            exe = paddle.static.Executor()
+            outs = exe.run(
+                feed={
+                    'x': self.x,
+                    'y': self.y,
+                    'l1_w': self.l1_w,
+                    'l2_w': self.l2_w,
+                },
+                fetch_list=[res2, gradients[0], gradients[1]],
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.gelu' not in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        return outs
+
+    def test_prim_all(self):
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_skip_dynamic.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.autograd.ir_backward import grad
+from paddle.decomposition import decomp
+from paddle.framework import core
+
+paddle.enable_static()
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2, 1024, 1024]
+        self.shape_y = [2, 1024, 1024]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', [-1, 1024, 1024], dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            x1 = paddle.nn.functional.relu(x)
+            y.stop_gradient = False
+            z = paddle.divide(x1, y)
+            res = paddle.nn.functional.gelu(z)
+            [res2] = decomp.decompose(main_program, [res])
+            gradients = grad(res2, (x, y))
+            exe = paddle.static.Executor()
+            outs = exe.run(
+                feed={
+                    'x': self.x,
+                    'y': self.y,
+                },
+                fetch_list=[res2, gradients[0], gradients[1]],
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "all":
+            core._set_prim_all_enabled(False)
+            assert 'pd_op.gelu' not in whole_ops
+        return outs
+
+    def test_prim_all_dynamic(self):
+        paddle.set_flags({"FLAGS_prim_skip_dynamic": True})
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py

@@ -0,0 +1,296 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def sum_net1(x):
+    return paddle.sum(x, axis=1, keepdim=False)
+
+
+def add_net(x, y):
+    return x + y
+
+
+def batch_norm_net1(x, y, z):
+    mean = paddle.zeros([40], dtype="float32")
+    var = paddle.ones([40], dtype='float32')
+    return paddle.nn.functional.batch_norm(x, mean, var, y, z)
+
+
+def reduce_as_net(x, y):
+    return paddle.reduce_as(x, y)
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    backend = "CINN" if use_cinn else None
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        backend=backend,
+        full_graph=True,
+    )
+
+
+class TestPrimBaseWithGrad(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.sum_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = sum_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_backward_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x)
+        res.backward()
+        x_grad = x.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_backward_enabled(False)
+        return res, x_grad
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+class TestPrimTwoWithGrad(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.add_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = add_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_backward_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        x_grad = x.gradient()
+        y_grad = y.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x, y)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_backward_enabled(False)
+        return res, [x_grad, y_grad]
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+class TestPrimBaseOneGradTwoInputs(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "reduce_as_grad"
+        self.dtype = "float32"
+        self.y_shape = [200, 40]
+        self.init_y_shape = [None, 200]
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = reduce_as_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_backward_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        if self.y_without_grad:
+            grad = x.gradient()
+        else:
+            grad = y.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x, y)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_backward_enabled(False)
+        return res, [grad]
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+class TestPrimThreeWithGrad(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.batch_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, None, None, 60]
+        self.y_shape = [40]
+        self.init_y_shape = [None]
+        self.z_shape = [40]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net1
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_backward_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y, stop_gradient=False)
+        z = paddle.to_tensor(self.z, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='float32'),
+                    InputSpec(shape=self.init_z_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y, z)
+        res.backward()
+        x_grad = x.gradient()
+        y_grad = y.gradient()
+        z_grad = z.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x, y, z)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_backward_enabled(False)
+        return res, [x_grad, y_grad, z_grad]
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_sub_graph_fghij_backward_dynamic_shape.py

@@ -0,0 +1,485 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from test_prim_sub_graph_backward_dynamic_shape import (
+    TestPrimBaseWithGrad,
+    TestPrimThreeWithGrad,
+    TestPrimTwoWithGrad,
+    apply_to_static,
+)
+
+import paddle
+from paddle.framework import core
+from paddle.static import InputSpec
+
+
+def floor_net(x):
+    return paddle.floor(x)
+
+
+def fmax_net(x, y):
+    return paddle.fmax(x, y)
+
+
+def fmin_net(x, y):
+    return paddle.fmin(x, y)
+
+
+def gather_net(x, y):
+    return paddle.gather(x, y, 1)
+
+
+def gather_nd_net(x, y):
+    return paddle.gather_nd(x, y)
+
+
+def gelu_net1(x):
+    return paddle.nn.functional.gelu(x, approximate=True)
+
+
+def gelu_net2(x):
+    return paddle.nn.functional.gelu(x, approximate=False)
+
+
+def group_norm_net1(x, y, z, epsilon=1e-5, num_groups=10):
+    return paddle._C_ops.group_norm(x, y, z, epsilon, num_groups, "NCHW")
+
+
+def group_norm_net2(x, epsilon=1e-5, num_groups=10):
+    return paddle._C_ops.group_norm(x, None, None, epsilon, num_groups, "NCHW")
+
+
+def group_norm_net3(x, y, z, epsilon=1e-5, num_groups=10):
+    return paddle._C_ops.group_norm(x, y, z, epsilon, num_groups, "NHWC")
+
+
+def group_norm_net4(x, epsilon=1e-5, num_groups=10):
+    return paddle._C_ops.group_norm(x, None, None, epsilon, num_groups, "NHWC")
+
+
+def hardsigmoid_net(x):
+    return paddle.nn.functional.hardsigmoid(x)
+
+
+def hardswish_net(x):
+    return paddle.nn.functional.hardswish(x)
+
+
+class TestPrimFloorWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.floor_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = floor_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFmaxWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmax_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.y[2, 10, 5:] = np.nan
+        self.net = fmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFmaxWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmax_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [30, 200, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.x[3, 100, 20:] = np.nan
+        self.net = fmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFmaxWithGrad3(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmax_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [30, 200, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.y[2, 10, 5:] = np.nan
+        self.net = fmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFmaxWithGrad4(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmax_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [30, 200, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.x[2, 10, 5:] = np.nan
+        self.y[2, 10, 5:] = np.nan
+        self.x[10, 9, :] = self.y[10, 9, :]
+        self.net = fmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFminWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmin_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.y[2, 10, 5:] = np.nan
+        self.net = fmin_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFminWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmin_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [30, 200, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.x[3, 100, 20:] = np.nan
+        self.net = fmin_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFminWithGrad3(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmin_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [30, 200, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.y[2, 10, 5:] = np.nan
+        self.net = fmin_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimFminWithGrad4(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.fmin_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [30, 200, 40]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.x[2, 10, 5:] = np.nan
+        self.y[2, 10, 5:] = np.nan
+        self.x[10, 9, :] = self.y[10, 9, :]
+        self.net = fmin_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimGatherWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.dtype = "float32"
+        self.x_shape = [10, 88, 10]
+        self.init_x_shape = [None, None, 10]
+        self.y_shape = [3]
+        self.init_y_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([1, 3, 5], dtype="int32")
+        self.net = gather_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_backward_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='int32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        x_grad = x.gradient()
+        if flag == "prim":
+            core._set_prim_backward_enabled(False)
+        return res, [x_grad]
+
+
+class TestPrimGatherWithGrad2(TestPrimGatherWithGrad1):
+    def setUp(self):
+        np.random.seed(2024)
+        self.dtype = "float32"
+        self.x_shape = [10, 88, 10]
+        self.init_x_shape = [None, 88, None]
+        self.y_shape = [3]
+        self.init_y_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([1, 3, 5], dtype="int32")
+        self.net = gather_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimGatherNdWithGrad(TestPrimGatherWithGrad1):
+    def setUp(self):
+        np.random.seed(2024)
+        self.dtype = "float32"
+        self.x_shape = [100, 100]
+        self.init_x_shape = [None, None, 10]
+        self.y_shape = [2, 2]
+        self.init_y_shape = [None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([[1, 1], [2, 1]], dtype="int32")
+        self.net = gather_nd_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimGeluWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.gelu_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = gelu_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimGeluWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.gelu_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = gelu_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimGeluWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.gelu_grad"
+        self.dtype = "float16"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.uniform(-1, 1, size=self.x_shape).astype(self.dtype)
+        self.net = gelu_net1
+        self.enable_cinn = False
+        self.rtol = 1e-5
+        self.atol = 0.0005
+
+    def test_prim_all_dynamic(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+        place = core.CUDAPlace(0)
+        if not core.is_float16_supported(place):
+            return
+
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.rtol, atol=self.atol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.rtol, atol=self.atol)
+
+
+class TestPrimGeluWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.gelu_grad"
+        self.dtype = "float16"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.uniform(-1, 1, size=self.x_shape).astype(self.dtype)
+        self.net = gelu_net2
+        self.enable_cinn = False
+        self.rtol = 1e-5
+        self.atol = 0.0005
+
+    def test_prim_all_dynamic(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+        place = core.CUDAPlace(0)
+        if not core.is_float16_supported(place):
+            return
+
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.rtol, atol=self.atol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.rtol, atol=self.atol)
+
+
+class TestPrimGroupNormWithGrad1(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.group_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 60, 50, 60]
+        self.init_x_shape = [None, None, None, 60]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.uniform(-0.5, 0.5, self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = group_norm_net1
+        self.enable_cinn = False
+        self.tol = 1e-3
+
+
+# class TestPrimGroupNormWithGrad2(TestPrimBaseWithGrad):
+#     def setUp(self):
+#         np.random.seed(2023)
+#         self.op_name = "pd_op.group_norm_grad"
+#         self.dtype = "float32"
+#         self.x_shape = [30, 60, 50, 60]
+#         self.init_x_shape = [None, 60, None, None]
+#         self.x = np.random.uniform(-0.5, 0.5, self.x_shape).astype(self.dtype)
+#         self.net = group_norm_net2
+#         self.enable_cinn = False
+#         self.tol = 1e-5
+
+
+class TestPrimGroupNormWithGrad3(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.group_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 60, 50, 60]
+        self.init_x_shape = [None, 60, None, None]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.uniform(-0.5, 0.5, self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = group_norm_net3
+        self.enable_cinn = False
+        self.tol = 1e-2
+
+
+# class TestPrimGroupNormWithGrad4(TestPrimBaseWithGrad):
+#     def setUp(self):
+#         np.random.seed(2023)
+#         self.op_name = "pd_op.group_norm_grad"
+#         self.dtype = "float32"
+#         self.x_shape = [30, 60, 50, 60]
+#         self.init_x_shape = [None, 60, None, None]
+#         self.x = np.random.uniform(-0.2, 0.2, self.x_shape).astype(self.dtype)
+#         self.net = group_norm_net4
+#         self.enable_cinn = False
+#         self.tol = 1e-2
+
+
+class TestPrimHardsigmoidWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.hardsigmoid_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = hardsigmoid_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimHardswishWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.hardswish_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = hardswish_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_sub_graph_uvwxyz_backward_dynamic_shape.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from test_prim_sub_graph_backward_dynamic_shape import (
+    TestPrimBaseWithGrad,
+    TestPrimTwoWithGrad,
+)
+
+import paddle
+
+
+def unsqueeze_net(x):
+    return paddle.unsqueeze(x, axis=[1, 2])
+
+
+def where_net(x, y):
+    return paddle.where(x > y, x, y)
+
+
+class TestPrimUnsqueezeWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.unsqueeze_grad"
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = unsqueeze_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimWhereWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.where_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimWhereWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.where_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [30, 30, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimWhereWithGrad3(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.where_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [30, 30, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_sink_decomp.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle.autograd.ir_backward import grad
+from paddle.base import core
+from paddle.decomposition import decompose
+
+paddle.enable_static()
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.shape_y = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+        self.prog = None
+
+    def base_net(self, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        elif flag == "backward":
+            core._set_prim_backward_enabled(True)
+        elif flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            divide_out = paddle.divide(x, y)
+            sum_out = paddle.mean(divide_out, axis=0)
+            [new_out] = decompose(main_program, [sum_out])
+            gradients = grad(new_out, (x, y))
+
+            exe = paddle.static.Executor()
+            [fwd, dx, dy] = exe.run(
+                feed={'x': self.x, 'y': self.y}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        self.prog = main_program
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' in whole_ops
+            )
+        elif flag == "backward":
+            core._set_prim_backward_enabled(False)
+            assert (
+                'pd_op.mean' in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        elif flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        else:
+            assert (
+                'pd_op.mean' in whole_ops and 'pd_op.divide_grad' in whole_ops
+            )
+        return fwd, dx, dy
+
+    def test_prim_forward(self):
+        res_ref = self.base_net()
+        res = self.base_net("forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_equal(ref, actual)
+
+    def test_prim_backward(self):
+        res_ref = self.base_net()
+        res = self.base_net("backward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+    def test_prim_all(self):
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+    def test_has_decomp(self):
+        _ = self.base_net()
+        for op in self.prog.global_block().ops:
+            if op.name() == "pd_op.divide":
+                self.assertEqual(core.has_decomp_rule(op), False)
+            if op.name() == "pd_op.mean":
+                self.assertEqual(core.has_decomp_rule(op), True)
+
+
+class TestReluSink(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.prog = None
+
+    def base_net(self, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            x.stop_gradient = False
+            sum_out = F.relu(x)
+            [new_out] = decompose(main_program, [sum_out])
+            gradients = grad(new_out, x)
+
+            exe = paddle.static.Executor()
+            [fwd, dx] = exe.run(
+                feed={'x': self.x}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        self.prog = main_program
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert 'pd_op.relu' not in whole_ops
+        else:
+            assert 'pd_op.relu' in whole_ops
+        return fwd, dx
+
+    def test_relu_forward(self):
+        res_ref = self.base_net()
+        res = self.base_net("forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_equal(ref, actual)
+
+
+class TestGeluSink(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.prog = None
+
+    def base_net(self, approximate=True, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            x.stop_gradient = False
+            sum_out = F.gelu(x, approximate=approximate)
+            [new_out] = decompose(main_program, [sum_out])
+            gradients = grad(new_out, x)
+
+            exe = paddle.static.Executor()
+            [fwd, dx] = exe.run(
+                feed={'x': self.x}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        self.prog = main_program
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert 'pd_op.gelu' not in whole_ops
+        else:
+            assert 'pd_op.gelu' in whole_ops
+        return fwd, dx
+
+    def test_gelu_forward_true(self):
+        res_ref = self.base_net(approximate=True)
+        res = self.base_net(approximate=True, flag="forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+    def test_gelu_approximate_false(self):
+        res_ref = self.base_net(approximate=False)
+        res = self.base_net(approximate=False, flag="forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+class TestHardSwishSink(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.prog = None
+
+    def base_net(self, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        elif flag == "backward":
+            core._set_prim_backward_enabled(True)
+        elif flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            x.stop_gradient = False
+            sum_out = F.hardswish(x)
+            [new_out] = decompose(main_program, [sum_out])
+            gradients = grad(new_out, x)
+
+            exe = paddle.static.Executor()
+            [fwd, dx] = exe.run(
+                feed={'x': self.x}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        self.prog = main_program
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert 'pd_op.hardswish' not in whole_ops
+        elif flag == "backward":
+            core._set_prim_backward_enabled(False)
+            assert (
+                'pd_op.hardswish' in whole_ops
+                and 'pd_op.hardswish_grad' not in whole_ops
+            )
+        elif flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.hardswish' not in whole_ops
+                and 'pd_op.hardswish_grad' not in whole_ops
+            )
+        else:
+            assert (
+                'pd_op.hardswish' in whole_ops
+                and 'pd_op.hardswish_grad' in whole_ops
+            )
+        return fwd, dx
+
+    def test_prim_forward(self):
+        res_ref = self.base_net()
+        res = self.base_net("forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-3, atol=1e-3)
+
+    def test_prim_backward(self):
+        res_ref = self.base_net()
+        res = self.base_net("backward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-3, atol=1e-3)
+
+    def test_prim_all(self):
+        res_ref = self.base_net()
+        res = self.base_net("all")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_vjp_prim.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.base.core import call_vjp
+
+paddle.enable_static()
+
+
+def get_ir_divide_program():
+    paddle.enable_static()
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.tensor.fill_constant(
+            shape=[1, 4], dtype='float32', value=2.0
+        )
+        x.stop_gradient = False
+        y = paddle.tensor.fill_constant(shape=[4], dtype='float32', value=1.0)
+        y.stop_gradient = False
+        dout = paddle.tensor.fill_constant(
+            shape=[1, 4], dtype='float32', value=1.0
+        )
+        dout.stop_gradient = False
+        out = paddle.divide(x, y)
+
+    return main_program
+
+
+def get_ir_sum_program():
+    paddle.enable_static()
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.tensor.fill_constant(
+            shape=[4, 5], dtype='float32', value=2.0
+        )
+        x.stop_gradient = False
+        dout = paddle.tensor.fill_constant(shape=[], dtype='float32', value=1.0)
+        dout.stop_gradient = False
+        out = paddle.sum(x)
+    return main_program
+
+
+class TestVjpPrim(unittest.TestCase):
+    def test_divide_grad_prim_case1(self):
+        pir_program = get_ir_divide_program()
+        paddle.framework.core._set_prim_backward_enabled(True)
+        with paddle.pir_utils.IrGuard():
+            dout = pir_program.global_block().ops[-2].result(0)
+            out_grads = [[dout]]
+            stop_gradients = [[False], [False]]
+            divide_op = pir_program.global_block().ops[-1]
+            with paddle.pir.core.program_guard(pir_program):
+                grad_outs = call_vjp(
+                    divide_op,
+                    [[value] for value in divide_op.operands_source()],
+                    [[value] for value in divide_op.results()],
+                    out_grads,
+                    stop_gradients,
+                )
+            print(pir_program)
+            reshape_op2 = pir_program.global_block().ops[-1]
+            reshape_op1 = pir_program.global_block().ops[-2]
+            self.assertEqual(len(grad_outs), 2)
+            self.assertEqual(len(pir_program.global_block().ops), 11)
+            self.assertTrue(reshape_op2.result(0).is_same(grad_outs[0][0]))
+            self.assertTrue(reshape_op1.result(0).is_same(grad_outs[1][0]))
+            paddle.framework.core._set_prim_backward_enabled(False)
+
+    def test_divide_grad_no_prim(self):
+        pir_program = get_ir_divide_program()
+        paddle.framework.core._set_prim_backward_enabled(False)
+        dout = pir_program.global_block().ops[-2].result(0)
+        out_grads = [[dout]]
+        stop_gradients = [[False], [False]]
+        divide_op = pir_program.global_block().ops[-1]
+        with paddle.pir.core.program_guard(pir_program):
+            grad_outs = call_vjp(
+                divide_op,
+                [[value] for value in divide_op.operands_source()],
+                [[value] for value in divide_op.results()],
+                out_grads,
+                stop_gradients,
+            )
+        self.assertEqual(len(grad_outs), 2)
+        self.assertEqual(
+            grad_outs[0][0].get_defining_op().name(), "pd_op.divide_grad"
+        )
+        self.assertEqual(
+            grad_outs[1][0].get_defining_op().name(), "pd_op.divide_grad"
+        )
+        self.assertEqual(len(pir_program.global_block().ops), 5)
+
+    def test_sum_grad_prim(self):
+        pir_program = get_ir_sum_program()
+        paddle.framework.core._set_prim_backward_enabled(True)
+        with paddle.pir_utils.IrGuard():
+            dout = pir_program.global_block().ops[-3].result(0)
+            out_grads = [[dout]]
+            stop_gradients = [[False]]
+            sum_op = pir_program.global_block().ops[-1]
+            with paddle.pir.core.program_guard(pir_program):
+                grad_outs = call_vjp(
+                    sum_op,
+                    [[value] for value in sum_op.operands_source()],
+                    [[value] for value in sum_op.results()],
+                    out_grads,
+                    stop_gradients,
+                )
+            expand_op = pir_program.global_block().ops[-1]
+            self.assertEqual(len(grad_outs), 1)
+            self.assertEqual(len(pir_program.global_block().ops), 8)
+            self.assertTrue(expand_op.result(0).is_same(grad_outs[0][0]))
+            all_op_names = [
+                "pd_op.full",
+                "pd_op.full",
+                "pd_op.full_int_array",
+                "pd_op.sum",
+                "pd_op.full_int_array",
+                "pd_op.reshape",
+                "pd_op.full_int_array",
+                "pd_op.expand",
+            ]
+            for idx, op in enumerate(pir_program.global_block().ops):
+                self.assertEqual(op.name(), all_op_names[idx])
+            paddle.framework.core._set_prim_backward_enabled(False)
+
+    def test_sum_grad_no_prim(self):
+        pir_program = get_ir_sum_program()
+        paddle.framework.core._set_prim_backward_enabled(False)
+        dout = pir_program.global_block().ops[-2].result(0)
+        out_grads = [[dout]]
+        stop_gradients = [[False]]
+        sum_op = pir_program.global_block().ops[-1]
+        with paddle.pir.core.program_guard(pir_program):
+            grad_outs = call_vjp(
+                sum_op,
+                [[value] for value in sum_op.operands_source()],
+                [[value] for value in sum_op.results()],
+                out_grads,
+                stop_gradients,
+            )
+        self.assertEqual(len(grad_outs), 1)
+        self.assertEqual(
+            grad_outs[0][0].get_defining_op().name(), "pd_op.sum_grad"
+        )
+        self.assertEqual(len(pir_program.global_block().ops), 5)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/prim/CMakeLists.txt

@@ -0,0 +1,11 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()
+
+add_subdirectory(vjp)

test/prim/prim/vjp/CMakeLists.txt

@@ -0,0 +1,13 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()
+
+set_tests_properties(test_comp_high_grad PROPERTIES TIMEOUT 100)
+
+add_subdirectory(eager)

test/prim/prim/vjp/eager/CMakeLists.txt

@@ -0,0 +1,12 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+set(GC_ENVS FLAGS_eager_delete_tensor_gb=0.0)
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()
+
+set_tests_properties(test_comp_eager_cumprod_grad PROPERTIES TIMEOUT 120)

test/prim/prim/vjp/eager/test_comp_eager_acos_double_grad.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.uniform(-0.95, 0.95, size=(10, 10)),
+            np.random.rand(10, 10),
+            np.float32,
+        ),
+    ],
+)
+class TestAcosDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_acos_double_grad_comp_dygraph(self):
+        paddle.disable_static()
+        core.set_prim_eager_enabled(True)
+
+        x = paddle.to_tensor(self.primal, dtype='float32', stop_gradient=False)
+        x.stop_gradient = False
+        y = paddle.acos(x)
+        dx = paddle.grad(y, x, create_graph=True, retain_graph=True)[0]
+        ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)[0]
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_add_grad.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal0', 'primal1', 'dtype'),
+    [
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(3, 1, 1),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 1),
+            np.float32,
+        ),
+    ],
+)
+class TestAddGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.dtype)
+        cls.primal1 = cls.primal1.astype(cls.dtype)
+
+    def test_add_grad_comp(self):
+        def actual(primal0, primal1):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.add(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        def desired(primal0, primal1):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.add(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        dx, dy = actual(self.primal0, self.primal1)
+
+        ddx, ddy = desired(self.primal0, self.primal1)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        np.testing.assert_allclose(
+            actual=dy,
+            desired=ddy,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_batch_norm_grad.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle.base import core
+
+np.random.seed(2023)
+
+
+class Arg:
+    dout = None
+
+
+def generate_data(shape, dtype="float32"):
+    np_data = np.random.random(shape).astype(dtype)
+    return np_data
+
+
+class Attr:
+    def __init__(self) -> None:
+        self.dtype = "float32"
+        self.shape = [8, 8, 16, 16]
+        self.training = True
+        self.momentum = 0.9
+        self.epsilon = 1e-05
+        self.data_format = "NCHW"
+        self.use_global_stats = None
+
+    def set_dtype(self, dtype) -> None:
+        self.dtype = dtype
+
+    def set_shape(self, shape) -> None:
+        self.shape = shape
+
+    def set_training(self, training) -> None:
+        self.training = training
+
+    def set_momentum(self, momentum) -> None:
+        self.momentum = momentum
+
+    def set_epsilon(self, epsilon) -> None:
+        self.epsilon = epsilon
+
+    def set_data_format(self, data_format) -> None:
+        self.data_format = data_format
+
+    def set_use_global_stats(self, use_global_stats) -> None:
+        self.use_global_stats = use_global_stats
+
+
+attrs = Attr()
+
+
+def fn(
+    x,
+    running_mean,
+    running_variance,
+    weight,
+    bias,
+    training,
+    momentum,
+    epsilon,
+    data_format,
+    use_global_stats,
+):
+    z = F.batch_norm(
+        x,
+        running_mean,
+        running_variance,
+        weight,
+        bias,
+        training=training,
+        momentum=momentum,
+        epsilon=epsilon,
+        data_format=data_format,
+        use_global_stats=use_global_stats,
+    )
+    out = z * paddle.to_tensor(Arg.dout)
+    res = paddle.mean(out)
+    return res
+
+
+def expect_grad(
+    x,
+    running_mean,
+    running_variance,
+    weight,
+    bias,
+    training,
+    momentum,
+    epsilon,
+    data_format,
+    use_global_stats,
+):
+    x.stop_gradient = False
+    res = fn(
+        x,
+        running_mean,
+        running_variance,
+        weight,
+        bias,
+        training,
+        momentum,
+        epsilon,
+        data_format,
+        use_global_stats,
+    )
+    gradients = paddle.grad(res, x)
+    return gradients
+
+
+def cal_composite(inputs, running_mean, running_variance, weight, bias):
+    paddle.enable_static()
+
+    startup_program = paddle.static.Program()
+    main_program = paddle.static.Program()
+    with paddle.static.program_guard(main_program, startup_program):
+        x1 = paddle.static.data(
+            'x1', shape=inputs.shape, dtype=str(inputs.dtype)
+        )
+        x1.stop_gradient = False
+        x2 = paddle.static.data(
+            'x2', shape=running_mean.shape, dtype=str(running_mean.dtype)
+        )
+        x3 = paddle.static.data(
+            'x3',
+            shape=running_variance.shape,
+            dtype=str(running_variance.dtype),
+        )
+        x4 = paddle.static.data(
+            'x4', shape=weight.shape, dtype=str(weight.dtype)
+        )
+        x5 = paddle.static.data('x5', shape=bias.shape, dtype=str(bias.dtype))
+        y = fn(
+            x1,
+            x2,
+            x3,
+            x4,
+            x5,
+            attrs.training,
+            attrs.momentum,
+            attrs.epsilon,
+            attrs.data_format,
+            attrs.use_global_stats,
+        )
+        z = paddle.static.gradients([y], [x1])
+
+    exe = paddle.static.Executor()
+    exe.run(startup_program)
+    res = exe.run(
+        main_program,
+        feed={
+            'x1': inputs,
+            'x2': running_mean,
+            'x3': running_variance,
+            'x4': weight,
+            'x5': bias,
+        },
+        fetch_list=[z],
+    )
+    paddle.disable_static()
+    return res
+
+
+class TestCompositeBatchNorm(unittest.TestCase):
+    def setUp(self):
+        self.dtypes = ["float32"]
+        self.training = [False, True]
+        self.shapes = [[8, 8, 16, 16], [2, 1, 2, 3]]
+        self.momentum = [0.1, 0.9]
+        self.epsilon = [1e-05, 2e-05]
+        self.data_formats = ["NCHW"]
+        self.use_global_stats = [None, True, False]
+
+    def compare_backward(self):
+        if attrs.training is True and attrs.use_global_stats is False:
+            # in this case, origin bn grad kernel is not the same as forward kernel.
+            return
+        np_data = generate_data(attrs.shape, attrs.dtype)
+        tensor_data = paddle.to_tensor(np_data)
+        Arg.dout = np.random.random(np_data.shape).astype(attrs.dtype)
+        C = np_data.shape[1]
+
+        running_mean = paddle.zeros(C, dtype=attrs.dtype)
+        running_variance = paddle.ones(C, dtype=attrs.dtype)
+        weight = paddle.ones(C, dtype=attrs.dtype) * 2
+        bias = paddle.ones(C, dtype=attrs.dtype)
+
+        expect = expect_grad(
+            tensor_data,
+            running_mean,
+            running_variance,
+            weight,
+            bias,
+            attrs.training,
+            attrs.momentum,
+            attrs.epsilon,
+            attrs.data_format,
+            attrs.use_global_stats,
+        )[0].numpy()
+        np_running_mean = np.zeros(C, dtype=attrs.dtype)
+        np_running_variance = np.ones(C, dtype=attrs.dtype)
+        np_weight = np.ones(C, dtype=attrs.dtype) * 2
+        np_bias = np.ones(C, dtype=attrs.dtype)
+
+        actual = cal_composite(
+            np_data, np_running_mean, np_running_variance, np_weight, np_bias
+        )[0]
+
+        assert expect.dtype == actual.dtype
+
+        np.testing.assert_allclose(
+            expect,
+            actual,
+            rtol=1e-5,
+            atol=1e-5,
+        )
+
+    def test_backward_prim_dygraph_vjp(self):
+        core.set_prim_eager_enabled(True)
+        for i in self.training:
+            for j in self.dtypes:
+                for m in self.momentum:
+                    attrs.set_training(i)
+                    attrs.set_dtype(j)
+                    attrs.set_momentum(m)
+                    self.compare_backward()
+
+        for n in self.shapes:
+            for t in self.use_global_stats:
+                attrs.set_shape(n)
+                attrs.set_use_global_stats(t)
+                self.compare_backward()
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_cast_grad.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'src_dtype', 'dst_type'),
+    [
+        (
+            np.random.rand(10, 10),
+            np.random.rand(10, 10),
+            np.float32,
+            np.float64,
+        ),
+        (
+            np.random.rand(10, 10),
+            np.random.rand(10, 10),
+            np.float64,
+            np.float32,
+        ),
+        (
+            np.random.rand(10, 10),
+            np.random.rand(10, 10),
+            np.float32,
+            np.float32,
+        ),
+    ],
+)
+class TestCastGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.src_dtype)
+        cls.cotangent = cls.cotangent.astype(cls.src_dtype)
+
+    def test_cast_grad_comp(self):
+        core.set_prim_eager_enabled(True)
+
+        def actual(primal, cotangent):
+            x = paddle.to_tensor(primal)
+            x.stop_gradient = False
+            v = paddle.to_tensor(cotangent)
+            y = paddle.cast(x, self.dst_type)
+            x_cotangent = paddle.grad(y, x, v)
+            return x_cotangent
+
+        def desired(primal, cotangent):
+            return (cotangent * np.ones_like(primal)).astype(primal.dtype)
+
+        actual = actual(self.primal, self.cotangent)
+        desired = desired(self.primal, self.cotangent)
+        from paddle.base.data_feeder import vartype_to_str
+        from paddle.pir.core import datatype_to_str
+
+        if actual[0].dtype in datatype_to_str:
+            TO_NUMPY_DTYPE = datatype_to_str
+        else:
+            TO_NUMPY_DTYPE = vartype_to_str
+
+        self.assertEqual(TO_NUMPY_DTYPE[actual[0].dtype], desired.dtype)
+        np.testing.assert_allclose(
+            actual=actual[0],
+            desired=desired,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_cos_double_grad.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestCosDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_cos_double_grad_comp_dygraph(self):
+        def actual(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            core._set_prim_backward_blacklist("cos_grad")
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cos(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        def desired(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cos(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal),
+            desired=desired(self.primal),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_cumprod_grad.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'dtype'),
+    [
+        (
+            np.random.uniform(1, 5, (50,)),
+            np.float32,
+        ),
+        (
+            np.random.rand(10, 10),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4, 5),
+            np.float32,
+        ),
+        (
+            np.random.randint(1, 100, (2, 3, 4)),
+            np.int64,
+        ),
+    ],
+)
+class TestCumprodGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        cls.zero_nums = [0, 1, 10, int(np.prod(cls.primal.shape))]
+
+    def test_cumprod_grad_comp(self):
+        def actual(primal, dim):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cumprod(x, dim=dim)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=True, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        def desired(primal, dim):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cumprod(x, dim=dim)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=False, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        for zero_num in self.zero_nums:
+            shape = self.primal.shape
+            x = self.primal.flatten()
+            indices = random.sample(range(x.size), zero_num)
+            for i in indices:
+                x[i] = 0
+            x = np.reshape(x, shape)
+            for i in range(len(self.primal.shape)):
+                np.testing.assert_allclose(
+                    actual=actual(x, i),
+                    desired=desired(x, i),
+                    rtol=1e-6,
+                    atol=0,
+                )
+        core.set_prim_eager_enabled(False)
+
+
+@param.parameterized_class(
+    ('primal', 'dtype'),
+    [
+        (
+            np.random.uniform(1, 5, ()),
+            np.float32,
+        ),
+    ],
+)
+class TestCumprodGradComp0D(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_cumprod_grad_comp_0d(self):
+        def actual(primal, dim):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cumprod(x, dim=dim)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=True, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        def desired(primal, dim):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cumprod(x, dim=dim)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=False, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, 0),
+            desired=desired(self.primal, 0),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_div_grad.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal0', 'primal1', 'dtype'),
+    [
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 1),
+            np.float32,
+        ),
+    ],
+)
+class TestDivGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.dtype)
+        cls.primal1 = cls.primal1.astype(cls.dtype)
+
+    def test_div_grad_comp(self):
+        def actual(primal0, primal1):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.divide(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        def desired(primal0, primal1):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.divide(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        dx, dy = actual(self.primal0, self.primal1)
+
+        ddx, ddy = desired(self.primal0, self.primal1)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        np.testing.assert_allclose(
+            actual=dy,
+            desired=ddy,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_exp_grad.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import autograd
+import autograd.numpy
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestExpGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core.set_prim_eager_enabled(True)
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    @classmethod
+    def tearDownClass(cls):
+        core.set_prim_eager_enabled(False)
+
+    def test_exp_grad_comp(self):
+        def actual(primal, cotangent):
+            primal = paddle.to_tensor(primal)
+            primal.stop_gradient = False
+            return paddle.grad(
+                paddle.exp(primal), primal, paddle.to_tensor(cotangent)
+            )[0]
+
+        def desired(primal, cotangent):
+            cotangent = (
+                np.ones_like(cotangent, dtype=primal.dtype)
+                if cotangent is None
+                else cotangent
+            )
+            return autograd.make_vjp(autograd.numpy.exp)(primal)[0](cotangent)
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent),
+            desired=desired(self.primal, self.cotangent),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_stop_gradients(self):
+        with self.assertRaises(ValueError):
+            primal = paddle.to_tensor(self.primal)
+            primal.stop_gradient = True
+            return paddle.grad(
+                paddle.exp(primal), primal, paddle.to_tensor(self.cotangent)
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_expand_grad.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('name', 'primal', 'cotangent', 'shape', 'dtype'),
+    (
+        (
+            'same_shape',
+            np.random.rand(10, 10),
+            np.random.rand(10, 10),
+            (10, 10),
+            np.float32,
+        ),
+        (
+            'same_rank',
+            np.random.rand(1, 10),
+            np.random.rand(10, 10),
+            (10, 10),
+            np.float32,
+        ),
+        (
+            'same_rank',
+            np.random.rand(10, 1, 10, 1),
+            np.random.rand(10, 10, 10, 10),
+            (10, 10, 10, 10),
+            np.float32,
+        ),
+        (
+            'diff_rank',
+            np.random.rand(1, 10, 1),
+            np.random.rand(10, 10, 10, 10),
+            (10, 10, 10, 10),
+            np.float32,
+        ),
+    ),
+)
+class TestExpandGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    @classmethod
+    def tearDownClass(cls):
+        core.set_prim_eager_enabled(False)
+
+    def test_comp(self):
+        def func(primal, cotangent, shape):
+            primal = paddle.to_tensor(primal)
+            primal.stop_gradient = False
+            cotangent = paddle.to_tensor(cotangent)
+            return paddle.grad(paddle.expand(primal, shape), primal, cotangent)[
+                0
+            ]
+
+        def actual(primal, cotangent, shape):
+            core.set_prim_eager_enabled(True)
+            return func(primal, cotangent, shape)
+
+        def desired(primal, cotangent, shape):
+            core.set_prim_eager_enabled(False)
+            return func(primal, cotangent, shape)
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, self.shape),
+            desired=desired(self.primal, self.cotangent, self.shape),
+            rtol=1e-6,
+            atol=0,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_gather_grad.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal0', 'index', 'axis', 'x_dtype', 'index_dtype', 'v'),
+    [
+        (
+            np.random.rand(100),
+            np.array([1, 3, 5]),
+            0,
+            np.float32,
+            np.int32,
+            np.random.rand(3),
+        ),
+        (
+            np.random.rand(10, 20),
+            np.array([1, 3, 5]),
+            0,
+            np.float64,
+            np.int64,
+            np.random.rand(3, 20),
+        ),
+        (
+            np.random.rand(10, 20),
+            np.array([1, 1, 3]),
+            0,
+            np.float32,
+            np.int32,
+            np.random.rand(3, 20),
+        ),
+        (
+            np.random.rand(3, 88, 30),
+            np.array([1, 3, 5]),
+            1,
+            np.float32,
+            np.int32,
+            np.random.rand(3, 3, 30),
+        ),
+        (
+            np.random.rand(10, 88, 10),
+            np.array([1, 3, 5]),
+            0,
+            np.float32,
+            np.int32,
+            np.random.rand(3, 88, 10),
+        ),
+    ],
+)
+class TestGatherGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.x_dtype)
+        cls.index = cls.index.astype(cls.index_dtype)
+        cls.v = cls.v.astype(cls.x_dtype)
+
+    @classmethod
+    def tearDownClass(cls):
+        core.set_prim_eager_enabled(False)
+
+    def test_exp_grad_comp(self):
+        def actual(primal0, index, axis):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(
+                primal0, dtype=primal0.dtype, stop_gradient=False
+            )
+            index = paddle.to_tensor(index, dtype=index.dtype)
+            x.stop_gradient = False
+            index.stop_gradient = True
+            out = paddle.gather(x, index, axis)
+            res = paddle.grad(out, [x], create_graph=False, retain_graph=True)
+            return res[0].numpy()
+
+        def desired(primal0, index, axis):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(
+                primal0, dtype=primal0.dtype, stop_gradient=False
+            )
+            index = paddle.to_tensor(index, dtype=index.dtype)
+            x.stop_gradient = False
+            index.stop_gradient = True
+            out = paddle.gather(x, index, axis)
+            res = paddle.grad(out, [x], create_graph=False, retain_graph=True)
+            return res[0].numpy()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal0, self.index, self.axis),
+            desired=desired(self.primal0, self.index, self.axis),
+            rtol=1e-6,
+            atol=0,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_index_add_double_grad.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+
+
+@param.parameterized_class(
+    ('x', 'index', 'axis', 'value', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.randn(4, 3, 2),  # x
+            np.random.randint(-3, 3, size=(16,)),  # index
+            1,  # axis
+            np.random.randint(0, 3, size=(4, 16, 2)),  # value
+            np.random.rand(4, 3, 2),  # cotangent
+            np.float32,  # dtype
+        ),
+    ],
+)
+class TestTakeAlongAxisTanhDoubleGrad(unittest.TestCase):
+    def test_index_add_tanh_double_grad(self):
+        x_tensor = paddle.to_tensor(
+            self.x, dtype=self.dtype, stop_gradient=False
+        )
+        value_tensor = paddle.to_tensor(
+            self.value, dtype=self.dtype, stop_gradient=False
+        )
+        index_tensor = paddle.to_tensor(self.index, dtype="int64")
+        dout_tensor = paddle.to_tensor(
+            self.cotangent, dtype=self.dtype, stop_gradient=False
+        )
+        out = paddle.index_add(x_tensor, index_tensor, self.axis, value_tensor)
+
+        out = paddle.tanh(out)
+
+        dx = paddle.grad(out, x_tensor, dout_tensor, create_graph=True)[0]
+
+        ddx = paddle.grad(dx, dout_tensor, create_graph=True)[0]
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_matmul_double_grad.py

@@ -0,0 +1,556 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+# vector * vector out.shape = (1)
+# matrix * vector out.shape = (2)
+# vector * matrix out.shape = (3)
+# batched matrix * batched matrix 4 for trans out.shape = (2, 3, 5)
+# batched matrix * broadcasted vector out.shape = (2, 3)
+# batched matrix * broadcasted matrix out.shape = (2, 3, 5, 4)
+
+TOLERANCE = {
+    "float16": {"rtol": 1e-3, "atol": 1e-3},
+    "float32": {"rtol": 1e-6, "atol": 1e-6},
+    "float64": {"rtol": 1e-15, "atol": 1e-15},
+}
+
+
+@param.parameterized_class(
+    ('primal0', 'primal1', 'trans_0', 'trans_1', 'dtype'),
+    [
+        (
+            np.random.rand(2),
+            np.random.rand(2),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3),
+            np.random.rand(3),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(2, 3),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(3, 2),
+            False,
+            True,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 4, 5),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 4, 5),
+            True,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 5, 4),
+            False,
+            True,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 5, 4),
+            True,
+            True,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 1, 5, 2),
+            np.random.rand(1, 3, 2, 4),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(2),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 3),
+            np.random.rand(3),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(2, 3),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(3, 2),
+            False,
+            True,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 4, 5),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 4, 5),
+            True,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 5, 4),
+            False,
+            True,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 5, 4),
+            True,
+            True,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 1, 5, 2),
+            np.random.rand(1, 3, 2, 4),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(2),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 3),
+            np.random.rand(3),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(2, 3),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2),
+            np.random.rand(3, 2),
+            False,
+            True,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 5, 4),
+            False,
+            True,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 4, 5),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 4, 5),
+            True,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 5, 4),
+            True,
+            True,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 1, 5, 2),
+            np.random.rand(1, 3, 2, 4),
+            False,
+            False,
+            np.float64,
+        ),
+    ],
+)
+class TestMatmulDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.dtype)
+        cls.primal1 = cls.primal1.astype(cls.dtype)
+        cls.trans_0 = cls.trans_0
+        cls.trans_1 = cls.trans_1
+
+    def setUp(self):
+        paddle.enable_static()
+
+    def tearDown(self):
+        paddle.disable_static()
+
+    def test_matmul_grad_comp(self):
+        def actual(primal0, primal1, trans_0, trans_1, dtype_):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype=dtype_, stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype=dtype_, stop_gradient=False)
+            out = paddle.matmul(x, y, trans_0, trans_1)
+            dout = paddle.ones_like(out, dtype=dtype_)
+            dout.stop_gradient = False
+            res = paddle.grad(
+                [out], [x, y], dout, create_graph=True, retain_graph=True
+            )
+            res_double = paddle.grad(
+                res, [x, y, dout], create_graph=True, retain_graph=True
+            )
+            return (
+                res_double[0].numpy(),
+                res_double[1].numpy(),
+                res_double[2].numpy(),
+            )
+
+        def desired(primal0, primal1, trans_0, trans_1, dtype_):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype=dtype_, stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype=dtype_, stop_gradient=False)
+            out = paddle.matmul(x, y, trans_0, trans_1)
+            dout = paddle.ones_like(out, dtype=dtype_)
+            dout.stop_gradient = False
+            res = paddle.grad(
+                out, [x, y], dout, create_graph=True, retain_graph=True
+            )
+            res_double = paddle.grad(
+                res, [x, y, dout], create_graph=True, retain_graph=True
+            )
+            return (
+                res_double[0].numpy(),
+                res_double[1].numpy(),
+                res_double[2].numpy(),
+            )
+
+        d_type = "float32"
+        if self.primal0.dtype == np.float16:
+            d_type = "float16"
+        elif self.primal0.dtype == np.float64:
+            d_type = "float64"
+
+        if paddle.device.get_device() == "cpu" and d_type == "float16":
+            # matmul fp16 cpu not supposed
+            pass
+        else:
+            dx, dy, ddout = actual(
+                self.primal0, self.primal1, self.trans_0, self.trans_1, d_type
+            )
+
+            dx_, dy_, ddout_ = desired(
+                self.primal0, self.primal1, self.trans_0, self.trans_1, d_type
+            )
+
+            np.testing.assert_allclose(
+                actual=dx,
+                desired=dx_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=dy,
+                desired=dy_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=ddout,
+                desired=ddout_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+
+
+@param.parameterized_class(
+    ('primal0', 'primal1', 'trans_0', 'trans_1', 'dtype'),
+    [
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(4),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 2, 3),
+            np.random.rand(2, 3, 2),
+            False,
+            False,
+            np.float16,
+        ),
+        (
+            np.random.rand(2, 2, 3),
+            np.random.rand(2, 3, 2),
+            False,
+            False,
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 2, 3),
+            np.random.rand(2, 3, 2),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 4, 3),
+            np.random.rand(2, 5, 4),
+            True,
+            True,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 2, 3),
+            np.random.rand(1, 3, 2),
+            False,
+            False,
+            np.float64,
+        ),
+        (
+            np.random.rand(2, 1, 5, 2),
+            np.random.rand(1, 3, 2, 4),
+            False,
+            False,
+            np.float32,
+        ),
+    ],
+)
+class TestMatmulTripleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.dtype)
+        cls.primal1 = cls.primal1.astype(cls.dtype)
+        cls.trans_0 = cls.trans_0
+        cls.trans_1 = cls.trans_1
+
+    def setUp(self):
+        paddle.enable_static()
+
+    def tearDown(self):
+        paddle.disable_static()
+
+    def test_matmul_grad_comp(self):
+        def actual(primal0, primal1, trans_0, trans_1, dtype_):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype=dtype_, stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype=dtype_, stop_gradient=False)
+            out = paddle.matmul(x, y, trans_0, trans_1)
+            dout = paddle.ones_like(out, dtype=dtype_)
+            dout.stop_gradient = False
+            ddx = paddle.ones_like(x, dtype=dtype_)
+            ddx.stop_gradient = False
+            ddy = paddle.ones_like(y, dtype=dtype_)
+            ddy.stop_gradient = False
+            res = paddle.grad(
+                [out], [x, y], dout, create_graph=True, retain_graph=True
+            )
+            res_double = paddle.grad(
+                res,
+                [x, y, dout],
+                [ddx, ddy],
+                create_graph=True,
+                retain_graph=True,
+            )
+
+            res_triple = paddle.grad(
+                res_double,
+                [x, y, dout, ddx, ddy],
+                create_graph=False,
+                retain_graph=False,
+            )
+            return (
+                res_double[0].numpy(),
+                res_double[1].numpy(),
+                res_double[2].numpy(),
+                res_triple[0].numpy(),
+                res_triple[1].numpy(),
+            )
+
+        def desired(primal0, primal1, trans_0, trans_1, dtype_):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype=dtype_, stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype=dtype_, stop_gradient=False)
+            out = paddle.matmul(x, y, trans_0, trans_1)
+            dout = paddle.ones_like(out, dtype=dtype_)
+            dout.stop_gradient = False
+            ddx = paddle.ones_like(x, dtype=dtype_)
+            ddx.stop_gradient = False
+            ddy = paddle.ones_like(y, dtype=dtype_)
+            ddy.stop_gradient = False
+            res = paddle.grad(
+                [out], [x, y], [dout], create_graph=True, retain_graph=True
+            )
+            res_double = paddle.grad(
+                res,
+                [x, y, dout],
+                [ddx, ddy],
+                create_graph=True,
+                retain_graph=True,
+            )
+            res_triple = paddle.grad(
+                res_double,
+                [x, y, dout, ddx, ddy],
+                create_graph=False,
+                retain_graph=True,
+            )
+            return (
+                res_double[0].numpy(),
+                res_double[1].numpy(),
+                res_double[2].numpy(),
+                res_triple[0].numpy(),
+                res_triple[1].numpy(),
+            )
+
+        d_type = "float32"
+        if self.primal0.dtype == np.float16:
+            d_type = "float16"
+        elif self.primal0.dtype == np.float64:
+            d_type = "float64"
+
+        if paddle.device.get_device() == "cpu" and d_type == "float16":
+            # matmul fp16 cpu not supposed
+            pass
+        else:
+            dx, dy, ddout, dx2, dy2 = actual(
+                self.primal0, self.primal1, self.trans_0, self.trans_1, d_type
+            )
+
+            dx_, dy_, ddout_, dx2_, dy2_ = desired(
+                self.primal0, self.primal1, self.trans_0, self.trans_1, d_type
+            )
+
+            np.testing.assert_allclose(
+                actual=dx,
+                desired=dx_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=dy,
+                desired=dy_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=ddout,
+                desired=ddout_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=dx2,
+                desired=dx2_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+            np.testing.assert_allclose(
+                actual=dy2,
+                desired=dy2_,
+                rtol=TOLERANCE[d_type]['rtol'],
+                atol=TOLERANCE[d_type]['atol'],
+            )
+
+    core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_min_grad.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'axis', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(16, 32), [1], np.random.rand(16, 32), np.float32),
+        (np.random.rand(16, 32), [0], np.random.rand(16, 32), np.float32),
+    ],
+)
+class TestMinGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_min_grad_comp(self):
+        def actual(primal0, axis):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.min(x, axis)
+            res = paddle.grad(out, [x], create_graph=False)
+            return res[0].numpy()
+
+        def desired(primal0, axis):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.min(x, axis)
+            res = paddle.grad(out, [x], create_graph=False)
+            return res[0].numpy()
+
+        dx = actual(self.primal, self.axis)
+
+        ddx = desired(self.primal, self.axis)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_multiply_grad.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('name', 'primals', 'stop_gradients', 'cotangents', 'dtype'),
+    (
+        (
+            'test_normal_case',
+            (np.random.rand(2, 3, 4), np.random.rand(2, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_broadcast_diff_rank',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(3, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_broadcast_same_rank',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(2, 1, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_stop_gradient',
+            (np.random.rand(2, 3, 1, 4), np.random.rand(2, 1, 3, 4)),
+            (False, True),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+        (
+            'test_reduce_axe_empty',
+            (np.random.rand(2, 3, 3, 4), np.random.rand(2, 1, 3, 4)),
+            (False, False),
+            (np.random.rand(2, 3, 3, 4),),
+            np.float32,
+        ),
+    ),
+)
+class TestMultiplyGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primals = tuple(primal.astype(cls.dtype) for primal in cls.primals)
+        cls.cotangents = tuple(co.astype(cls.dtype) for co in cls.cotangents)
+
+    def as_tuple(self, x):
+        return (x,) if isinstance(x, paddle.Tensor) else x
+
+    def vjp(self):
+        primals, cotangents = self.primals, self.cotangents
+        primals = tuple(paddle.to_tensor(primal) for primal in primals)
+        for primal, flag in zip(primals, self.stop_gradients):
+            primal.stop_gradient = flag
+        cotangents = tuple(paddle.to_tensor(co) for co in cotangents)
+        out = self.as_tuple(paddle.multiply(*primals))
+        grads = paddle.grad(out, primals, cotangents, allow_unused=True)
+        return [g for g in grads if g is not None]
+
+    def test_comp(self):
+        core.set_prim_eager_enabled(True)
+        actual = self.vjp()
+
+        core.set_prim_eager_enabled(False)
+        desired = self.vjp()
+
+        for i, j in zip(actual, desired):
+            np.testing.assert_allclose(
+                i,
+                j,
+                rtol=1e-6,
+                atol=0,
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_p_norm_grad.py

@@ -0,0 +1,216 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('shape', 'porder', 'axis', 'keepdim', 'dtype'),
+    [
+        [[], -float("inf"), -1, True, "float32"],
+        [[], -float("inf"), -1, False, "float32"],
+        [[], -float("inf"), 0, True, "float32"],
+        [[], -float("inf"), 0, False, "float32"],
+        [[], float("inf"), -1, True, "float32"],
+        [[], float("inf"), -1, False, "float32"],
+        [[], float("inf"), 0, True, "float32"],
+        [[], float("inf"), 0, False, "float32"],
+        [[], 0, -1, True, "float32"],
+        [[], 0, -1, False, "float32"],
+        [[], 0, 0, True, "float32"],
+        [[], 0, 0, False, "float32"],
+        [[], 1, -1, True, "float32"],
+        [[], 1, -1, False, "float32"],
+        [[], 1, 0, True, "float32"],
+        [[], 1, 0, False, "float32"],
+        [[], 2, -1, True, "float32"],
+        [[], 2, -1, False, "float32"],
+        [[], 2, 0, True, "float32"],
+        [[], 2, 0, False, "float32"],
+        [[], 0.3, -1, True, "float32"],
+        [[], 0.3, -1, False, "float32"],
+        [[], 0.3, 0, True, "float32"],
+        [[], 0.3, 0, False, "float32"],
+        [[], 1.5, -1, True, "float32"],
+        [[], 1.5, -1, False, "float32"],
+        [[], 1.5, 0, True, "float32"],
+        [[], 1.5, 0, False, "float32"],
+        [[], 3.7, -1, True, "float32"],
+        [[], 3.7, -1, False, "float32"],
+        [[], 3.7, 0, True, "float32"],
+        [[], 3.7, 0, False, "float32"],
+        [[1], -float("inf"), -1, True, "float32"],
+        [[1], -float("inf"), -1, False, "float32"],
+        [[1], -float("inf"), 0, True, "float32"],
+        [[1], -float("inf"), 0, False, "float32"],
+        [[1], float("inf"), -1, True, "float32"],
+        [[1], float("inf"), -1, False, "float32"],
+        [[1], float("inf"), 0, True, "float32"],
+        [[1], float("inf"), 0, False, "float32"],
+        [[1], 0, -1, True, "float32"],
+        [[1], 0, -1, False, "float32"],
+        [[1], 0, 0, True, "float32"],
+        [[1], 0, 0, False, "float32"],
+        [[1], 1, -1, True, "float32"],
+        [[1], 1, -1, False, "float32"],
+        [[1], 1, 0, True, "float32"],
+        [[1], 1, 0, False, "float32"],
+        [[1], 2, -1, True, "float32"],
+        [[1], 2, -1, False, "float32"],
+        [[1], 2, 0, True, "float32"],
+        [[1], 2, 0, False, "float32"],
+        [[1], 0.3, -1, True, "float32"],
+        [[1], 0.3, -1, False, "float32"],
+        [[1], 0.3, 0, True, "float32"],
+        [[1], 0.3, 0, False, "float32"],
+        [[1], 1.5, -1, True, "float32"],
+        [[1], 1.5, -1, False, "float32"],
+        [[1], 1.5, 0, True, "float32"],
+        [[1], 1.5, 0, False, "float32"],
+        [[1], 3.7, -1, True, "float32"],
+        [[1], 3.7, -1, False, "float32"],
+        [[1], 3.7, 0, True, "float32"],
+        [[1], 3.7, 0, False, "float32"],
+        [[100, 100], -float("inf"), -1, True, "float32"],
+        [[100, 100], -float("inf"), -1, False, "float32"],
+        [[100, 100], -float("inf"), 0, True, "float32"],
+        [[100, 100], -float("inf"), 0, False, "float32"],
+        [[100, 100], float("inf"), -1, True, "float32"],
+        [[100, 100], float("inf"), -1, False, "float32"],
+        [[100, 100], float("inf"), 0, True, "float32"],
+        [[100, 100], float("inf"), 0, False, "float32"],
+        [[100, 100], 0, -1, True, "float32"],
+        [[100, 100], 0, -1, False, "float32"],
+        [[100, 100], 0, 0, True, "float32"],
+        [[100, 100], 0, 0, False, "float32"],
+        [[100, 100], 1, -1, True, "float32"],
+        [[100, 100], 1, -1, False, "float32"],
+        [[100, 100], 1, 0, True, "float32"],
+        [[100, 100], 1, 0, False, "float32"],
+        [[100, 100], 2, -1, True, "float32"],
+        [[100, 100], 2, -1, False, "float32"],
+        [[100, 100], 2, 0, True, "float32"],
+        [[100, 100], 2, 0, False, "float32"],
+        [[100, 100], 0.3, -1, True, "float32"],
+        [[100, 100], 0.3, -1, False, "float32"],
+        [[100, 100], 0.3, 0, True, "float32"],
+        [[100, 100], 0.3, 0, False, "float32"],
+        [[100, 100], 1.5, -1, True, "float32"],
+        [[100, 100], 1.5, -1, False, "float32"],
+        [[100, 100], 1.5, 0, True, "float32"],
+        [[100, 100], 1.5, 0, False, "float32"],
+        [[100, 100], 3.7, -1, True, "float32"],
+        [[100, 100], 3.7, -1, False, "float32"],
+        [[100, 100], 3.7, 0, True, "float32"],
+        [[100, 100], 3.7, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], -float("inf"), -1, True, "float32"],
+        [[3, 4, 5, 6, 8], -float("inf"), -1, False, "float32"],
+        [[3, 4, 5, 6, 8], -float("inf"), 0, True, "float32"],
+        [[3, 4, 5, 6, 8], -float("inf"), 0, False, "float32"],
+        [[3, 4, 5, 6, 8], float("inf"), -1, True, "float32"],
+        [[3, 4, 5, 6, 8], float("inf"), -1, False, "float32"],
+        [[3, 4, 5, 6, 8], float("inf"), 0, True, "float32"],
+        [[3, 4, 5, 6, 8], float("inf"), 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 0, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 0, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 0, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 0, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 1, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 1, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 1, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 1, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 2, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 2, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 2, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 2, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 0.3, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 0.3, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 0.3, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 0.3, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 1.5, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 1.5, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 1.5, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 1.5, 0, False, "float32"],
+        [[3, 4, 5, 6, 8], 3.7, -1, True, "float32"],
+        [[3, 4, 5, 6, 8], 3.7, -1, False, "float32"],
+        [[3, 4, 5, 6, 8], 3.7, 0, True, "float32"],
+        [[3, 4, 5, 6, 8], 3.7, 0, False, "float32"],
+    ],
+)
+class TestPNormGradComp(unittest.TestCase):
+    def test_p_norm_grad_comp(self):
+        paddle.disable_static()
+        ndim = len(self.shape)
+        norm_axis = self.axis
+        if norm_axis < 0:
+            norm_axis += ndim
+
+        # skip invalid case
+        if 0 < norm_axis < ndim:
+            primal = np.random.randn(*self.shape)
+            cot_shape = list(self.shape)
+            if self.keepdim:
+                cot_shape[self.axis] = 1
+            else:
+                cot_shape.pop(self.axis)
+
+            cotagent = np.random.randn(*cot_shape)
+            np.testing.assert_allclose(
+                actual=self.actual(
+                    primal,
+                    self.porder,
+                    self.axis,
+                    self.keepdim,
+                    self.dtype,
+                    cotagent,
+                ),
+                desired=self.desired(
+                    primal,
+                    self.porder,
+                    self.axis,
+                    self.keepdim,
+                    self.dtype,
+                    cotagent,
+                ),
+                rtol=1e-6,
+                atol=0,
+            )
+
+    def actual(self, primal, porder, axis, keepdim, dtype, cotagent):
+        core.set_prim_eager_enabled(False)
+        x = paddle.to_tensor(primal, dtype=dtype, stop_gradient=False)
+        y = paddle.norm(x, porder, axis, keepdim=keepdim)
+        v = paddle.to_tensor(cotagent, dtype=dtype, stop_gradient=False)
+        v.stop_gradient = False
+        x_cotangent = paddle.grad(y, x, v)
+        return x_cotangent[0]
+
+    def desired(self, primal, porder, axis, keepdim, dtype, cotagent):
+        core.set_prim_eager_enabled(True)
+        x = paddle.to_tensor(primal, dtype=dtype, stop_gradient=False)
+        y = paddle.norm(x, porder, axis, keepdim=keepdim)
+        v = paddle.to_tensor(cotagent, dtype=dtype, stop_gradient=False)
+        v.stop_gradient = False
+        x_cotangent = paddle.grad(y, x, v)
+        return x_cotangent[0]
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_pow_grad.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+sys.path.append('../../../../legacy_test/')
+import unittest
+
+import numpy as np
+import parameterized as param
+from op_test import OpTest, convert_float_to_uint16
+
+import paddle
+from paddle.base import core
+
+
+class TestPowOp(OpTest):
+    def setUp(self):
+        self.op_type = "pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self.dtype = self.get_dtype()
+        self.init_test_data()
+        self.if_enable_cinn()
+        self.inputs = {'X': self.x}
+        self.attrs = {'factor': self.factor}
+
+        self.outputs = {'Out': np.power(self.x, self.factor)}
+
+    def get_dtype(self):
+        return "float64"
+
+    def test_check_output(self):
+        if self.dtype == np.uint16:
+            place = core.CUDAPlace(0)
+            self.check_output_with_place(place, check_pir=True)
+        else:
+            self.check_output(check_pir=True)
+
+    def test_check_grad(self):
+        if self.dtype == np.uint16:
+            place = core.CUDAPlace(0)
+            self.check_grad_with_place(
+                place,
+                ['X'],
+                'Out',
+                check_prim=True,
+                check_pir=True,
+            )
+        else:
+            self.check_grad(
+                ['X'],
+                'Out',
+                check_prim=True,
+                check_pir=True,
+            )
+
+    def init_test_data(self):
+        if self.dtype == np.uint16:
+            x = np.random.random((5, 1, 4, 5)).astype(np.float32)
+            # x = np.array([4,5,6]).astype(np.float32)
+            self.x = convert_float_to_uint16(x)
+        else:
+            self.x = np.random.random((5, 1, 4, 5)).astype(self.dtype)
+            # self.x = np.array([4,5,6]).astype(self.dtype)
+        self.factor = 2
+
+    def if_enable_cinn(self):
+        pass
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestPowDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_cos_double_grad_comp_dygraph(self):
+        def actual(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            core._set_prim_backward_blacklist("pow_grad")
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.pow(x, 2.7)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        def desired(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.pow(x, 2.7)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal),
+            desired=desired(self.primal),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_reshape_double_grad.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights ddxerved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'shape', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.rand(10, 1, 10),
+            [10, 10],
+            np.random.rand(10, 10),
+            np.float32,
+        ),
+        (np.random.rand(2, 60), [12, 10], np.random.rand(12, 10), np.float32),
+    ],
+)
+class TestReshapeDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_reshape_double_grad_comp(self):
+        def actual(primal0, shape):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            # disable rshape_grad to trigger the composite double_grad
+            core._set_prim_backward_blacklist("reshape_grad")
+
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.reshape(x, shape)
+            # wrap by tanh for >= 2 order derivative
+            out = paddle.tanh(out)
+            dx = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            ddx = paddle.grad(dx, [x], create_graph=True, retain_graph=True)
+            return ddx[0].numpy()
+
+        def desired(primal0, shape):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.reshape(x, shape)
+            # wrap by tanh for >= 2 order derivative
+            out = paddle.tanh(out)
+            dx = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            ddx = paddle.grad(dx, [x], create_graph=True, retain_graph=True)
+            return ddx[0].numpy()
+
+        actual_result = actual(self.primal, self.shape)
+
+        desired_result = desired(self.primal, self.shape)
+
+        np.testing.assert_allclose(
+            actual=actual_result,
+            desired=desired_result,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_reshape_grad.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'shape', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.rand(10, 1, 10),
+            [10, 10],
+            np.random.rand(10, 10),
+            np.float32,
+        ),
+        (np.random.rand(2, 60), [12, 10], np.random.rand(12, 10), np.float32),
+    ],
+)
+class TestReshapeGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_reshape_grad_comp(self):
+        def actual(primal0, shape):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.reshape(x, shape)
+            res = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            return res[0].numpy()
+
+        def desired(primal0, shape):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.reshape(x, shape)
+            res = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            return res[0].numpy()
+
+        dx = actual(self.primal, self.shape)
+
+        ddx = desired(self.primal, self.shape)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_sigmoid_grad.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+import paddle.nn.functional as F
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestExpGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        core.set_prim_eager_enabled(True)
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def setUp(self):
+        paddle.enable_static()
+
+    def tearDown(self):
+        paddle.disable_static()
+
+    def test_sigmoid_grad_comp(self):
+        def actual(primal, cotangent):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+
+            x = paddle.to_tensor(primal)
+            dout = paddle.to_tensor(cotangent)
+            x.stop_gradient = False
+            return paddle.grad(F.sigmoid(x), x, dout)[0]
+
+        def desired(primal, cotangent):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+
+            x = paddle.to_tensor(primal)
+            dout = paddle.to_tensor(cotangent)
+            x.stop_gradient = False
+            return paddle.grad(F.sigmoid(x), x, dout)[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent),
+            desired=desired(self.primal, self.cotangent),
+            rtol=1e-6,
+            atol=0,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_sin_double_grad.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestSinDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_sin_double_grad_comp_dygraph(self):
+        def actual(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            core._set_prim_backward_blacklist("sin_grad")
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.sin(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        def desired(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.sin(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal),
+            desired=desired(self.primal),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_sqrt_grad.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import autograd
+import autograd.numpy
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestSqrtGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_sqrt_grad_comp(self):
+        def actual(primal, cotangent):
+            paddle.disable_static()
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            v = paddle.to_tensor(
+                cotangent, dtype='float32', stop_gradient=False
+            )
+            y = paddle.sqrt(x)
+            return paddle.grad(y, x, v, create_graph=True, retain_graph=True)[0]
+
+        def desired(primal, cotangent):
+            return autograd.make_vjp(autograd.numpy.sqrt)(primal)[0](cotangent)
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent),
+            desired=desired(self.primal, self.cotangent),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_sub_grad.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal0', 'primal1', 'dtype'),
+    [
+        (
+            np.random.rand(2, 3, 4),
+            np.random.rand(2, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.random.rand(2, 3, 1, 1),
+            np.float32,
+        ),
+    ],
+)
+class TestSubGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal0 = cls.primal0.astype(cls.dtype)
+        cls.primal1 = cls.primal1.astype(cls.dtype)
+
+    def test_sub_grad_comp(self):
+        def actual(primal0, primal1):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.subtract(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        def desired(primal0, primal1):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            y = paddle.to_tensor(primal1, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y.stop_gradient = False
+            out = paddle.subtract(x, y)
+            res = paddle.grad(out, [x, y], create_graph=True, retain_graph=True)
+            return res[0].numpy(), res[1].numpy()
+
+        dx, dy = actual(self.primal0, self.primal1)
+
+        ddx, ddy = desired(self.primal0, self.primal1)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        np.testing.assert_allclose(
+            actual=dy,
+            desired=ddy,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_sum_grad.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+
+def actual(primal, cotangent, axis, keep_dim):
+    core.set_prim_eager_enabled(False)
+    x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+    v = paddle.to_tensor(cotangent, dtype='float32', stop_gradient=False)
+    y = paddle.sum(x, axis=axis, keepdim=keep_dim)
+    x_cotangent = paddle.grad(y, x, v, create_graph=True, retain_graph=True)
+    return x_cotangent[0]
+
+
+def desired(primal, cotangent, axis, keep_dim):
+    core.set_prim_eager_enabled(True)
+    x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+    v = paddle.to_tensor(cotangent, dtype='float32', stop_gradient=False)
+    y = paddle.sum(x, axis=axis, keepdim=keep_dim)
+    x_cotangent = paddle.grad(y, x, v, create_graph=True, retain_graph=True)
+    return x_cotangent[0]
+
+
+class TestSumGradComp(unittest.TestCase):
+    def test_sum_grad_comp_1(self):
+        self.primal = np.random.rand(10, 10)
+        self.cotangent = np.array(np.random.rand())
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, [], False),
+            desired=desired(self.primal, self.cotangent, [], False),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_sum_grad_comp_2(self):
+        self.primal = np.random.rand(4, 3, 2)
+        self.cotangent = np.random.rand(4, 2)
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, 1, False),
+            desired=desired(self.primal, self.cotangent, 1, False),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_sum_grad_comp_3(self):
+        self.primal = np.random.rand(4, 3, 2)
+        self.cotangent = np.random.rand(4, 1, 2)
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, 1, True),
+            desired=desired(self.primal, self.cotangent, 1, True),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_sum_grad_comp_4(self):
+        self.primal = np.random.rand(4, 3, 2, 5)
+        self.cotangent = np.random.rand(4, 1, 2, 1)
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, [1, 3], True),
+            desired=desired(self.primal, self.cotangent, [1, 3], True),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_sum_grad_comp_5(self):
+        self.primal = np.random.rand(4, 3, 2, 5)
+        self.cotangent = np.random.rand(4, 2)
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, [1, 3], False),
+            desired=desired(self.primal, self.cotangent, [1, 3], False),
+            rtol=1e-6,
+            atol=0,
+        )
+
+    def test_sum_grad_comp_6(self):
+        self.primal = np.random.rand(3, 2, 5)
+        self.cotangent = np.random.rand(3, 1, 1)
+        paddle.disable_static()
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal, self.cotangent, [-2, -1], True),
+            desired=desired(self.primal, self.cotangent, [-2, -1], True),
+            rtol=1e-6,
+            atol=0,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_take_along_axis_double_grad.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+
+
+@param.parameterized_class(
+    ('arr', 'indices', 'axis', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.randn(4, 3, 2),  # arr
+            np.random.randint(0, 3, size=(1, 3, 2)),  # indices
+            1,  # axis
+            np.random.rand(1, 3, 2),  # cotangent
+            np.float32,  # dtype
+        ),
+    ],
+)
+class TestTakeAlongAxisTanhDoubleGrad(unittest.TestCase):
+    def test_take_along_axis_tanh_double_grad(self):
+        arr_tensor = paddle.to_tensor(
+            self.arr, dtype=self.dtype, stop_gradient=False
+        )
+        indices_tensor = paddle.to_tensor(self.indices, dtype="int64")
+        dout_tensor = paddle.to_tensor(
+            self.cotangent, dtype=self.dtype, stop_gradient=False
+        )
+        out = paddle.take_along_axis(arr_tensor, indices_tensor, axis=self.axis)
+
+        out = paddle.tanh(out)
+
+        dx = paddle.grad(out, arr_tensor, dout_tensor, create_graph=True)[0]
+
+        ddx = paddle.grad(dx, dout_tensor, create_graph=True)[0]
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_tanh_grad.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'dtype'),
+    [
+        (
+            np.random.rand(2, 3, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 3, 4),
+            np.float32,
+        ),
+    ],
+)
+class TestTanhGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_tanh_grad_comp(self):
+        def actual(primal):
+            paddle.disable_static()
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.tanh(x)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=True, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        def desired(primal):
+            paddle.disable_static()
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.tanh(x)
+            x_cotangent = paddle.grad(
+                y, x, create_graph=True, retain_graph=True
+            )
+            return x_cotangent[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal),
+            desired=desired(self.primal),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/prim/vjp/eager/test_comp_eager_transpose_grad.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+core.set_prim_eager_enabled(True)
+
+
+@param.parameterized_class(
+    ('primal', 'axis', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.rand(
+                100,
+            ),
+            [0],
+            np.random.rand(100),
+            np.float32,
+        ),
+        (
+            np.random.rand(3, 4, 10),
+            [0, 2, 1],
+            np.random.rand(3, 10, 4),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4, 5),
+            [0, 2, 3, 1],
+            np.random.rand(2, 4, 5, 3),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4, 5, 6),
+            [4, 2, 3, 1, 0],
+            np.random.rand(6, 4, 5, 3, 2),
+            np.float32,
+        ),
+        (
+            np.random.rand(2, 3, 4, 5, 6, 1),
+            [4, 2, 3, 1, 0, 5],
+            np.random.rand(6, 4, 5, 3, 2, 1),
+            np.float32,
+        ),
+        # (np.random.rand(),
+        #  [],
+        # np.random.rand(),
+        # np.float32),
+    ],
+)
+class TestTransposeGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if isinstance(cls.primal, np.ndarray):
+            cls.primal = cls.primal.astype(cls.dtype)
+
+    def test_transpose_grad_comp(self):
+        def actual(primal0, shape):
+            core.set_prim_eager_enabled(True)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.transpose(x, shape)
+            res = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            return res[0].numpy()
+
+        def desired(primal0, shape):
+            core.set_prim_eager_enabled(False)
+            paddle.disable_static()
+            x = paddle.to_tensor(primal0, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            out = paddle.transpose(x, shape)
+            res = paddle.grad(out, [x], create_graph=True, retain_graph=True)
+            return res[0].numpy()
+
+        dx = actual(self.primal, self.axis)
+
+        ddx = desired(self.primal, self.axis)
+
+        np.testing.assert_allclose(
+            actual=dx,
+            desired=ddx,
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/process/CMakeLists.txt

@@ -0,0 +1,10 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS})
+endforeach()

test/prim/process/test_prim_amp.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.base import core, framework
+from paddle.nn import BatchNorm
+
+np.random.seed(2023)
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2D(2, 4, (3, 3), bias_attr=False)
+        self.bn = BatchNorm(4, act="relu")
+
+    def forward(self, x):
+        y = self.conv(x)
+        out = self.bn(y)
+        res = F.max_pool2d(out, kernel_size=2, stride=2, padding=0)
+        return res
+
+
+class TestPrimAMPO1(unittest.TestCase):
+    """
+    Test PrimeNet with @to_static + prim v.s Dygraph in AMPO1.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.x = paddle.randn([4, 2, 6, 6], dtype="float32")
+        self.x.stop_gradient = False
+        self.atol = 1e-3
+        self.rtol = 1e-3
+        if paddle.is_compiled_with_xpu():
+            self.atol = 5e-3
+            self.rtol = 5e-3
+
+    def train(self, use_prim):
+        core._set_prim_all_enabled(use_prim)
+        paddle.seed(2022)
+        net = PrimeNet()
+        sgd = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=net.parameters()
+        )
+
+        if use_prim:
+            net = paddle.jit.to_static(
+                net, build_strategy=False, full_graph=True
+            )
+        with paddle.amp.auto_cast(level='O1'):
+            out = net(self.x)
+            loss = paddle.mean(out)
+            loss.backward()
+            sgd.step()
+            sgd.clear_grad()
+            return loss
+
+    def test_amp_01(self):
+        if not isinstance(framework._current_expected_place(), core.CPUPlace):
+            expected = self.train(False)
+            actual = self.train(True)
+            np.testing.assert_allclose(
+                expected,
+                actual,
+                rtol=self.rtol,
+                atol=self.atol,
+            )
+
+    def test_amp_O1_infer(self):
+        if not isinstance(framework._current_expected_place(), core.CPUPlace):
+            net = PrimeNet()
+            core._set_prim_all_enabled(False)
+            net.eval()
+            static_net = paddle.jit.to_static(
+                net, build_strategy=False, full_graph=True
+            )
+            res = static_net(self.x)
+
+            # set prim all enabled
+            core._set_prim_all_enabled(True)
+            net.eval()
+            static_net = paddle.jit.to_static(
+                net, build_strategy=False, full_graph=True
+            )
+            with paddle.amp.auto_cast(level='O1'):
+                res_amp = static_net(self.x)
+
+            np.testing.assert_allclose(
+                res,
+                res_amp,
+                rtol=self.rtol,
+                atol=self.atol,
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/prim/test_comp_get_grad_op_desc_prim_disabled.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+
+from paddle.base import core
+
+core._set_prim_backward_enabled(False)
+
+import parameterized as param
+
+import paddle
+from paddle.base import core, framework
+
+
+@param.parameterized_class(
+    (
+        'fwd_type',
+        'inputs',
+        'outputs',
+        'no_grad_var',
+        'grad_sub_block',
+        'desired_ops',
+    ),
+    (
+        ('tanh', {'X': ['x']}, {'Out': ['y']}, set(), (), ('tanh_grad',)),
+        ('empty', {}, {'Out': ['y']}, set(), (), ()),
+    ),
+)
+class TestGetGradOpDescPrimEnabled(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        paddle.enable_static()
+        block = framework.Block(framework.Program(), 0)
+        block.append_op(
+            type=cls.fwd_type,
+            inputs={
+                n: [block.create_var(name=v, stop_gradient=False) for v in vs]
+                for n, vs in cls.inputs.items()
+            },
+            outputs={
+                n: [block.create_var(name=v, stop_gradient=False) for v in vs]
+                for n, vs in cls.outputs.items()
+            },
+        )
+        cls.fwd = block.ops[0].desc
+
+    @classmethod
+    def tearDownClass(cls):
+        paddle.disable_static()
+
+    def test_get_grad_op_desc(self):
+        actual = tuple(
+            desc.type()
+            for desc in core.get_grad_op_desc(
+                self.fwd, self.no_grad_var, self.grad_sub_block
+            )[0]
+        )
+        self.assertEqual(actual, self.desired_ops)
+
+
+if __name__ == '__main__':
+    unittest.main()