chore: import upstream snapshot with attribution

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()