chore: import upstream snapshot with attribution

python/paddle/static/__init__.py

@@ -0,0 +1,119 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#   Copyright (c) 2021 NVIDIA Corporation. 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 ..base import Scope  # noqa: F401
+from ..base.backward import append_backward, gradients
+from ..base.compiler import (
+    BuildStrategy,
+    CompiledProgram,
+    IpuCompiledProgram,
+    IpuStrategy,
+)
+from ..base.executor import Executor, global_scope, scope_guard
+from ..base.framework import (
+    Operator,  # noqa: F401
+    Parameter,  # noqa: F401
+    Program,
+    Variable,
+    cpu_places,
+    cuda_places,
+    default_main_program,
+    default_startup_program,
+    device_guard,
+    ipu_shard_guard,
+    name_scope,
+    program_guard,
+    set_ipu_shard,
+    xpu_places,
+)
+from ..base.libpaddle import NativeMetaTensor as MetaTensor  # noqa: F401
+from ..base.param_attr import WeightNormParamAttr
+from ..tensor.creation import create_global_var, create_parameter
+from . import amp, nn  # noqa: F401
+from .input import (
+    InputSpec,
+    data,
+    setitem,  # noqa: F401
+)
+from .io import (
+    deserialize_persistables,
+    deserialize_program,
+    is_persistable,  # noqa: F401
+    load,
+    load_from_file,
+    load_inference_model,
+    load_program_state,
+    load_vars,  # noqa: F401
+    normalize_program,
+    save,
+    save_inference_model,
+    save_to_file,
+    save_vars,  # noqa: F401
+    serialize_persistables,
+    serialize_program,
+    set_program_state,
+)
+from .nn.common import ExponentialMovingAverage, py_func
+from .nn.control_flow import Print
+from .nn.metric import accuracy, auc, ctr_metric_bundle
+from .python_op import register_op  # noqa: F401
+
+__all__ = [
+    'append_backward',
+    'gradients',
+    'Executor',
+    'global_scope',
+    'scope_guard',
+    'BuildStrategy',
+    'CompiledProgram',
+    'ipu_shard_guard',
+    'IpuCompiledProgram',
+    'IpuStrategy',
+    'Print',
+    'py_func',
+    'name_scope',
+    'program_guard',
+    'WeightNormParamAttr',
+    'ExponentialMovingAverage',
+    'default_main_program',
+    'default_startup_program',
+    'Program',
+    'data',
+    'InputSpec',
+    'save',
+    'load',
+    'save_inference_model',
+    'load_inference_model',
+    'serialize_program',
+    'serialize_persistables',
+    'save_to_file',
+    'deserialize_program',
+    'deserialize_persistables',
+    'load_from_file',
+    'normalize_program',
+    'load_program_state',
+    'set_program_state',
+    'cpu_places',
+    'cuda_places',
+    'xpu_places',
+    'Variable',
+    'create_global_var',
+    'accuracy',
+    'auc',
+    'device_guard',
+    'create_parameter',
+    'set_ipu_shard',
+    'ctr_metric_bundle',
+]

python/paddle/static/amp/__init__.py

@@ -0,0 +1,28 @@
+#   Copyright (c) 2020 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 . import (  # noqa: F401
+    bf16,
+    debugging,
+    decorator,
+    fp16_lists,
+    fp16_utils,
+)
+from .decorator import decorate  # noqa: F401
+from .fp16_lists import AutoMixedPrecisionLists, CustomOpLists  # noqa: F401
+from .fp16_utils import (  # noqa: F401
+    cast_model_to_fp16,
+    cast_parameters_to_fp16,
+    fp16_guard,
+)

python/paddle/static/amp/amp_nn.py

@@ -0,0 +1,176 @@
+#   Copyright (c) 2020 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 paddle
+from paddle import _C_ops
+from paddle.base.data_feeder import check_type, check_variable_and_dtype
+from paddle.base.framework import (
+    Variable,
+    in_dynamic_or_pir_mode,
+)
+from paddle.base.layer_helper import LayerHelper
+
+
+def check_finite_and_unscale(x, scale, name=None, float_status=None):
+    """
+    Check if input X contains all finite data, if yes, scale it by input Scale.
+
+    $$Out = X / scale$$
+
+    If any tensor in X contains Inf or Nan, the Out will generate a indicator.
+    FoundInfinite will be 1 (True), and Out will not be scaled. In this case, the data of
+    Out should not be used, and its data may not be deterministic.
+    Otherwise, FoundInfinite will be 0 (False).
+
+    Args:
+        x(list|tuple): The input tensors of check_finite_and_unscale operator.
+        scale: The scale of check_finite_and_unscale operator.
+        float_status(Tensor): (Only used on NPU) The float status to check overflow.
+    """
+
+    if in_dynamic_or_pir_mode():
+        x, found_inf = _C_ops.check_finite_and_unscale_(x, scale)
+        return x, found_inf
+
+    helper = LayerHelper("check_finite_and_unscale", **locals())
+    found_inf = helper.create_variable_for_type_inference(dtype='bool')
+    check_type(x, 'x', (tuple, list), 'check_finite_and_unscale')
+    for e in x:
+        check_variable_and_dtype(
+            e,
+            "x",
+            ['float16', 'float32', 'float64', 'uint16'],
+            'check_finite_and_unscale',
+        )
+
+    inputs = {'X': x, 'Scale': scale}
+    outputs = {'Out': x, 'FoundInfinite': found_inf}
+    helper.append_op(
+        type='check_finite_and_unscale', inputs=inputs, outputs=outputs
+    )
+
+    return x, found_inf
+
+
+def update_loss_scaling(
+    x,
+    found_inf,
+    prev_loss_scaling,
+    num_good_steps,
+    num_bad_steps,
+    incr_every_n_steps,
+    decr_every_n_nan_or_inf,
+    incr_ratio,
+    decr_ratio,
+    stop_update=False,
+    name=None,
+):
+    """
+    Update loss scaling according to overall gradients. If all gradients is
+    finite after incr_every_n_steps, loss scaling will increase by incr_ratio.
+    Otherwise, loss scaling will decrease by decr_ratio after
+    decr_every_n_nan_or_inf steps and each step some gradients are infinite.
+
+    Args:
+        x(list|tuple): The input tensors of update_loss_scaling operator.
+        found_inf (Variable): A boolean variable indicates whether
+                                     there is any infinite gradient.
+        prev_loss_scaling (Variable): Previous loss scaling.
+        num_good_steps (Variable): A variable accumulates good steps in which
+                                   all gradients are finite.
+        num_bad_steps (Variable): A variable accumulates bad steps in which
+                                  some gradients are infinite.
+        incr_every_n_steps (int): A variable represents increasing loss
+                                       scaling every n consecutive steps with
+                                       finite gradients.
+        decr_every_n_nan_or_inf (int): A variable represents decreasing
+                                            loss scaling every n accumulated
+                                            steps with nan or inf gradients.
+        incr_ratio(float): The multiplier to use when increasing the loss
+                           scaling.
+        decr_ratio(float): The less-than-one-multiplier to use when decreasing
+                           loss scaling.
+    """
+    if in_dynamic_or_pir_mode():
+        _C_ops.update_loss_scaling_(
+            x,
+            found_inf,
+            prev_loss_scaling,
+            num_good_steps,
+            num_bad_steps,
+            incr_every_n_steps,
+            decr_every_n_nan_or_inf,
+            incr_ratio,
+            decr_ratio,
+            stop_update,
+        )
+        return x
+
+    check_variable_and_dtype(
+        prev_loss_scaling,
+        "prev_loss_scaling",
+        ['float32', 'float64'],
+        "update_loss_scaling",
+    )
+    check_type(x, 'x', (tuple, list), 'update_loss_scaling')
+    for e in x:
+        check_variable_and_dtype(
+            e,
+            "x",
+            ['float16', 'float32', 'float64', 'uint16'],
+            'update_loss_scaling',
+        )
+        if e.dtype in [paddle.float16, paddle.bfloat16]:
+            assert prev_loss_scaling.dtype == paddle.float32, (
+                "The dtype of prev_loss_scaling should be float32 when the dtype of x is float16 or bfloat16."
+            )
+        else:
+            assert prev_loss_scaling.dtype == e.dtype, (
+                "The dtype of prev_loss_scaling should be equal to the dtype of x."
+            )
+
+    helper = LayerHelper("update_loss_scaling", **locals())
+
+    inputs = {
+        'X': x,
+        'FoundInfinite': found_inf,
+        'PrevLossScaling': prev_loss_scaling,
+        'InGoodSteps': num_good_steps,
+        'InBadSteps': num_bad_steps,
+    }
+
+    outputs = {
+        'Out': x,
+        'LossScaling': prev_loss_scaling,
+        'OutGoodSteps': num_good_steps,
+        'OutBadSteps': num_bad_steps,
+    }
+
+    attrs = {
+        'incr_every_n_steps': incr_every_n_steps,
+        'decr_every_n_nan_or_inf': decr_every_n_nan_or_inf,
+        'incr_ratio': incr_ratio,
+        'decr_ratio': decr_ratio,
+    }
+
+    if isinstance(stop_update, Variable):
+        inputs['StopUpdate'] = stop_update
+    else:
+        attrs['stop_update'] = stop_update
+
+    helper.append_op(
+        type='update_loss_scaling', inputs=inputs, outputs=outputs, attrs=attrs
+    )
+
+    return x

python/paddle/static/amp/bf16/__init__.py

@@ -0,0 +1,28 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 . import (  # noqa: F401
+    amp_lists,
+    amp_utils,
+    decorator,
+)
+from .amp_lists import AutoMixedPrecisionListsBF16  # noqa: F401
+from .amp_utils import (  # noqa: F401
+    bf16_guard,
+    cast_model_to_bf16,
+    cast_parameters_to_bf16,
+    convert_float_to_uint16,
+    rewrite_program_bf16,
+)
+from .decorator import decorate_bf16  # noqa: F401

python/paddle/static/amp/bf16/amp_lists.py

@@ -0,0 +1,110 @@
+#   Copyright (c) 2021 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 copy
+
+from paddle.amp.amp_lists import BF16_WHITE_LIST
+from paddle.base import core
+
+from ..fp16_lists import (
+    black_list as black_list_fp16,
+    gray_list as gray_list_fp16,
+    white_list as white_list_fp16,
+)
+
+
+class AutoMixedPrecisionListsBF16:
+    """
+    AutoMixedPrecisionListsBF16 is a class for fp32/bf16 op types list. The lists are used for an
+    algorithm which determines op's execution mode (fp32 or bf16).It can update pre-defined
+    fp32 list and bf16 list according to users' custom fp32 bf16 lists.
+
+    Args:
+        custom_bf16_list (set): Users' custom bf16 list.
+        custom_fp32_list (set): Users' custom fp32 list.
+        custom_fp32_varnames (set): Users' custom fp32 variables' names.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.enable_static()
+            >>> with paddle.static.amp.bf16.bf16_guard():
+            ...     paddle.static.amp.bf16.AutoMixedPrecisionListsBF16(custom_fp32_list={'lstm'})
+    """
+
+    def __init__(
+        self,
+        custom_bf16_list=None,
+        custom_fp32_list=None,
+        custom_fp32_varnames=None,
+    ):
+        self._custom_bf16_list = custom_bf16_list
+        self._custom_fp32_list = custom_fp32_list
+        self.bf16_list = copy.copy(bf16_list)
+        self.fp32_list = copy.copy(fp32_list)
+        self.gray_list = copy.copy(gray_list)
+        self.bf16_initializer_list = copy.copy(bf16_initializer_list)
+        self.unsupported_list = copy.copy(unsupported_list)
+        self.fp32_varnames = copy.copy(custom_fp32_varnames)
+        self._update_list()
+
+    def _update_list(self):
+        """
+        Update fp32 and bf16 list according to users' custom list.
+        """
+        if self._custom_bf16_list and self._custom_fp32_list:
+            for op_name in self._custom_bf16_list:
+                if op_name in self._custom_fp32_list:
+                    raise ValueError(
+                        "Custom bf16 list overlap custom fp32 list"
+                    )
+        if self._custom_bf16_list:
+            for op_name in self._custom_bf16_list:
+                if op_name in self.fp32_list:
+                    self.fp32_list.remove(op_name)
+                elif op_name in self.gray_list:
+                    self.gray_list.remove(op_name)
+                self.bf16_list.add(op_name)
+        if self._custom_fp32_list:
+            for op_name in self._custom_fp32_list:
+                if op_name in self.bf16_list:
+                    self.bf16_list.remove(op_name)
+                elif op_name in self.gray_list:
+                    self.gray_list.remove(op_name)
+                self.fp32_list.add(op_name)
+                self.unsupported_list.add(op_name)
+
+
+bf16_initializer_list = {'fill_constant', 'uniform_random'}
+
+# always bf16
+bf16_list = BF16_WHITE_LIST
+
+# depends on the prev_op type
+gray_list = gray_list_fp16
+
+_, _, _sys_unsupported_bf16_list = core.op_supported_infos(
+    'CPU', core.VarDesc.VarType.BF16
+)
+unsupported_list = _sys_unsupported_bf16_list
+
+fp32_list = black_list_fp16.copy().copy()
+fp32_list |= white_list_fp16
+fp32_list |= gray_list_fp16
+
+fp32_list -= bf16_list
+fp32_list -= gray_list
+unsupported_list -= bf16_list
+unsupported_list -= gray_list

python/paddle/static/amp/bf16/amp_utils.py

@@ -0,0 +1,600 @@
+#   Copyright (c) 2021 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 collections
+import logging
+import struct
+
+import numpy as np
+
+import paddle
+from paddle.base import core, framework, global_scope
+from paddle.base.log_helper import get_logger
+from paddle.base.wrapped_decorator import signature_safe_contextmanager
+
+from ..fp16_utils import (
+    _rename_arg,
+    _rename_op_input,
+    find_true_post_op,
+    find_true_prev_op,
+)
+from .amp_lists import AutoMixedPrecisionListsBF16
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+_valid_types = [
+    core.VarDesc.VarType.DENSE_TENSOR,
+    core.VarDesc.VarType.SELECTED_ROWS,
+    core.VarDesc.VarType.DENSE_TENSOR_ARRAY,
+]
+
+_bf16_guard_pattern = "__use_bf16__"
+
+
+def convert_float_to_uint16(in_list):
+    in_list = np.asarray(in_list)
+    out = np.vectorize(
+        lambda x: struct.unpack('<I', struct.pack('<f', x))[0] >> 16,
+        otypes=[np.uint16],
+    )(in_list.flat)
+    return np.reshape(out, in_list.shape)
+
+
+def _dtype_to_str(dtype):
+    """
+    Convert specific variable type to its corresponding string.
+
+    Args:
+        dtype (VarType): Variable type.
+    """
+    if dtype == paddle.bfloat16:
+        return 'bf16'
+    else:
+        return 'fp32'
+
+
+def _insert_cast_op(block, op, idx, src_dtype, dest_dtype):
+    """
+    Insert cast op and rename args of input and output.
+
+    Args:
+        block (Program): The block in which the operator is.
+        op (Operator): The operator to insert cast op.
+        idx (int): The index of current operator.
+        src_dtype (VarType): The input variable dtype of cast op.
+        dest_dtype (VarType): The output variable dtype of cast op.
+
+    Returns:
+        num_cast_op (int): The number of cast ops that have been inserted.
+    """
+    num_cast_ops = 0
+
+    for in_name in op.input_names:
+        if src_dtype == paddle.float32 and op.type in [
+            'batch_norm',
+            'fused_bn_add_activation',
+            'layer_norm',
+        ]:
+            if in_name not in {'X', 'Z'}:
+                continue
+        for in_var_name in op.input(in_name):
+            in_var = block.var(in_var_name)
+            if in_var.type not in _valid_types or in_var.dtype == dest_dtype:
+                continue
+            if in_var.dtype == src_dtype:
+                cast_name = in_var.name + '.cast_' + _dtype_to_str(dest_dtype)
+                out_var = block.vars.get(cast_name)
+                if out_var is None or out_var.dtype != dest_dtype:
+                    out_var = block.create_var(
+                        name=cast_name,
+                        dtype=dest_dtype,
+                        persistable=False,
+                        stop_gradient=in_var.stop_gradient,
+                    )
+
+                    block._insert_op(
+                        idx,
+                        type="cast",
+                        inputs={"X": in_var},
+                        outputs={"Out": out_var},
+                        attrs={
+                            "in_dtype": in_var.dtype,
+                            "out_dtype": out_var.dtype,
+                        },
+                    )
+                    num_cast_ops += 1
+                _rename_arg(op, in_var.name, out_var.name)
+            else:
+                if op.has_attr('in_dtype'):
+                    op._set_attr('in_dtype', dest_dtype)
+    if src_dtype == paddle.float32 and dest_dtype == paddle.bfloat16:
+        for out_name in op.output_names:
+            if (
+                op.type
+                in ['batch_norm', 'fused_bn_add_activation', 'layer_norm']
+                and out_name != 'Y'
+            ):
+                continue
+            for out_var_name in op.output(out_name):
+                out_var = block.var(out_var_name)
+                if out_var.type not in _valid_types:
+                    continue
+                if out_var.dtype == paddle.float32:
+                    out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
+                    if op.has_attr('out_dtype'):
+                        op._set_attr('out_dtype', core.VarDesc.VarType.BF16)
+    return num_cast_ops
+
+
+def _insert_cast_post_op(
+    block, op, idx, src_dtype, dest_dtype, target_name, op_var_rename_map
+):
+    num_cast_ops = 0
+    target_var = block.var(target_name)
+    if target_var.type not in _valid_types or target_var.dtype == dest_dtype:
+        return num_cast_ops
+
+    assert target_var.dtype == src_dtype, (
+        f"The real dtype({_dtype_to_str(target_var.dtype)}) is not equal to the src dtype({_dtype_to_str(src_dtype)})"
+    )
+
+    cast_name = target_var.name + '.cast_' + _dtype_to_str(dest_dtype)
+    cast_var = block.vars.get(cast_name)
+    if cast_var is None or cast_var.dtype != dest_dtype:
+        cast_var = block.create_var(
+            name=cast_name,
+            dtype=dest_dtype,
+            persistable=False,
+            stop_gradient=target_var.stop_gradient,
+        )
+        block._insert_op(
+            idx,
+            type="cast",
+            inputs={"X": target_var},
+            outputs={"Out": cast_var},
+            attrs={"in_dtype": target_var.dtype, "out_dtype": cast_var.dtype},
+        )
+        num_cast_ops += 1
+        op_var_rename_map[block.idx][target_var.name] = cast_var.name
+
+    return num_cast_ops
+
+
+def _is_in_fp32_varnames(op, amp_lists):
+    if not amp_lists.fp32_varnames:
+        return False
+
+    for in_name in op.input_arg_names:
+        if in_name in amp_lists.fp32_varnames:
+            return True
+
+    for out_name in op.output_arg_names:
+        if out_name in amp_lists.fp32_varnames:
+            return True
+
+    return False
+
+
+def _need_keep_fp32(op, unsupported_op_list, use_bf16_guard):
+    if op.type in unsupported_op_list:
+        # the highest priority condition: If ops don't have bf16 computing kernels,
+        # they must be executed in fp32 calculation pattern.
+        return True
+
+    # process ops about learning rate
+    in_out_arg_names = []
+    in_out_arg_names.extend(list(op.input_arg_names))
+    in_out_arg_names.extend(list(op.output_arg_names))
+    for name in in_out_arg_names:
+        if "learning_rate" in name:
+            return True
+
+    if use_bf16_guard:
+        if op.has_attr("op_namescope") and (
+            _bf16_guard_pattern in op.attr("op_namescope")
+        ):
+            # op in bf16 guard
+            return False
+        else:
+            # op not in bf16 guard
+            return True
+    else:
+        return False
+
+
+@signature_safe_contextmanager
+def bf16_guard():
+    """
+    As for the pure bf16 training, if users set `use_bf16_guard` to True,
+    only those ops created in the context manager `bf16_guard` will be
+    transformed as float16 type.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import numpy as np
+            >>> import paddle
+            >>> import paddle.nn.functional as F
+            >>> paddle.enable_static()
+            >>> data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
+            >>> conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
+
+            >>> with paddle.static.amp.bf16.bf16_guard():
+            ...     bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
+            ...     pool = F.max_pool2d(bn, kernel_size=2, stride=2)
+            ...     hidden = paddle.static.nn.fc(pool, size=10)
+            ...     loss = paddle.mean(hidden)
+    """
+    with framework.name_scope(prefix=_bf16_guard_pattern):
+        yield
+
+
+def are_post_ops_bf16(post_ops, keep_fp32_ops):
+    for post_op in post_ops:
+        for op in post_op:
+            if op in keep_fp32_ops:
+                return False
+    return True
+
+
+def cast_initializers_to_bf16(
+    startup_prog,
+    amp_lists,
+    block,
+    all_ops,
+    keep_fp32_ops,
+    to_bf16_var_names=None,
+):
+    prepend_ops = startup_prog.global_block().ops
+    for op in prepend_ops:
+        if str(op.type) in amp_lists.bf16_initializer_list:
+            change_op = True
+            op_post_ops = []
+            op_out_vars = []
+            for out_name in op.output_names:
+                for out_var_name in op.output(out_name):
+                    out_var = block.var(out_var_name)
+                    post_op = find_true_post_op(all_ops, op, out_var_name, True)
+
+                    if out_var is None or out_var.type not in _valid_types:
+                        change_op = False
+                        break
+                    op_post_ops.append(post_op)
+                    op_out_vars.append(out_var)
+
+            if change_op and are_post_ops_bf16(op_post_ops, keep_fp32_ops):
+                for out_var in op_out_vars:
+                    if out_var.dtype == paddle.float32:
+                        out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
+                    if (
+                        to_bf16_var_names is not None
+                        and out_var.name in to_bf16_var_names
+                    ):
+                        to_bf16_var_names.remove(out_var.name)
+                if (
+                    op.has_attr('dtype')
+                    and op.attr('dtype') == core.VarDesc.VarType.FP32
+                ):
+                    op._set_attr('dtype', core.VarDesc.VarType.BF16)
+
+
+def cast_model_to_bf16(
+    program, startup_prog=None, amp_lists=None, use_bf16_guard=True
+):
+    """
+    Traverse all ops in the whole model and set their inputs and outputs
+    to the bf16 data type. This function will do some special processing for
+    the batch normalization, which will keep the batchnorm's computations in FP32.
+    Args:
+        program (Program): The used program.
+        amp_lists (AutoMixedPrecisionListsBF16): An AutoMixedPrecisionListsBF16 object.
+        use_bf16_guard(bool): Determine whether to use `bf16_guard` when
+                              constructing the program. Default True.
+    """
+
+    if amp_lists is None:
+        amp_lists = AutoMixedPrecisionListsBF16()
+    global_block = program.global_block()
+    keep_fp32_ops = set()
+    to_bf16_var_names = set()
+    to_bf16_pre_cast_ops = set()
+    origin_ops = []
+    for block in program.blocks:
+        origin_ops.extend(block.ops)
+
+    for block in program.blocks:
+        ops = block.ops
+        for op in ops:
+            if op.type == 'create_py_reader' or op.type == 'read':
+                continue
+            if _need_keep_fp32(op, amp_lists.unsupported_list, use_bf16_guard):
+                keep_fp32_ops.add(op)
+                continue  # processed below
+            for in_name in op.input_names:
+                if op.type in {
+                    'batch_norm',
+                    'fused_bn_add_activation',
+                    'layer_norm',
+                } and in_name not in {'X', 'Z'}:
+                    continue
+                for in_var_name in op.input(in_name):
+                    in_var = None
+                    try:
+                        in_var = block.var(in_var_name)
+                    except ValueError as e:
+                        _logger.debug(
+                            f"-- {e}, try to get it in the global block --"
+                        )
+                        in_var = global_block.var(in_var_name)
+                        if in_var is not None:
+                            _logger.debug(
+                                f"-- var {in_var_name} is got in the global block --"
+                            )
+
+                    if in_var is None or in_var.type not in _valid_types:
+                        continue
+
+                    if in_var.dtype == paddle.float32:
+                        in_var.desc.set_dtype(core.VarDesc.VarType.BF16)
+                        to_bf16_var_names.add(in_var_name)
+
+                    _logger.debug(
+                        f"-- op type: {op.type}, in var name: {in_var_name}, in var dtype: {in_var.dtype} --"
+                    )
+
+            for out_name in op.output_names:
+                if (
+                    op.type
+                    in {'batch_norm', 'fused_bn_add_activation', 'layer_norm'}
+                    and out_name != 'Y'
+                ):
+                    continue
+                for out_var_name in op.output(out_name):
+                    out_var = None
+                    try:
+                        out_var = block.var(out_var_name)
+                    except ValueError as e:
+                        _logger.debug(
+                            f"-- {e}, try to get it in the global block --"
+                        )
+                        out_var = global_block.var(out_var_name)
+                        if out_var is not None:
+                            _logger.debug(
+                                f"-- var {out_var_name} is got in the global block --"
+                            )
+
+                    if out_var is None or out_var.type not in _valid_types:
+                        continue
+
+                    if out_var.dtype == paddle.float32:
+                        out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
+
+                    _logger.debug(
+                        f"-- op type: {op.type}, out var name: {out_var_name}, out var dtype: {out_var.dtype} --"
+                    )
+            for attr_name in ['in_dtype', 'out_dtype', 'dtype']:
+                if (
+                    op.has_attr(attr_name)
+                    and op.attr(attr_name) == paddle.float32
+                ):
+                    op._set_attr(attr_name, core.VarDesc.VarType.BF16)
+
+        if startup_prog is not None:
+            cast_initializers_to_bf16(
+                startup_prog,
+                amp_lists,
+                global_block,
+                ops,
+                keep_fp32_ops,
+                to_bf16_var_names,
+            )
+
+    # process ops in keep_fp32_ops
+    op_var_rename_map = [
+        collections.OrderedDict() for _ in range(len(program.blocks))
+    ]
+    for block in program.blocks:
+        ops = block.ops
+        idx = 0
+        while idx < len(ops):
+            op = ops[idx]
+            num_cast_ops = 0
+            if op not in keep_fp32_ops:
+                if op in to_bf16_pre_cast_ops:
+                    in_var_cast_num = _insert_cast_op(
+                        block,
+                        op,
+                        idx,
+                        core.VarDesc.VarType.FP32,
+                        core.VarDesc.VarType.BF16,
+                    )
+                    num_cast_ops += in_var_cast_num
+            else:
+                pre_cast_num = _insert_cast_op(
+                    block,
+                    op,
+                    idx,
+                    core.VarDesc.VarType.BF16,
+                    core.VarDesc.VarType.FP32,
+                )
+                num_cast_ops += pre_cast_num
+                for out_var_name in op.output_arg_names:
+                    out_var = block.vars.get(out_var_name)
+                    if out_var is None or out_var.type not in _valid_types:
+                        continue
+                    if out_var.dtype == paddle.bfloat16:
+                        out_var.desc.set_dtype(core.VarDesc.VarType.FP32)
+                        post_ops = find_true_post_op(ops, op, out_var_name)
+                        for post_op in post_ops:
+                            if post_op in keep_fp32_ops:
+                                continue
+                            post_cast_num = _insert_cast_post_op(
+                                block,
+                                op,
+                                idx + pre_cast_num + 1,
+                                core.VarDesc.VarType.FP32,
+                                core.VarDesc.VarType.BF16,
+                                out_var_name,
+                                op_var_rename_map,
+                            )
+                            num_cast_ops += post_cast_num
+            idx += num_cast_ops + 1
+
+    _rename_op_input(program, op_var_rename_map, origin_ops, keep_fp32_ops)
+    return to_bf16_var_names
+
+
+def cast_parameters_to_bf16(place, program, scope=None, to_bf16_var_names=None):
+    """
+    Traverse all parameters in the whole model and set them to the BF16 data type.
+    Whereas, this function will keep parameters of batchnorms in FP32.
+    Args:
+        place(base.CPUPlace|base.CUDAPlace): `place` is used to restore the BF16 weight tensors.
+        program (Program): The used program.
+        scope(base.Scope, optional): `scope` is used to get the FP32 weight tensor values.
+                                      Default is None.
+        to_bf16_var_names(set|list, optional): The data types of vars in `to_bf16_var_names`
+                                               will be set to BF16. Usually, it is the returned
+                                               value of `cast_model_to_bf16` API.
+    """
+    all_parameters = []
+    for block in program.blocks:
+        all_parameters.extend(block.all_parameters())
+
+    bf16_var_names = to_bf16_var_names if to_bf16_var_names else set()
+    var_scope = scope if scope else global_scope()
+    for param in all_parameters:
+        if param.name in bf16_var_names:
+            _logger.debug(f"---- cast {param.name} to bf16 dtype ----")
+            param_t = var_scope.find_var(param.name).get_tensor()
+            data = np.array(param_t)
+            param_t.set(convert_float_to_uint16(data), place)
+
+
+def rewrite_program_bf16(main_prog, amp_lists=None):
+    """
+    Traverse all ops in current block and insert cast op according to
+    which set current op belongs to.
+
+    1. When an op belongs to the fp32 list, add it to fp32 set
+    2. When an op belongs to the bf16 list, add it to bf16 set
+    3. When an op belongs to the gray list. If one
+       of its inputs is the output of fp32 set op or fp32 list op,
+       add it to fp32 set. If all of its previous ops are not fp32
+       op and one of its inputs is the output of bf16 set op or
+       bf16 list op, add it to bf16 set.
+    4. When an op isn't in the lists, add it to fp32 op set.
+    5. Add necessary cast ops to make sure that fp32 set op will be
+       computed in fp32 mode, while bf16 set op will be computed in
+       bf16 mode.
+
+    Args:
+        main_prog (Program): The main program for training.
+    """
+    if amp_lists is None:
+        amp_lists = AutoMixedPrecisionListsBF16()
+    block = main_prog.global_block()
+    ops = block.ops
+    bf16_op_set = set()
+    fp32_op_set = set()
+    for op in ops:
+        # NOTE(zhiqiu): 'create_py_reader' and 'read' is used in non-iterable DataLoader,
+        # we don't need to handle reader op and the input of 'create_py_reader' is not
+        # in block, which may result in errors.
+        # See GeneratorLoader._init_non_iterable() for details.
+        if op.type == 'create_py_reader' or op.type == 'read':
+            continue
+
+        if amp_lists.fp32_varnames is not None and _is_in_fp32_varnames(
+            op, amp_lists
+        ):
+            fp32_op_set.add(op)
+            continue
+
+        if op.type in amp_lists.fp32_list:
+            fp32_op_set.add(op)
+        elif op.type in amp_lists.bf16_list:
+            bf16_op_set.add(op)
+        elif op.type in amp_lists.gray_list:
+            is_fp32_op = False
+            is_bf16_op = False
+            for in_name in op.input_names:
+                # if this op has inputs
+                if in_name:
+                    for in_var_name in op.input(in_name):
+                        in_var = block.var(in_var_name)
+                        # this in_var isn't the output of other op
+                        if in_var.op is None:
+                            continue
+                        elif in_var.op is op:
+                            prev_op = find_true_prev_op(ops, op, in_var_name)
+                            if prev_op is None:
+                                continue
+                        else:
+                            prev_op = in_var.op
+                        # if it's one of inputs
+                        if (
+                            prev_op in fp32_op_set
+                            or prev_op.type in amp_lists.fp32_list
+                        ):
+                            is_fp32_op = True
+                        elif (
+                            prev_op in bf16_op_set
+                            or prev_op.type in amp_lists.bf16_list
+                        ):
+                            is_bf16_op = True
+            if is_fp32_op:
+                fp32_op_set.add(op)
+            elif is_bf16_op:
+                bf16_op_set.add(op)
+            else:
+                pass
+        else:
+            # For numerical safe, we apply fp32 computation on ops that
+            # are not determined which list they should stay.
+            fp32_op_set.add(op)
+
+    idx = 0
+    while idx < len(ops):
+        op = ops[idx]
+        num_cast_ops = 0
+        if op in fp32_op_set:
+            num_cast_ops = _insert_cast_op(
+                block,
+                op,
+                idx,
+                core.VarDesc.VarType.BF16,
+                core.VarDesc.VarType.FP32,
+            )
+        elif op in bf16_op_set:
+            if (
+                op.has_attr('dtype')
+                and op.attr('dtype') == core.VarDesc.VarType.FP32
+            ):
+                op._set_attr('dtype', core.VarDesc.VarType.BF16)
+
+            num_cast_ops = _insert_cast_op(
+                block,
+                op,
+                idx,
+                core.VarDesc.VarType.FP32,
+                core.VarDesc.VarType.BF16,
+            )
+        else:
+            pass
+
+        idx += num_cast_ops + 1

python/paddle/static/amp/bf16/decorator.py

@@ -0,0 +1,343 @@
+# Copyright (c) 2021 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 types
+import warnings
+
+import paddle
+from paddle.base import core, default_main_program, program_guard, unique_name
+
+from .amp_lists import AutoMixedPrecisionListsBF16
+from .amp_utils import (
+    cast_model_to_bf16,
+    cast_parameters_to_bf16,
+    rewrite_program_bf16,
+)
+
+
+class OptimizerWithMixedPrecision:
+    """
+    Optimizer with mixed-precision (MP) training. This is a wrapper of a common
+    optimizer, plus the support of mixed-precision pre-training. The object
+    of this class almost has the same behavior as the common optimizer, with the
+    methods `minimize()`, `backward()`, `apply_gradients()` implemented.
+    Additionally, it enables the MP training automatically, i.e, the creation
+    and maintenance of master parameters, scaling of loss, etc.
+
+    Args:
+        optimizer (Optimizer): A common Optimizer object.
+        amp_lists (CustomOpLists): An CustomOpLists object.
+        use_pure_bf16(bool): Whether to use the pure bf16 training.
+        use_bf16_guard(bool): Whether to use `bf16_guard` when constructing the program.
+
+    """
+
+    def __init__(self, optimizer, amp_lists, use_pure_bf16, use_bf16_guard):
+        self._optimizer = optimizer
+        self._amp_lists = amp_lists
+        self._param_grads = None
+        self._train_program = None
+
+        self._learning_rate = optimizer._learning_rate
+        self._learning_rate_map = optimizer._learning_rate_map
+        self._use_pure_bf16 = use_pure_bf16
+        self._use_bf16_guard = use_bf16_guard
+        self._to_bf16_var_names = None
+
+    def _init_amp_var(self):
+        # Ensure the data type of learning rate vars is float32 (same as the
+        # master parameter dtype)
+        if isinstance(self._optimizer._learning_rate, float):
+            self._optimizer._learning_rate_map[default_main_program()] = (
+                paddle.static.create_global_var(
+                    name=unique_name.generate("learning_rate"),
+                    shape=[1],
+                    value=float(self._optimizer._learning_rate),
+                    dtype='float32',
+                    persistable=True,
+                )
+            )
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        """
+        Backward propagation or auto differentiation for gradients' computation.
+
+        Args:
+            loss (Variable): The loss Variable to minimize.
+            startup_program (Program|None): The startup Program for initializing
+                                       parameters in `parameter_list`.
+            parameter_list (list|None): A list of Variables to update.
+            no_grad_set (set|None): A set of Variables should be ignored.
+            callbacks (list|None): A list of callable objects to run when appending
+                                   backward operator for one parameter.
+
+        Returns:
+            A list of (param, grad), which is a tuple of a parameter and its
+            gradient respectively, and the scaled loss.
+        """
+        train_program = loss.block.program
+        self._train_program = train_program
+
+        with program_guard(self._train_program, startup_program):
+            self._init_amp_var()
+
+            if self._use_pure_bf16:
+                self._to_bf16_var_names = cast_model_to_bf16(
+                    self._train_program,
+                    startup_program,
+                    self._amp_lists,
+                    self._use_bf16_guard,
+                )
+            else:
+                rewrite_program_bf16(self._train_program, self._amp_lists)
+
+            if loss.dtype != core.VarDesc.VarType.FP32:
+                loss = loss.astype('float32')
+
+            params_grads = self._optimizer.backward(
+                loss, startup_program, parameter_list, no_grad_set, callbacks
+            )
+        return params_grads
+
+    def amp_init(
+        self, place, scope=None, test_program=None, use_bf16_test=False
+    ):
+        """
+        Init the amp training, such as cast fp32 parameters to bf16 type.
+
+        Args:
+            place(CPUPlace): place is used to initialize
+                bf16 parameters with fp32 values.
+            scope(Scope): The scope is used to find fp32 parameters.
+            test_program(Program): The program is used for testing.
+            use_bf16_test(bool): Whether to use bf16 testing.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
+                >>> import numpy as np
+                >>> import paddle
+                >>> import paddle.nn.functional as F
+                >>> paddle.enable_static()
+
+                >>> def run_example_code():
+                ...     place = paddle.CPUPlace()
+                ...     exe = paddle.static.Executor(place)
+                ...     data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
+                ...     conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
+                ...     # 1) Use bf16_guard to control the range of bf16 kernels used.
+                ...     with paddle.static.amp.bf16.bf16_guard():
+                ...         bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
+                ...         pool = F.max_pool2d(bn, kernel_size=2, stride=2)
+                ...         hidden = paddle.static.nn.fc(pool, size=10)
+                ...         loss = paddle.mean(hidden)
+                ...     # 2) Create the optimizer and set `multi_precision` to True.
+                ...     # Setting `multi_precision` to True can avoid the poor accuracy
+                ...     # or the slow convergence in a way.
+                ...     optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
+                ...     # 3) These ops in `custom_black_list` will keep in the float32 computation type.
+                ...     amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
+                ...     # 4) The entry of Paddle AMP.
+                ...     # Enable pure bf16 training by setting `use_pure_bf16` to True.
+                ...     optimizer = paddle.static.amp.bf16.decorate_bf16(
+                ...         optimizer,
+                ...         amp_list,
+                ...         use_pure_bf16=True,
+                ...     )
+                ...     # If you don't use the default_startup_program(), you should pass
+                ...     # your defined `startup_program` into `minimize`.
+                ...     optimizer.minimize(loss)
+                ...     exe.run(paddle.static.default_startup_program())
+                ...     # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
+                ...     # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
+                ...     optimizer.amp_init(place, scope=paddle.static.global_scope())
+
+                >>> run_example_code()
+
+        """
+        assert self._train_program is not None, (
+            "Please call the minimize method first."
+        )
+        if self._use_pure_bf16:
+            cast_parameters_to_bf16(
+                place, self._train_program, scope, self._to_bf16_var_names
+            )
+        if test_program is not None:
+            if self._use_pure_bf16:
+                cast_model_to_bf16(
+                    test_program,
+                    amp_lists=self._amp_lists,
+                    use_bf16_guard=self._use_bf16_guard,
+                )
+            elif use_bf16_test:
+                rewrite_program_bf16(test_program, amp_lists=self._amp_lists)
+
+    def apply_gradients(self, params_grads):
+        """
+        Apply gradients.
+
+        Args:
+            params_grads (list): A list of params.
+
+        Returns:
+            A list of optimize operators.
+        """
+
+        return self._optimizer.apply_gradients(params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        program = loss.block.program
+        with program_guard(program, startup_program):
+            optimize_ops = self.apply_gradients(params_grads)
+        return optimize_ops
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        """
+        Perform optimization by minimizing the given loss.
+
+        Args:
+            loss (Variable): The loss Variable.
+            startup_program (Program): startup_program for initializing parameters
+                in `parameter_list`.
+            parameter_list (list): list of Variables to update.
+            no_grad_set (set|None): set of Variables should be ignored.
+
+        Returns:
+            The scaled loss by scaling factor, the list of optimize ops, and a
+            list of scaled parameters and gradients.
+        """
+        opt_dict = self._optimizer.__class__.__dict__
+        if 'minimize' in opt_dict and isinstance(
+            opt_dict['minimize'], types.FunctionType
+        ):
+            warnings.warn(
+                "The decorated optimizer has its own `minimize` method, but it will not be executed."
+            )
+
+        params_grads = self.backward(
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        optimize_ops = self.apply_optimize(loss, startup_program, params_grads)
+
+        return optimize_ops, params_grads
+
+
+def decorate_bf16(
+    optimizer, amp_lists=None, use_pure_bf16=False, use_bf16_guard=None
+):
+    """
+    Decorate the given optimizer to adapt to the mixed-precision training.
+
+    Args:
+        optimizer(Optimizer): A common Optimizer.
+        amp_lists (CustomOpLists): An CustomOpLists object.
+        use_pure_bf16(bool): Whether to use the pure bf16 training. Default False.
+        use_bf16_guard(bool): Whether to use `bf16_guard` when constructing the program.
+                           Default None, which means that its value equals to `use_pure_bf16`.
+
+    Returns:
+        An optimizer acting like a normal one but with mixed-precision training
+        enabled.
+
+    Examples:
+        .. code-block:: pycon
+            :name: example-1
+
+            # fp32&bf16 list based strategy example
+            >>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
+            >>> import paddle
+            >>> import paddle.static as static
+
+            >>> paddle.enable_static()
+
+            >>> data = static.data(name='X', shape=[None, 1], dtype='float32')
+            >>> hidden = static.nn.fc(x=data, size=10)
+            >>> loss = paddle.mean(hidden)
+            >>> optimizer = paddle.optimizer.Adam(learning_rate=0.001)
+
+            >>> mp_optimizer = static.amp.bf16.decorate_bf16(optimizer=optimizer)
+
+            >>> ops, param_grads = mp_optimizer.minimize(loss)
+
+
+
+
+        .. code-block:: pycon
+            :name: example-2
+
+            # pure bf16 training example
+            >>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
+            >>> import numpy as np
+            >>> import paddle
+            >>> import paddle.nn.functional as F
+
+            >>> def run_example_code():
+            ...     place = paddle.CPUPlace()
+            ...     exe = paddle.static.Executor(place)
+            ...     data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
+            ...     conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
+            ...     # 1) Use bf16_guard to control the range of bf16 kernels used.
+            ...     with paddle.static.amp.bf16.bf16_guard():
+            ...         bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
+            ...         pool = F.max_pool2d(bn, kernel_size=2, stride=2)
+            ...         hidden = paddle.static.nn.fc(pool, size=10)
+            ...         loss = paddle.mean(hidden)
+            ...     # 2) Create the optimizer and set `multi_precision` to True.
+            ...     # Setting `multi_precision` to True can avoid the poor accuracy
+            ...     # or the slow convergence in a way.
+            ...     optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
+            ...     # 3) These ops in `custom_black_list` will keep in the float32 computation type.
+            ...     amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
+            ...     # 4) The entry of Paddle AMP.
+            ...     # Enable pure bf16 training by setting `use_pure_bf16` to True.
+            ...     optimizer = paddle.static.amp.bf16.decorate_bf16(
+            ...         optimizer,
+            ...         amp_list,
+            ...         use_pure_bf16=True,
+            ...     )
+            ...     # If you don't use the default_startup_program(), you should pass
+            ...     # your defined `startup_program` into `minimize`.
+            ...     optimizer.minimize(loss)
+            ...     exe.run(paddle.static.default_startup_program())
+            ...     # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
+            ...     # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
+            ...     optimizer.amp_init(place, scope=paddle.static.global_scope())
+            >>> run_example_code()
+
+    """
+    if amp_lists is None:
+        amp_lists = AutoMixedPrecisionListsBF16()
+
+    if use_bf16_guard is None:
+        use_bf16_guard = use_pure_bf16
+
+    mp_optimizer = OptimizerWithMixedPrecision(
+        optimizer, amp_lists, use_pure_bf16, use_bf16_guard
+    )
+
+    return mp_optimizer

python/paddle/static/amp/debugging.py

@@ -0,0 +1,279 @@
+# 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 copy
+import logging
+
+import paddle
+from paddle.base.log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class OperatorStatsUnit:
+    def __init__(self):
+        self.op_type = None
+        self.fp32_calls = 0
+        self.fp16_calls = 0
+        self.bf16_calls = 0
+        self.other_calls = 0
+
+    def update(self, dtype):
+        if dtype is None:
+            self.other_calls = self.other_calls + 1
+        else:
+            if dtype == paddle.float32:
+                self.fp32_calls = self.fp32_calls + 1
+            elif dtype == paddle.float16:
+                self.fp16_calls = self.fp16_calls + 1
+            elif dtype == paddle.bfloat16:
+                self.bf16_calls = self.bf16_calls + 1
+            else:
+                self.other_calls = self.other_calls + 1
+
+    def addto(self, another):
+        self.fp32_calls += another.fp32_calls
+        self.fp16_calls += another.fp16_calls
+        self.bf16_calls += another.bf16_calls
+        self.other_calls += another.other_calls
+
+    def convert_to_list(self):
+        return [
+            self.fp16_calls,
+            self.bf16_calls,
+            self.fp32_calls,
+            self.other_calls,
+        ]
+
+
+def _is_floating_point(dtype):
+    if dtype in [
+        paddle.base.core.VarDesc.VarType.FP64,
+        paddle.base.core.VarDesc.VarType.FP32,
+        paddle.base.core.VarDesc.VarType.FP16,
+        paddle.base.core.VarDesc.VarType.BF16,
+    ]:
+        return True
+    else:
+        return False
+
+
+def _get_var_dtype_from_block(block, op, arg_name, is_input):
+    var_names = op.input(arg_name) if is_input else op.output(arg_name)
+    assert isinstance(var_names, list)
+    if len(var_names) == 0:
+        return None
+
+    var_name = var_names[0]
+    try:
+        var = block._var_recursive(var_name)
+        return var.dtype
+    except:
+        _logger.warning(
+            "Operator < {} > gets {} < {} : {} > error!".format(
+                op.type, "input" if is_input else "output", arg_name, var_name
+            )
+        )
+        return None
+
+
+def _extract_compute_dtype(op, block):
+    var_name = None
+    compute_dtype = None
+    for in_name in op.input_names:
+        var_dtype = _get_var_dtype_from_block(block, op, in_name, True)
+        if var_dtype is None:
+            continue
+
+        if compute_dtype is None:
+            compute_dtype = var_dtype
+        else:
+            if compute_dtype != var_dtype:
+                if _is_floating_point(compute_dtype) and _is_floating_point(
+                    var_dtype
+                ):
+                    _logger.warning(
+                        f"Operator < {op.type} > has different input data types, input_names = {op.input_names}, output_names = {op.output_names}."
+                    )
+                elif _is_floating_point(var_dtype):
+                    # When there are multiple inputs, such as embedding
+                    # (ids is integer, w is floating-point), the kernel
+                    # dtype is normally decided by the input of floating-point.
+                    compute_dtype = var_dtype
+
+    for out_name in op.output_names:
+        var_dtype = _get_var_dtype_from_block(block, op, out_name, False)
+        if var_dtype is None:
+            continue
+
+        if compute_dtype is None:
+            # Kernel dtype is mostly decided by the input's dtype.
+            # When the operator has no input, it might have an attr
+            # such as dtype to specify the output's dtype.
+            compute_dtype = var_dtype
+        else:
+            if compute_dtype != var_dtype:
+                if _is_floating_point(compute_dtype) and _is_floating_point(
+                    var_dtype
+                ):
+                    _logger.warning(
+                        f"Operator < {op.type} > has different input / output data types, input_names = {op.input_names}, output_names = {op.output_names}."
+                    )
+    return compute_dtype
+
+
+def _merge_op_stats(op_stats_list):
+    merged_op_stats_dict = {}
+    for each_op_stats_dict in op_stats_list:
+        for op_type, unit in each_op_stats_dict.items():
+            if merged_op_stats_dict.get(op_type, None) is None:
+                merged_op_stats_dict[op_type] = copy.copy(unit)
+            else:
+                merged_op_stats_dict[op_type].addto(unit)
+    return merged_op_stats_dict
+
+
+def _get_op_stats_list(program):
+    def _is_special_ops_with_input_x(op_type):
+        # operators have input X and have inputs different dtypes.
+        special_op_list = ['cast', 'batch_norm', 'instance_norm', 'layer_norm']
+        if op_type in special_op_list:
+            return True
+        if op_type.replace("_grad", "") in special_op_list:
+            return True
+        return False
+
+    op_stats_list = []
+    for block in program.blocks:
+        block_op_stats_dict = {}
+        for op in block.ops:
+            if block_op_stats_dict.get(op.type, None) is None:
+                unit = OperatorStatsUnit()
+                block_op_stats_dict[op.type] = unit
+            else:
+                unit = block_op_stats_dict[op.type]
+
+            if op.type in [
+                'create_py_reader',
+                'read',
+                'create_double_buffer_reader',
+            ]:
+                compute_dtype = None
+            elif _is_special_ops_with_input_x(op.type):
+                # Not check the input and output dtype difference for this operators.
+                compute_dtype = _get_var_dtype_from_block(block, op, 'X', True)
+            elif "Param" in op.input_names:
+                # Specify compute_dtype for optimizers.
+                compute_dtype = _get_var_dtype_from_block(
+                    block, op, 'Param', True
+                )
+            else:
+                compute_dtype = _extract_compute_dtype(op, block)
+            unit.update(dtype=compute_dtype)
+        op_stats_list.append(block_op_stats_dict)
+    return op_stats_list
+
+
+def collect_operator_stats(program=None, print_subblocks=False):
+    """
+    Collect the number of operators for different data types through parsing
+    the program. The statistical data are categorized according to four data
+    types, namely float32, float16, bfloat16 and others.
+
+    Args:
+        program(Program, optional): The program to parse. Default None, and the default main_program will be parsed.
+        print_subblocks(bool, optional): Whether to print the operator stats for each subblock. Default False.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> class SimpleConvNet(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.conv = paddle.nn.Conv2D(in_channels=1, out_channels=6, kernel_size=3)
+            ...         self.linear = paddle.nn.Linear(in_features=26, out_features=10)
+            ...
+            ...     def forward(self, x):
+            ...         out = self.conv(x)
+            ...         out = paddle.nn.functional.relu(out)
+            ...         out = self.linear(out)
+            ...         out = paddle.nn.functional.softmax(out)
+            ...         return out
+
+            >>> main_program = paddle.static.Program()
+            >>> startup_program = paddle.static.Program()
+            >>> with paddle.utils.unique_name.guard():
+            ...     with paddle.static.program_guard(main_program, startup_program):
+            ...         model = SimpleConvNet()
+            ...         x = paddle.static.data(name='input', shape=[None, 1, 28, 28], dtype='float32')
+            ...         out = model(x)
+            ...         loss = paddle.mean(out)
+            ...         optimizer = paddle.optimizer.AdamW()
+            ...         optimizer = paddle.static.amp.decorate(optimizer)
+            ...         optimizer.minimize(loss)
+            >>> paddle.static.amp.debugging.collect_operator_stats(main_program)
+            <------------------------------------------------ op list of all blocks ------------------------------------------------->
+            <------------------------------------------------------- op list -------------------------------------------------------->
+            <--------------- Op Name ---------------- | -- FP16 Calls --- | -- BF16 Calls --- | --- FP32 Calls--- | -- Other Calls -->
+            adamw                                   |  0                |  0                |  4                |  0
+            cast                                    |  5                |  0                |  6                |  0
+            check_finite_and_unscale                |  0                |  0                |  1                |  0
+            conv2d                                  |  1                |  0                |  0                |  0
+            conv2d_grad                             |  1                |  0                |  0                |  0
+            elementwise_add                         |  2                |  0                |  0                |  0
+            elementwise_add_grad                    |  2                |  0                |  0                |  0
+            elementwise_mul                         |  0                |  0                |  1                |  0
+            elementwise_mul_grad                    |  0                |  0                |  1                |  0
+            fill_constant                           |  0                |  0                |  1                |  0
+            matmul_v2                               |  1                |  0                |  0                |  0
+            matmul_v2_grad                          |  1                |  0                |  0                |  0
+            memcpy                                  |  0                |  0                |  0                |  1
+            reduce_mean                             |  0                |  0                |  1                |  0
+            reduce_mean_grad                        |  0                |  0                |  1                |  0
+            relu                                    |  1                |  0                |  0                |  0
+            relu_grad                               |  1                |  0                |  0                |  0
+            reshape2                                |  0                |  0                |  1                |  0
+            reshape2_grad                           |  0                |  0                |  1                |  0
+            softmax                                 |  0                |  0                |  1                |  0
+            softmax_grad                            |  0                |  0                |  1                |  0
+            update_loss_scaling                     |  0                |  0                |  1                |  0
+            <----------------------------------------------------- op count: 22 ----------------------------------------------------->
+    """
+
+    def _convert_to_list(op_stats_unit_dict):
+        for key, value in op_stats_unit_dict.items():
+            op_stats_unit_dict[key] = value.convert_to_list()
+        return op_stats_unit_dict
+
+    if program is None:
+        program = paddle.static.default_main_program()
+
+    op_stats_list = _get_op_stats_list(program)
+    merged_op_stats = _merge_op_stats(op_stats_list)
+    if print_subblocks and len(op_stats_list) > 1:
+        for i in range(len(op_stats_list)):
+            print("<{:-^120}>".format(" op list of block " + str(i) + " "))
+            paddle.amp.debugging._print_operator_stats(
+                _convert_to_list(op_stats_list[i])
+            )
+    print("<{:-^120}>".format(" op list of all blocks "))
+    paddle.amp.debugging._print_operator_stats(
+        _convert_to_list(merged_op_stats)
+    )

python/paddle/static/amp/fp16_lists.py

@@ -0,0 +1,273 @@
+#   Copyright (c) 2019 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 copy
+import logging
+
+import paddle
+from paddle.amp.amp_lists import (
+    EXTRA_BLACK_LIST,
+    FP16_BLACK_LIST,
+    FP16_WHITE_LIST,
+)
+from paddle.base import core
+from paddle.base.log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+black_list = FP16_BLACK_LIST
+_extra_black_list = EXTRA_BLACK_LIST
+white_list = FP16_WHITE_LIST
+
+
+def check_amp_dtype(dtype):
+    """
+    Check amp_dtype: float16 or bfloat16
+    """
+    if isinstance(dtype, str):
+        dtype = dtype.lower()
+    if dtype not in ['float16', 'bfloat16']:
+        raise ValueError(
+            "If enable AMP, dtype should be 'float16' or 'bfloat16'."
+        )
+    return dtype
+
+
+def get_low_precision_vartype(dtype):
+    if isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
+        return dtype
+    elif isinstance(dtype, str):
+        dtype = dtype.lower()
+        if dtype == "float16":
+            var_type = core.VarDesc.VarType.FP16
+        elif dtype == "bfloat16":
+            var_type = core.VarDesc.VarType.BF16
+        else:
+            raise ValueError(
+                "If enable AMP, dtype should be 'float16' or 'bfloat16'."
+            )
+        return var_type
+    else:
+        raise TypeError(
+            f"The type of dtype is expected to be string or core.VarDesc.VarType, but received {type(dtype)}."
+        )
+
+
+def get_low_precision_dtypestr(dtype):
+    if isinstance(dtype, str):
+        return check_amp_dtype(dtype)
+    elif isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
+        if dtype == paddle.float16:
+            return "float16"
+        elif dtype == paddle.bfloat16:
+            return "bfloat16"
+        else:
+            raise ValueError(
+                "If enable AMP, dtype should be core.VarDesc.VarType.FP16 or core.VarDesc.VarType.BF16."
+            )
+    else:
+        raise TypeError(
+            f"The type of dtype is expected to be string or core.VarDesc.VarType, but received {type(dtype)}."
+        )
+
+
+def _get_sys_unsupported_list(dtype):
+    var_type = get_low_precision_vartype(dtype)
+
+    # The set of ops that don't support fp16 calculation
+    device = None
+    if core.is_compiled_with_xpu():
+        device = 'XPU'
+    elif isinstance(
+        paddle.framework._current_expected_place(), paddle.CustomPlace
+    ):
+        device = paddle.framework._current_expected_place().get_device_type()
+    else:
+        device = 'GPU'
+    all_ops, _, sys_unsupported_list = core.op_supported_infos(device, var_type)
+
+    # sys_unsupported_list will include the following ops.
+    supported_fp16_list = {
+        "conditional_block",
+        "conditional_block_infer",
+        "select_input",
+        "while",
+        "cast",
+        "tensor_array_to_tensor",
+        "lod_array_length",
+        "write_to_array",
+    }
+    sys_unsupported_list -= supported_fp16_list
+
+    return device, sys_unsupported_list, all_ops
+
+
+def _get_unsupported_list(dtype):
+    # The set of ops that don't support fp16 calculation
+    _, _sys_unsupported_list, _sys_all_list = _get_sys_unsupported_list(dtype)
+    return _sys_unsupported_list, _sys_all_list
+
+
+# The three sets listed below are changed dynamically. They don't contain all
+# paddle ops currently.
+
+# The set of ops that support fp16 calculation and are considered numerically-
+# safe and performance-critical. These ops are always converted to fp16.
+
+_only_supported_fp16_list = {'resnet_unit', 'fused_bn_add_activation'}
+
+
+def _get_white_list(dtype):
+    white_list_for_dtype = copy.copy(FP16_WHITE_LIST)
+    if dtype == 'float16':
+        white_list_for_dtype = white_list_for_dtype | _only_supported_fp16_list
+    return white_list_for_dtype
+
+
+def _get_black_list():
+    _black_list = copy.copy(FP16_BLACK_LIST)
+    _black_list = _black_list | EXTRA_BLACK_LIST
+    return _black_list
+
+
+class AutoMixedPrecisionLists:
+    """
+    AutoMixedPrecisionLists is a class for black/white list. It can update
+    pre-defined black list and white list according to users' custom black
+    white lists. The lists are used for an algorithm which determines op's
+    execution mode (fp32, fp16 or bf16).
+
+    Args:
+        custom_white_list (set): Users' custom white list.
+        custom_black_list (set): Users' custom black list.
+        custom_black_varnames (set): Users' custom black variables' names.
+        dtype (str): the low precision dtype, which can be set to 'float16' or 'bfloat16'.
+    """
+
+    def __init__(
+        self,
+        custom_white_list=None,
+        custom_black_list=None,
+        custom_black_varnames=None,
+        dtype="float16",
+    ):
+        self.amp_dtype = check_amp_dtype(dtype)
+        self._custom_white_list = custom_white_list
+        self._custom_black_list = custom_black_list
+        self.white_list = copy.copy(_get_white_list(self.amp_dtype))
+        self.black_list = copy.copy(_get_black_list())
+        self.gray_list = copy.copy(gray_list)
+        unsupported_list, sys_all_list = _get_unsupported_list(self.amp_dtype)
+        self.unsupported_list = copy.copy(unsupported_list)
+        self.all_list = copy.copy(sys_all_list)
+        self.black_varnames = copy.copy(custom_black_varnames)
+        self._update_list()
+
+    def _update_list(self):
+        """
+        Update black and white list according to users' custom list.
+        """
+        _logger.debug(f"---- custom_white_list {self._custom_white_list} ---- ")
+        _logger.debug(f"---- custom_black_list {self._custom_black_list} ---- ")
+        _logger.debug(f"---- custom_black_varnames {self.black_varnames} ---- ")
+        if self._custom_white_list and self._custom_black_list:
+            for op_name in self._custom_white_list:
+                if op_name in self._custom_black_list:
+                    raise ValueError(
+                        f"The given custom_white_list overlaps custom_black_list with < {op_name} >!"
+                    )
+        if self._custom_white_list:
+            for op_name in self._custom_white_list:
+                if op_name in self.black_list:
+                    self.black_list.remove(op_name)
+                elif op_name in self.gray_list:
+                    self.gray_list.remove(op_name)
+                self.white_list.add(op_name)
+        if self._custom_black_list:
+            for op_name in self._custom_black_list:
+                if op_name in self.white_list:
+                    self.white_list.remove(op_name)
+                elif op_name in self.gray_list:
+                    self.gray_list.remove(op_name)
+                self.black_list.add(op_name)
+                self.unsupported_list.add(op_name)
+        device, sys_unsupported_list, _ = _get_sys_unsupported_list(
+            self.amp_dtype
+        )
+        actual_unsupported_list = []
+        for op_name in sys_unsupported_list:
+            if op_name in self.white_list:
+                actual_unsupported_list.append(op_name)
+        if len(actual_unsupported_list) > 0:
+            _logger.warning(
+                f"On current {device}, {self.amp_dtype} is not supported for operators < {actual_unsupported_list} > in white_list!"
+            )
+
+
+# This set contains two types of ops. All ops supported fp16 calculation. One
+# of two types is considered numerically-safe, but may be made unsafe by an
+# upstream blacklist op. Another type do not have numerically-significant
+# effects, like stack, flatten2.
+gray_list = {
+    'elementwise_add',
+    'elementwise_sub',
+    'elementwise_mul',
+    'elementwise_div',
+    'elementwise_max',
+    'elementwise_min',
+    'elementwise_pow',
+    'elementwise_mod',
+    'elementwise_floordiv',
+    'batch_norm',
+    'layer_norm',
+    'tanh',
+    'sigmoid',
+    'top_k',
+    'pool2d',
+    'pool3d',
+    'dropout',
+    'relu',
+    'relu6',
+    'leaky_relu',
+    'soft_relu',
+    'flatten2',
+    'stack',
+    'unstack',
+    'uniform_random',
+    'uniform_random_batch_size_like',
+    'gaussian_random',
+    'slice',
+    'rank',
+    'scale',
+    'transpose2',
+    'reshape2',
+    'gather',
+    'fill_constant',
+    'get_tensor_from_selected_rows',
+    'sign',
+    'cast',
+    'fused_bn_add_activation',
+    'c_identity',
+    'c_concat',
+    'all_reduce',
+    'concat',
+    'split',
+    'fused_feedforward',
+    'fused_attention',
+    'fused_multi_transformer',
+}
+
+CustomOpLists = AutoMixedPrecisionLists

python/paddle/static/amp/function_overload.py

@@ -0,0 +1,142 @@
+# 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.
+
+# The implementation refers to https://arpitbhayani.me/blogs/function-overloading.
+# Note: it is customed for paddle.static.amp.decorate function.
+
+import inspect
+import logging
+from enum import Enum
+
+from paddle.base.log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class FunctionType(Enum):
+    FP16_ONLY = 0
+    COMMON = 1
+
+
+class Function:
+    """
+    Function is a wrap over standard python function
+    An instance of this Function class is also callable
+    just like the python function that it wrapped.
+    When the instance is "called" like a function it fetches
+    the function to be invoked from the virtual namespace and then
+    invokes the same.
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(self, *args, **kwargs):
+        """
+        Overriding the __call__ function which makes the
+        instance callable.
+        """
+        # fetching the function to be invoked from the virtual namespace
+        # through the arguments.
+        fn = Namespace.get_instance().get(*args, **kwargs)
+        # invoking the wrapped function and returning the value.
+        return fn(*args, **kwargs)
+
+
+class Namespace:
+    """
+    Namespace is the singleton class that is responsible
+    for holding all the functions.
+    """
+
+    __instance = None
+
+    def __init__(self):
+        if self.__instance is None:
+            self.function_map = {}
+            Namespace.__instance = self
+        else:
+            raise Exception("cannot instantiate Namespace again.")
+
+    @staticmethod
+    def get_instance():
+        if Namespace.__instance is None:
+            Namespace()
+        return Namespace.__instance
+
+    def register(self, fn, key):
+        """
+        Register the function in the virtual namespace and return
+        an instance of callable Function that wraps the function fn.
+
+        Args:
+            fn (function): the native python function handle.
+            key (FunctionType): the specified type.
+        """
+        assert isinstance(key, FunctionType), (
+            f"The type of  key is expected to be FunctionType, but received {type(key)}."
+        )
+        func = Function(fn)
+        self.function_map[key] = fn
+        return func
+
+    def get(self, *args, **kwargs):
+        """
+        Get the matching function from the virtual namespace according to the actual arguments.
+        Return None if it did not find any matching function.
+        """
+        _logger.debug(f"get function: args={args}, kwargs={kwargs}")
+        satisfied_function_keys = set(self.function_map.keys())
+        num_actual_args = len(args) + len(kwargs)
+        for func_key in self.function_map.keys():
+            if func_key not in satisfied_function_keys:
+                continue
+            fn = self.function_map[func_key]
+            specs = inspect.getfullargspec(fn)
+            if len(specs) < len(args) + len(kwargs):
+                # Remove the not satisfied function according to the number of actual arguments.
+                _logger.debug(
+                    f"fn={fn} (key={func_key}) is not satisfied and removed."
+                )
+                satisfied_function_keys.remove(func_key)
+                continue
+            if len(kwargs) > 0:
+                # Remove the not satisfied function according to argument keys in kwargs.
+                for arg_name, value in kwargs.items():
+                    if arg_name not in specs.args:
+                        _logger.debug(
+                            f"fn={fn} (key={func_key}) is not satisfied and removed."
+                        )
+                        satisfied_function_keys.remove(func_key)
+                        break
+        if len(satisfied_function_keys) == 1:
+            key = next(iter(satisfied_function_keys))
+        elif len(args) >= 3 and isinstance(args[2], float):
+            key = FunctionType.FP16_ONLY
+        else:
+            key = FunctionType.COMMON
+        return self.function_map.get(key)
+
+
+def overload(key):
+    """overload is the decorator that wraps the function
+    and returns a callable object of type Function.
+    """
+
+    def decorator(fn):
+        return Namespace.get_instance().register(fn, key)
+
+    return decorator

python/paddle/static/input.py

@@ -0,0 +1,481 @@
+# Copyright (c) 2020 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 os
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import Self
+
+import paddle
+from paddle.base import Variable, core
+from paddle.base.data_feeder import check_type
+from paddle.base.framework import (
+    convert_nptype_to_datatype_or_vartype,
+    in_pir_mode,
+    static_only,
+)
+from paddle.base.layer_helper import LayerHelper
+from paddle.base.libpaddle import DataType
+from paddle.base.libpaddle.pir import (
+    get_current_insertion_point,
+    set_insertion_point,
+)
+
+from ..base.variable_index import _setitem_static
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import (
+        DTypeLike,
+        ShapeLike,
+        Size1,
+        TensorIndex,
+        TensorLike,
+    )
+
+__all__ = []
+
+
+def evaluate_flag(val) -> bool:
+    return str(val).lower() not in ('false', 'off', '0', 'none')
+
+
+@static_only
+def data(
+    name: str,
+    shape: ShapeLike,
+    dtype: DTypeLike | None = None,
+    lod_level: int = 0,
+) -> paddle.Tensor:
+    """
+
+    This function creates a variable on the global block. The global variable
+    can be accessed by all the following operators in the graph. The variable
+    is a placeholder that could be fed with input, such as Executor can feed
+    input into the variable. When `dtype` is None, the dtype
+    will get from the global dtype by `paddle.get_default_dtype()`.
+
+    Args:
+        name (str): The name/alias of the variable, see :ref:`api_guide_Name`
+            for more details.
+        shape (list|tuple): List|Tuple of integers declaring the shape. You can
+            set None or -1 at a dimension to indicate the dimension can be of any
+            size. For example, it is useful to set changeable batch size as None or -1.
+        dtype (np.dtype|str, optional): The type of the data. Supported
+            dtype: bool, float16, float32, float64, int8, int16, int32, int64,
+            uint8. Default: None. When `dtype` is not set, the dtype will get
+            from the global dtype by `paddle.get_default_dtype()`.
+        lod_level (int, optional): The LoD level of the DenseTensor. Usually users
+            don't have to set this value. Default: 0.
+
+    Returns:
+        Variable: The global variable that gives access to the data.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("This has diff in xdoctest env")
+            >>> import numpy as np
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            # Creates a variable with fixed size [3, 2, 1]
+            # User can only feed data of the same shape to x
+            # the dtype is not set, so it will set "float32" by
+            # paddle.get_default_dtype(). You can use paddle.get_default_dtype() to
+            # change the global dtype
+            >>> x = paddle.static.data(name='x', shape=[3, 2, 1])
+
+            # Creates a variable with changeable batch size -1.
+            # Users can feed data of any batch size into y,
+            # but size of each data sample has to be [2, 1]
+            >>> y = paddle.static.data(name='y', shape=[-1, 2, 1], dtype='float32')
+
+            >>> z = x + y
+
+            # In this example, we will feed x and y with np-ndarray "1"
+            # and fetch z, like implementing "1 + 1 = 2" in PaddlePaddle
+            >>> feed_data = np.ones(shape=[3, 2, 1], dtype=np.float32)
+
+            >>> exe = paddle.static.Executor(paddle.framework.CPUPlace())
+            >>> out = exe.run(
+            ...     paddle.static.default_main_program(),
+            ...     feed={
+            ...         'x': feed_data,
+            ...         'y': feed_data,
+            ...     },
+            ...     fetch_list=[z.name],
+            ... )
+
+            # np-ndarray of shape=[3, 2, 1], dtype=float32, whose elements are 2
+            >>> print(out)
+            [array([[[2.],
+                    [2.]],
+                   [[2.],
+                    [2.]],
+                   [[2.],
+                    [2.]]], dtype=float32)]
+
+    """
+
+    def _reset_data_op_insertion_point():
+        default_main_program = paddle.pir.core.default_main_program()
+        ops = default_main_program.global_block().ops
+        if len(ops) == 0:
+            return
+        for op in ops:
+            if op.name() != 'pd_op.data':
+                paddle.pir.set_insertion_point(op)
+                return
+
+    helper = LayerHelper('data', **locals())
+    check_type(name, 'name', (bytes, str), 'data')
+    check_type(shape, 'shape', (list, tuple), 'data')
+
+    shape = list(shape)
+    for i in range(len(shape)):
+        if shape[i] is None:
+            shape[i] = -1
+        if isinstance(shape[i], int) and shape[i] < 0 and shape[i] != -1:
+            raise ValueError(
+                f"Only -1 can be used in shape to indicate unknown dimension, but received {shape[i]}"
+            )
+
+    if dtype is None:
+        dtype = paddle.get_default_dtype()
+
+    if core.is_compiled_with_custom_device("iluvatar_gpu") and os.environ.get(
+        'FLAG_FORCE_FLOAT32', ''
+    ).lower() in ['1', 'true', 'on']:
+        dtype_str = dtype if isinstance(dtype, str) else str(dtype)
+        if dtype_str in ('float64', np.float64, 'f8'):
+            import warnings
+
+            warnings.warn(
+                f"Variable '{name}' dtype 'float64' is not supported on iluvatar gpu, "
+                "forcibly using 'float32'.",
+                UserWarning,
+                stacklevel=2,
+            )
+            dtype = 'float32'
+        elif dtype_str in ('complex128', np.complex128, 'c16'):
+            import warnings
+
+            warnings.warn(
+                f"Variable '{name}' dtype 'complex128' is not supported on iluvatar gpu, "
+                "forcibly using 'complex64'.",
+                UserWarning,
+                stacklevel=2,
+            )
+            dtype = 'complex64'
+
+    if in_pir_mode():
+        ir_dtype = dtype
+        if not isinstance(ir_dtype, DataType):
+            ir_dtype = paddle.pir.core.convert_nptype_to_datatype(dtype)
+        prev_insertion_point = get_current_insertion_point()
+        _reset_data_op_insertion_point()
+        out = paddle._pir_ops.data(name, shape, ir_dtype, core.Place())
+        set_insertion_point(prev_insertion_point)
+        return out
+
+    out = helper.create_global_variable(
+        name=name,
+        shape=shape,
+        dtype=dtype,
+        type=core.VarDesc.VarType.DENSE_TENSOR,
+        stop_gradient=True,
+        lod_level=lod_level,
+        is_data=True,
+        need_check_feed=True,
+    )
+
+    is_pir_mode = os.environ.get("FLAGS_enable_pir_in_executor", None)
+    if evaluate_flag(is_pir_mode):
+        helper = LayerHelper('data', **locals())
+        if not isinstance(dtype, core.VarDesc.VarType):
+            dtype = convert_nptype_to_datatype_or_vartype(dtype)
+        helper.append_op(
+            type='data',
+            inputs={},
+            outputs={'out': out},
+            attrs={
+                'shape': shape,
+                'dtype': dtype,
+                'place': 0,
+                'name': name,
+            },
+        )
+    return out
+
+
+class InputSpec:
+    """
+    InputSpec describes the signature information of the model input, such as ``shape`` , ``dtype`` , ``name`` .
+
+    This interface is often used to specify input tensor information of models in high-level API.
+    It's also used to specify the tensor information for each input parameter of the forward function
+    decorated by `@paddle.jit.to_static`.
+
+    Args:
+        shape (tuple(integers)|list[integers]): List|Tuple of integers
+            declaring the shape. You can set "None" or -1 at a dimension
+            to indicate the dimension can be of any size. For example,
+            it is useful to set changeable batch size as "None" or -1.
+        dtype (np.dtype|str, optional): The type of the data. Supported
+            dtype: bool, float16, float32, float64, int8, int16, int32, int64,
+            uint8. Default: float32.
+        name (str): The name/alias of the variable, see :ref:`api_guide_Name`
+            for more details.
+        stop_gradient (bool, optional): A boolean that mentions whether gradient should flow. Default is False, means don't stop calculate gradients.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.static import InputSpec
+
+            >>> input = InputSpec([None, 784], 'float32', 'x')
+            >>> label = InputSpec([None, 1], 'int64', 'label')
+
+            >>> print(input)
+            InputSpec(shape=(-1, 784), dtype=paddle.float32, name=x, stop_gradient=False)
+
+            >>> print(label)
+            InputSpec(shape=(-1, 1), dtype=paddle.int64, name=label, stop_gradient=False)
+    """
+
+    def __init__(
+        self,
+        shape: ShapeLike,
+        dtype: DTypeLike = 'float32',
+        name: str | None = None,
+        stop_gradient: bool = False,
+    ) -> None:
+        # replace `None` in shape  with -1
+        self.shape = self._verify(shape)
+        # convert dtype into united representation
+        if dtype is not None:
+            if isinstance(dtype, (np.dtype, str)):
+                dtype = convert_nptype_to_datatype_or_vartype(dtype)
+
+        self.dtype = dtype
+        self.name = name
+        self.stop_gradient = stop_gradient
+
+    def _create_feed_layer(self):
+        return data(self.name, shape=self.shape, dtype=self.dtype)
+
+    def __repr__(self) -> str:
+        return f'{type(self).__name__}(shape={self.shape}, dtype={self.dtype}, name={self.name}, stop_gradient={self.stop_gradient})'
+
+    @classmethod
+    def from_tensor(cls, tensor: Tensor, name: str | None = None) -> Self:
+        """
+        Generates a InputSpec based on the description of input tensor.
+
+        Args:
+            tensor(Tensor): the source tensor to generate a InputSpec instance
+
+        Returns:
+            A InputSpec instance generated from Tensor.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static import InputSpec
+
+                >>> paddle.disable_static()
+
+                >>> x = paddle.ones([2, 2], dtype="float32")
+                >>> x_spec = InputSpec.from_tensor(x, name='x')
+                >>> print(x_spec)
+                InputSpec(shape=(2, 2), dtype=paddle.float32, name=x, stop_gradient=False)
+
+        """
+        if isinstance(tensor, (Variable, core.eager.Tensor, paddle.pir.Value)):
+            return cls(tensor.shape, tensor.dtype, name or tensor.name)
+        else:
+            raise ValueError(
+                f"Input `tensor` should be a Tensor, but received {type(tensor).__name__}."
+            )
+
+    @classmethod
+    def from_numpy(
+        cls, ndarray: npt.NDArray[Any], name: str | None = None
+    ) -> Self:
+        """
+        Generates a InputSpec based on the description of input np.ndarray.
+
+        Args:
+            tensor(Tensor): the source numpy ndarray to generate a InputSpec instance
+
+        Returns:
+            A InputSpec instance generated from Tensor.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import numpy as np
+                >>> from paddle.static import InputSpec
+
+                >>> x = np.ones([2, 2], np.float32)
+                >>> x_spec = InputSpec.from_numpy(x, name='x')
+                >>> print(x_spec)
+                InputSpec(shape=(2, 2), dtype=paddle.float32, name=x, stop_gradient=False)
+
+        """
+        return cls(ndarray.shape, ndarray.dtype, name)
+
+    def batch(self, batch_size: int | Size1) -> Self:
+        """
+        Inserts `batch_size` in front of the `shape`.
+
+        Args:
+            batch_size(int): the inserted integer value of batch size.
+
+        Returns:
+            The original InputSpec instance by inserting `batch_size` in front of `shape`.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> from paddle.static import InputSpec
+
+                >>> x_spec = InputSpec(shape=[64], dtype='float32', name='x')
+                >>> x_spec.batch(4)
+                >>> print(x_spec)
+                InputSpec(shape=(4, 64), dtype=paddle.float32, name=x, stop_gradient=False)
+
+        """
+        if isinstance(batch_size, (list, tuple)):
+            if len(batch_size) != 1:
+                raise ValueError(
+                    f"Length of batch_size: {batch_size} shall be 1, but received {len(batch_size)}."
+                )
+            batch_size = batch_size[0]
+        elif not isinstance(batch_size, int):
+            raise TypeError(
+                f"type(batch_size) shall be `int`, but received {type(batch_size).__name__}."
+            )
+
+        new_shape = [batch_size, *list(self.shape)]
+        self.shape = tuple(new_shape)
+
+        return self
+
+    def unbatch(self) -> Self:
+        """
+        Removes the first element of `shape`.
+
+        Returns:
+            The original InputSpec instance by removing the first element of `shape` .
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> from paddle.static import InputSpec
+
+                >>> x_spec = InputSpec(shape=[4, 64], dtype='float32', name='x')
+                >>> x_spec.unbatch()
+                >>> print(x_spec)  # InputSpec(shape=(64,), dtype=paddle.float32, name=x)
+                InputSpec(shape=(64,), dtype=paddle.float32, name=x, stop_gradient=False)
+
+        """
+        if len(self.shape) == 0:
+            raise ValueError(
+                "Not support to unbatch a InputSpec when len(shape) == 0."
+            )
+
+        self.shape = self._verify(self.shape[1:])
+        return self
+
+    def _verify(self, shape):
+        """
+        Verifies the input shape and modifies `None` into `-1`.
+        """
+        if not isinstance(shape, (list, tuple)):
+            raise TypeError(
+                f"Type of `shape` in InputSpec should be one of (tuple, list), but received {type(shape).__name__}."
+            )
+
+        for i, ele in enumerate(shape):
+            if ele is not None:
+                if not isinstance(ele, int):
+                    raise ValueError(
+                        f"shape[{i}] should be an `int`, but received `{type(ele).__name__}`:{ele}."
+                    )
+            if ele is None or ele < -1:
+                shape[i] = -1
+
+        return tuple(shape)
+
+    def __hash__(self) -> int:
+        # Note(Aurelius84): `name` is not considered as a field to compute hashkey.
+        # Because it's no need to generate a new program in following cases while using
+        # @paddle.jit.to_static.
+        #
+        # Case 1:
+        #      foo(x_var)
+        #      foo(y_var)
+        #  x_var and y_var hold same shape and dtype, they should share a same program.
+        #
+        #
+        # Case 2:
+        #      foo(x_var)
+        #      foo(x_np)  # x_np is a numpy.ndarray.
+        #  x_var and x_np hold same shape and dtype, they should also share a same program.
+        return hash((tuple(self.shape), self.dtype, self.stop_gradient))
+
+    def __eq__(self, other: Self) -> bool:
+        slots = ['shape', 'dtype', 'name', 'stop_gradient']
+        return type(self) is type(other) and all(
+            getattr(self, attr) == getattr(other, attr) for attr in slots
+        )
+
+    def __ne__(self, other) -> bool:
+        return not self == other
+
+
+def setitem(
+    x: Tensor,
+    index: TensorIndex,
+    value: TensorLike,
+) -> Tensor:
+    """
+    x(Tensor): input Tensor.
+    index(Scalar|Tuple|List|Tensor): Where should be set value.
+    value(Scalar|Tensor): The value which is going to be set.
+
+    [How to write index?]
+    1. ':' -> slice(),
+       (1) a[:]=v -> setitem(a, slice(None,None,None), v)
+       (2) a[1::2] -> setitem(a, slice(1,None,2), v)
+
+    2. if there are multiple indexes for axes, use TUPLE (Not LIST) to pack them.
+       (1) a[1, 2]=v -> setitem(a, (1, 2), v)
+       (2) a[[1,2],[2,3]]=v -> setitem(a, ([1,2],[2,3]), v)
+       (3) a[1,:, 3] = v -> setitem(a, (1, slice(None,None,None),3), v)
+       (4) a[1, ..., 2]=v -> setitem(a, (1, ..., 2), v)
+
+    3. You can always use TUPLE as index input, even there is only one index.
+       (1) a[Tensor([10,10])]=v -> setitem(a, (Tensor([10,10]),), v)
+       (2) a[1] = v -> setitem(a, (1,), v)
+    """
+    return _setitem_static(x, index, value)

python/paddle/static/io_utils.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 os
+import warnings
+
+import paddle
+from paddle import pir
+from paddle.base import (
+    CompiledProgram,
+    Variable,
+)
+
+
+def _check_args(caller, args, supported_args=None, deprecated_args=None):
+    supported_args = [] if supported_args is None else supported_args
+    deprecated_args = [] if deprecated_args is None else deprecated_args
+    for arg in args:
+        if arg in deprecated_args:
+            raise ValueError(
+                f"argument '{arg}' in function '{caller}' is deprecated, only {supported_args} are supported."
+            )
+        elif arg not in supported_args:
+            raise ValueError(
+                f"function '{caller}' doesn't support argument '{arg}',\n only {supported_args} are supported."
+            )
+
+
+def _check_vars(name, var_list):
+    if not isinstance(var_list, list):
+        var_list = [var_list]
+    if not all(isinstance(var, (Variable, pir.Value)) for var in var_list):
+        raise ValueError(
+            f"'{name}' should be a Variable or a list of Variable."
+        )
+
+
+def _normalize_path_prefix(path_prefix):
+    """
+    convert path_prefix to absolute path.
+    """
+    if not isinstance(path_prefix, str):
+        raise ValueError("'path_prefix' should be a string.")
+    if path_prefix.endswith("/"):
+        raise ValueError("'path_prefix' should not be a directory")
+    path_prefix = os.path.normpath(path_prefix)
+    path_prefix = os.path.abspath(path_prefix)
+    return path_prefix
+
+
+def _get_valid_program(program=None):
+    """
+    return default main program if program is None.
+    """
+    if program is None:
+        program = paddle.static.default_main_program()
+    elif isinstance(program, CompiledProgram):
+        program = program._program
+        if program is None:
+            raise TypeError(
+                "The type of input program is invalid, expected type is Program, but received None"
+            )
+        warnings.warn(
+            "The input is a CompiledProgram, this is not recommended."
+        )
+    if not isinstance(program, paddle.static.Program):
+        raise TypeError(
+            f"The type of input program is invalid, expected type is base.Program, but received {type(program)}"
+        )
+    return program
+
+
+def _safe_load_pickle(file, encoding="ASCII"):
+    from paddle.framework.restricted_unpickler import RestrictedUnpickler
+
+    load_dict = RestrictedUnpickler(file, encoding=encoding).load()
+    return load_dict

python/paddle/static/log_helper.py

@@ -0,0 +1,57 @@
+#   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 logging
+
+
+def get_logger(name, level, fmt=None):
+    """
+    Get logger from logging with given name, level and format without
+    setting logging basicConfig. For setting basicConfig in paddle
+    will disable basicConfig setting after import paddle.
+
+    Args:
+        name (str): The logger name.
+        level (logging.LEVEL): The base level of the logger
+        fmt (str): Format of logger output
+
+    Returns:
+        logging.Logger: logging logger with given settings
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import logging
+            >>> logger = paddle.static.log_helper.get_logger(
+            ...     __name__,
+            ...     logging.INFO,
+            ...     fmt='%(asctime)s-%(levelname)s: %(message)s',
+            ... )
+    """
+
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    handler = logging.StreamHandler()
+
+    if fmt:
+        formatter = logging.Formatter(fmt=fmt, datefmt='%a %b %d %H:%M:%S')
+        handler.setFormatter(formatter)
+
+    logger.addHandler(handler)
+
+    # stop propagate for propagating may print
+    # log multiple times
+    logger.propagate = False
+    return logger

python/paddle/static/nn/__init__.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2020 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 ...tensor.creation import create_parameter  # noqa: F401
+from .common import (
+    batch_norm,
+    bilinear_tensor_product,
+    continuous_value_model,  # noqa: F401
+    conv2d,
+    conv2d_transpose,
+    conv3d,
+    conv3d_transpose,
+    deform_conv2d,
+    embedding,
+    fc,
+    group_norm,
+    instance_norm,
+    layer_norm,
+    prelu,
+    py_func,
+    row_conv,
+    sparse_embedding,
+    spectral_norm,
+)
+from .control_flow import case, cond, switch_case, while_loop
+from .loss import nce
+from .sequence_lod import (
+    sequence_conv,
+    sequence_expand,
+    sequence_first_step,
+    sequence_last_step,
+    sequence_pool,
+    sequence_softmax,
+)
+from .static_pylayer import static_pylayer
+
+__all__ = [
+    'fc',
+    'batch_norm',
+    'bilinear_tensor_product',
+    'embedding',
+    'case',
+    'cond',
+    'static_pylayer',
+    'conv2d',
+    'conv2d_transpose',
+    'conv3d',
+    'conv3d_transpose',
+    'deform_conv2d',
+    'group_norm',
+    'instance_norm',
+    'layer_norm',
+    'nce',
+    'prelu',
+    'py_func',
+    'row_conv',
+    'spectral_norm',
+    'switch_case',
+    'while_loop',
+    'sparse_embedding',
+    'sequence_conv',
+    'sequence_softmax',
+    'sequence_pool',
+    'sequence_first_step',
+    'sequence_last_step',
+    'sequence_expand',
+    'prelu',
+]

python/paddle/static/nn/loss.py

@@ -0,0 +1,266 @@
+#   Copyright (c) 2020 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 numpy as np
+
+from paddle.base.framework import static_only
+
+# TODO: define loss functions of neural network
+from paddle.base.layer_helper import LayerHelper
+from paddle.base.param_attr import ParamAttr
+from paddle.nn.initializer import Assign
+
+from ...base.data_feeder import check_variable_and_dtype
+
+__all__ = []
+
+
+# FIXME(wuyi): let docstring_checker.py understand @autodoc.
+# For now, the comments in c++ use types like Tensor, but in python side
+# the type is often "Variable", and arguments may vary.
+@static_only
+def nce(
+    input,
+    label,
+    num_total_classes,
+    sample_weight=None,
+    param_attr=None,
+    bias_attr=None,
+    num_neg_samples=None,
+    name=None,
+    sampler="uniform",
+    custom_dist=None,
+    seed=0,
+    is_sparse=False,
+):
+    """
+    :api_attr: Static Graph
+
+    Compute and return the noise-contrastive estimation training loss. See `Noise-contrastive estimation: A new estimation principle
+    for unnormalized statistical models <http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf>`_.
+    By default this operator uses a uniform distribution for sampling.
+
+    Args:
+        input (Tensor): Input tensor, 2-D tensor with shape [batch_size, dim],
+            and data type is float32 or float64.
+        label (Tensor): Input label, 2-D tensor with shape [batch_size, num_true_class],
+            and data type is int64.
+        num_total_classes (int): Total number of classes in all samples.
+        sample_weight (Tensor|None): A Tensor of shape [batch_size, 1]
+            storing a weight for each sample. The default weight for each
+            sample is 1.0.
+        param_attr (ParamAttr|None): To specify the weight parameter attribute.
+            Default: None, which means the default weight parameter property is
+            used. See usage for details in :ref:`api_paddle_ParamAttr` .
+        bias_attr (ParamAttr|None): To specify the bias parameter attribute.
+            Default: None, which means the default bias parameter property is
+            used. See usage for details in :ref:`api_paddle_ParamAttr` .
+        num_neg_samples (int): The number of negative classes. The default value is 10.
+        name(str|None): For detailed information, please refer to
+            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
+        sampler (str, optional): The sampler used to sample class from negative classes.
+                       It can be 'uniform', 'log_uniform' or 'custom_dist'.
+                       default: 'uniform'.
+        custom_dist (nd.array|None): A numpy ndarray with size=num_total_classes.
+                       It is used when sampler is set to 'custom_dist'.
+                       custom_dist[i] is the probability of i-th class to be sampled.
+                       default: None.
+        seed (int, optional): The seed used in sampler. Default 0, means no random seed.
+        is_sparse(bool, optional): The flag indicating whether to use sparse update,
+            the weight@GRAD and bias@GRAD will be changed to SelectedRows. Default False.
+
+    Returns:
+        Tensor: The output nce loss.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("paddle.static.nn.nce doesn't support PIR mode")
+            >>> import paddle
+            >>> import numpy as np
+
+            >>> paddle.enable_static()
+
+            >>> window_size = 5
+            >>> words = []
+            >>> for i in range(window_size):
+            ...     words.append(paddle.static.data(name='word_{0}'.format(i), shape=[-1, 1], dtype='int64'))
+
+            >>> dict_size = 10000
+            >>> label_word = int(window_size / 2) + 1
+
+            >>> embs = []
+            >>> for i in range(window_size):
+            ...     if i == label_word:
+            ...         continue
+            ...
+            ...     emb = paddle.static.nn.embedding(input=words[i], size=[dict_size, 32], param_attr='embed', is_sparse=True)
+            ...     embs.append(emb)
+
+            >>> embs = paddle.concat(x=embs, axis=1)  # concat from 4 * [(-1, 1, 32)] to (-1, 4, 32)
+            >>> embs = paddle.reshape(x=embs, shape=(-1, 4 * 32))  # reshape to (batch_size = -1, dim = 4*32)
+            >>> loss = paddle.static.nn.nce(
+            ...     input=embs, label=words[label_word], num_total_classes=dict_size, param_attr='nce.w_0', bias_attr='nce.b_0'
+            ... )
+
+            # or use custom distribution
+            >>> dist = np.array([0.05, 0.5, 0.1, 0.3, 0.05])
+            >>> loss = paddle.static.nn.nce(
+            ...     input=embs,
+            ...     label=words[label_word],
+            ...     num_total_classes=5,
+            ...     param_attr='nce.w_1',
+            ...     bias_attr='nce.b_1',
+            ...     num_neg_samples=3,
+            ...     sampler="custom_dist",
+            ...     custom_dist=dist,
+            ... )
+    """
+    helper = LayerHelper('nce', **locals())
+    check_variable_and_dtype(input, 'input', ['float32', 'float64'], 'nce')
+    check_variable_and_dtype(label, 'label', ['int64'], 'nce')
+
+    if input.ndim != 2:
+        raise ValueError(
+            f'The rank of `input` must be 2, but received {input.ndim}.'
+        )
+
+    dim = input.shape[1]
+    num_true_class = label.shape[1]
+    w = helper.create_parameter(
+        attr=helper.param_attr,
+        shape=[num_total_classes, dim],
+        is_bias=False,
+        dtype=input.dtype,
+    )
+    inputs = {}
+    if helper.bias_attr:
+        b = helper.create_parameter(
+            attr=helper.bias_attr,
+            shape=[num_total_classes, 1],
+            is_bias=True,
+            dtype=input.dtype,
+        )
+        inputs['Bias'] = b
+    cost = helper.create_variable_for_type_inference(dtype=input.dtype)
+    sample_logits = helper.create_variable_for_type_inference(dtype=input.dtype)
+    sample_labels = helper.create_variable_for_type_inference(dtype=label.dtype)
+
+    inputs['Input'] = input
+    inputs['Label'] = label
+    inputs['Weight'] = w
+    inputs['SampleWeight'] = sample_weight if sample_weight is not None else []
+
+    if sampler == "uniform":
+        sampler = 0
+    elif sampler == "log_uniform":
+        sampler = 1
+    elif sampler == "custom_dist":
+        assert custom_dist is not None
+
+        custom_dist_len = num_total_classes
+        alias_probs_ = [0] * custom_dist_len
+        alias_ = [0] * custom_dist_len
+        bigs = []
+        littles = []
+        for i in range(custom_dist_len):
+            normal_prob = custom_dist[i] * custom_dist_len
+            if normal_prob - 1.0 > 0:
+                bigs.append((i, normal_prob))
+            elif 1.0 - normal_prob > 0:
+                littles.append((i, normal_prob))
+            else:
+                alias_probs_[i] = normal_prob
+                alias_[i] = -1
+
+        while len(bigs) and len(littles):
+            big = bigs.pop(0)
+            little = littles.pop(0)
+
+            big_idx = big[0]
+            big_prob = big[1]
+
+            alias_probs_[little[0]] = little[1]
+            alias_[little[0]] = big_idx
+            big_left = big[1] + little[1] - 1
+            if big_left - 1.0 > 0:
+                bigs.append((big_idx, big_left))
+            elif 1.0 - big_left > 0:
+                littles.append((big_idx, big_left))
+            else:
+                alias_probs_[big_idx] = big_left
+                alias_[big_idx] = -1
+
+        if len(bigs):
+            big = bigs.pop(0)
+            alias_probs_[big[0]] = 1.0
+            alias_[big[0]] = -1
+        if len(littles):
+            little = littles.pop(0)
+            alias_probs_[little[0]] = 1.0
+            alias_[little[0]] = -1
+
+        def _init_by_numpy_array(numpy_array):
+            ret = helper.create_parameter(
+                attr=ParamAttr(),
+                shape=numpy_array.shape,
+                dtype=numpy_array.dtype,
+                default_initializer=Assign(numpy_array),
+            )
+            ret.stop_gradient = True
+            return ret
+
+        inputs['CustomDistProbs'] = _init_by_numpy_array(
+            np.array(custom_dist).astype('float32')
+        )
+        inputs['CustomDistAlias'] = _init_by_numpy_array(
+            np.array(alias_).astype('int32')
+        )
+        inputs['CustomDistAliasProbs'] = _init_by_numpy_array(
+            np.array(alias_probs_).astype('float32')
+        )
+        sampler = 2
+    else:
+        raise Exception("Unsupported sampler type.")
+
+    if num_neg_samples is None:
+        num_neg_samples = 10
+    else:
+        num_neg_samples = int(num_neg_samples)
+
+    remote_prefetch = is_sparse
+    print(
+        "With sparse mode, if your models has only small parameter prefetch may cause speed down"
+    )
+
+    attrs = {
+        'num_total_classes': int(num_total_classes),
+        'num_neg_samples': num_neg_samples,
+        'seed': seed,
+        'sampler': sampler,
+        'is_sparse': is_sparse,
+        'remote_prefetch': remote_prefetch,
+    }
+
+    helper.append_op(
+        type='nce',
+        inputs=inputs,
+        outputs={
+            'Cost': cost,
+            'SampleLogits': sample_logits,
+            'SampleLabels': sample_labels,
+        },
+        attrs=attrs,
+    )
+    return cost / (num_neg_samples + 1)

python/paddle/static/nn/metric.py

@@ -0,0 +1,627 @@
+#   Copyright (c) 2018 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.
+"""
+All layers just related to metric.
+"""
+
+import numpy as np
+
+import paddle
+from paddle import _C_ops, _legacy_C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.framework import (
+    Variable,
+    _create_tensor,
+    in_dygraph_mode,
+    in_pir_mode,
+)
+from paddle.base.layer_helper import LayerHelper
+from paddle.nn.initializer import ConstantInitializer
+
+__all__ = []
+
+
+def accuracy(input, label, k=1, correct=None, total=None):
+    """
+
+    accuracy layer.
+    Refer to the https://en.wikipedia.org/wiki/Precision_and_recall
+    This function computes the accuracy using the input and label.
+    If the correct label occurs in top k predictions, then correct will increment by one.
+
+    Note:
+        the dtype of accuracy is determined by input. the input and label dtype can be different.
+
+    Args:
+        input(Tensor): The input of accuracy layer, which is the predictions of network. A Tensor with type float32,float64.
+            The shape is ``[sample_number, class_dim]`` .
+        label(Tensor): The label of dataset.  Tensor with type int32,int64. The shape is ``[sample_number, 1]`` .
+        k(int, optional): The top k predictions for each class will be checked. Data type is int64 or int32. Default is 1.
+        correct(Tensor, optional): The correct predictions count. A Tensor with type int64 or int32. Default is None.
+        total(Tensor, optional): The total entries count. A tensor with type int64 or int32. Default is None.
+
+    Returns:
+        Tensor, The correct rate. A Tensor with type float32.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import numpy as np
+            >>> import paddle
+            >>> import paddle.static as static
+            >>> import paddle.nn.functional as F
+            >>> paddle.seed(2023)
+            >>> paddle.enable_static()
+            >>> data = static.data(name="input", shape=[-1, 32, 32], dtype="float32")
+            >>> label = static.data(name="label", shape=[-1, 1], dtype="int64")
+            >>> fc_out = static.nn.fc(x=data, size=10)
+            >>> predict = F.softmax(x=fc_out)
+            >>> result = static.accuracy(input=predict, label=label, k=5)
+            >>> place = paddle.CPUPlace()
+            >>> exe = static.Executor(place)
+            >>> exe.run(static.default_startup_program())
+            >>> np.random.seed(1107)
+            >>> x = np.random.rand(3, 32, 32).astype("float32")
+            >>> y = np.array([[1], [0], [1]])
+            >>> output = exe.run(
+            ...     feed={"input": x, "label": y},
+            ...     fetch_list=[result],
+            ... )
+            >>> print(output)
+            [array(0.33333334, dtype=float32)]
+
+    """
+    if in_dygraph_mode():
+        if correct is None:
+            correct = _create_tensor(dtype="int32")
+        if total is None:
+            total = _create_tensor(dtype="int32")
+
+        _k = np.array(k).item(0) if isinstance(k, Variable) else k
+        topk_out, topk_indices = _legacy_C_ops.top_k_v2(
+            input, 'k', _k, 'sorted', False
+        )
+        _acc, _, _ = _legacy_C_ops.accuracy(
+            topk_out, topk_indices, label, correct, total
+        )
+        return _acc
+    elif in_pir_mode():
+        topk_out, topk_indices = paddle.topk(input, k=k, sorted=False)
+        _acc, _, _ = _C_ops.accuracy(topk_out, topk_indices, label)
+        return _acc
+
+    helper = LayerHelper("accuracy", **locals())
+    check_variable_and_dtype(
+        input, 'input', ['float16', 'uint16', 'float32', 'float64'], 'accuracy'
+    )
+    topk_out = helper.create_variable_for_type_inference(dtype=input.dtype)
+    topk_indices = helper.create_variable_for_type_inference(dtype="int64")
+    inputs = {"X": [input]}
+    if isinstance(k, Variable):
+        inputs['K'] = [k]
+    else:
+        attrs = {'k': k}
+    attrs['sorted'] = False
+    helper.append_op(
+        type="top_k_v2",
+        inputs=inputs,
+        attrs=attrs,
+        outputs={"Out": [topk_out], "Indices": [topk_indices]},
+    )
+    acc_out = helper.create_variable_for_type_inference(dtype="float32")
+    if correct is None:
+        correct = helper.create_variable_for_type_inference(dtype="int32")
+    if total is None:
+        total = helper.create_variable_for_type_inference(dtype="int32")
+    helper.append_op(
+        type="accuracy",
+        inputs={"Out": [topk_out], "Indices": [topk_indices], "Label": [label]},
+        outputs={
+            "Accuracy": [acc_out],
+            "Correct": [correct],
+            "Total": [total],
+        },
+    )
+    return acc_out
+
+
+def auc(
+    input,
+    label,
+    curve='ROC',
+    num_thresholds=2**12 - 1,
+    topk=1,
+    slide_steps=1,
+    ins_tag_weight=None,
+):
+    """
+    **Area Under the Curve (AUC) Layer**
+
+    This implementation computes the AUC according to forward output and label.
+    It is used very widely in binary classification evaluation.
+
+    Note: If input label contains values other than 0 and 1, it will be cast
+    to `bool`. Find the relevant definitions `here <https://en.wikipedia.org\
+    /wiki/Receiver_operating_characteristic#Area_under_the_curve>`_.
+
+    There are two types of possible curves:
+
+        1. ROC: Receiver operating characteristic;
+        2. PR: Precision Recall
+
+    Args:
+        input(Tensor): A floating-point 2D Tensor, values are in the range
+                         [0, 1]. Each row is sorted in descending order. This
+                         input should be the output of topk. Typically, this
+                         Tensor indicates the probability of each label.
+                         A Tensor with type float32,float64.
+        label(Tensor): A 2D int Tensor indicating the label of the training
+                         data. The height is batch size and width is always 1.
+                         A Tensor with type int32,int64.
+        curve(str, optional): Curve type, can be 'ROC' or 'PR'. Default 'ROC'.
+        num_thresholds(int, optional): The number of thresholds to use when discretizing
+                             the roc curve. Default 4095.
+        topk(int, optional): only topk number of prediction output will be used for auc.
+        slide_steps(int, optional): when calc batch auc, we can not only use step currently but the previous steps can be used. slide_steps=1 means use the current step, slide_steps=3 means use current step and the previous second steps, slide_steps=0 use all of the steps.
+        ins_tag_weight(Tensor, optional): A 2D int Tensor indicating the data's tag weight, 1 means real data, 0 means fake data. Default None, and it will be assigned to a tensor of value 1.
+                         A Tensor with type float32,float64.
+
+    Returns:
+        Tensor: A tuple representing the current AUC. Data type is Tensor, supporting float32, float64.
+        The return tuple is auc_out, batch_auc_out, [batch_stat_pos, batch_stat_neg, stat_pos, stat_neg ]
+
+            auc_out: the result of the accuracy rate
+            batch_auc_out: the result of the batch accuracy
+            batch_stat_pos: the statistic value for label=1 at the time of batch calculation
+            batch_stat_neg: the statistic value for label=0 at the time of batch calculation
+            stat_pos: the statistic for label=1 at the time of calculation
+            stat_neg: the statistic for label=0 at the time of calculation
+
+
+    Examples:
+        .. code-block:: pycon
+            :name: example-1
+
+            >>> # doctest: +SKIP("This has diff in xdoctest env")
+            >>> import paddle
+            >>> import numpy as np
+            >>> paddle.enable_static()
+
+            >>> paddle.seed(2023)
+            >>> data = paddle.static.data(name="input", shape=[-1, 32,32], dtype="float32")
+            >>> label = paddle.static.data(name="label", shape=[-1], dtype="int64")
+            >>> fc_out = paddle.static.nn.fc(x=data, size=2)
+            >>> predict = paddle.nn.functional.softmax(x=fc_out)
+            >>> result=paddle.static.auc(input=predict, label=label)
+
+            >>> place = paddle.CPUPlace()
+            >>> exe = paddle.static.Executor(place)
+
+            >>> exe.run(paddle.static.default_startup_program())
+            >>> np.random.seed(1107)
+            >>> x = np.random.rand(3,32,32).astype("float32")
+            >>> y = np.array([1,0,1])
+            >>> output= exe.run(feed={"input": x,"label": y},
+            ...                 fetch_list=[result[0]])
+            >>> print(output)
+            [array(1.)]
+
+
+        .. code-block:: pycon
+            :name: example-2
+
+            # you can learn the usage of ins_tag_weight by the following code.
+
+            >>> # doctest: +SKIP("This has diff in xdoctest env")
+            >>> import paddle
+            >>> import numpy as np
+            >>> paddle.enable_static()
+
+            >>> paddle.seed(2023)
+            >>> data = paddle.static.data(name="input", shape=[-1, 32,32], dtype="float32")
+            >>> label = paddle.static.data(name="label", shape=[-1], dtype="int64")
+            >>> ins_tag_weight = paddle.static.data(name='ins_tag_weight', shape=[-1,16], dtype='float64')
+            >>> fc_out = paddle.static.nn.fc(x=data, size=2)
+            >>> predict = paddle.nn.functional.softmax(x=fc_out)
+            >>> result=paddle.static.auc(input=predict, label=label, ins_tag_weight=ins_tag_weight)
+
+            >>> place = paddle.CPUPlace()
+            >>> exe = paddle.static.Executor(place)
+
+            >>> exe.run(paddle.static.default_startup_program())
+            >>> np.random.seed(1107)
+            >>> x = np.random.rand(3,32,32).astype("float32")
+            >>> y = np.array([1,0,1])
+            >>> z = np.array([1,0,1]).astype("float64")
+            >>> output= exe.run(feed={"input": x,"label": y, "ins_tag_weight":z},
+            ...                 fetch_list=[result[0]])
+            >>> print(output)
+            [array(1.)]
+
+    """
+    if in_pir_mode():
+        if ins_tag_weight is None:
+            ins_tag_weight = paddle.full(
+                shape=[1, 1], dtype="float32", fill_value=1.0
+            )
+        check_variable_and_dtype(input, 'input', ['float32', 'float64'], 'auc')
+        check_variable_and_dtype(label, 'label', ['int32', 'int64'], 'auc')
+        check_variable_and_dtype(
+            ins_tag_weight, 'ins_tag_weight', ['float32', 'float64'], 'auc'
+        )
+        stat_pos = paddle.zeros(shape=[1, num_thresholds + 1], dtype="int64")
+        stat_neg = paddle.zeros(shape=[1, num_thresholds + 1], dtype="int64")
+        auc_out, batch_stat_pos, batch_stat_neg = _C_ops.auc(
+            input,
+            label,
+            stat_pos,
+            stat_neg,
+            ins_tag_weight,
+            curve,
+            num_thresholds,
+            0,
+        )
+        return (
+            auc_out,
+            batch_stat_pos,
+            batch_stat_neg,
+        )
+    helper = LayerHelper("auc", **locals())
+
+    if ins_tag_weight is None:
+        ins_tag_weight = paddle.tensor.fill_constant(
+            shape=[1, 1], dtype="float32", value=1.0
+        )
+    check_variable_and_dtype(input, 'input', ['float32', 'float64'], 'auc')
+    check_variable_and_dtype(label, 'label', ['int32', 'int64'], 'auc')
+    check_variable_and_dtype(
+        ins_tag_weight, 'ins_tag_weight', ['float32', 'float64'], 'auc'
+    )
+    auc_out = helper.create_variable_for_type_inference(dtype="float64")
+    batch_auc_out = helper.create_variable_for_type_inference(dtype="float64")
+    # make tp, tn, fp, fn persistable, so that can accumulate all batches.
+
+    # for batch auc
+    # we create slide_step+1 buckets, the first slide_steps buckets store
+    # historical batch-level values, and the last bucket stores the sum values of
+    # previous slide_step buckets.
+    # The index of bucket that the newest batch will use is determined by batch_id mod slide_steps,
+    # and batch_id is store in the last position of following variable
+    batch_stat_pos = helper.create_global_variable(
+        persistable=True,
+        dtype='int64',
+        shape=[(1 + slide_steps) * (num_thresholds + 1) + 1],
+    )
+    batch_stat_neg = helper.create_global_variable(
+        persistable=True,
+        dtype='int64',
+        shape=[(1 + slide_steps) * (num_thresholds + 1) + 1],
+    )
+
+    # for global auc
+    # Needn't maintain the batch id
+    stat_pos = helper.create_global_variable(
+        persistable=True, dtype='int64', shape=[1, num_thresholds + 1]
+    )
+    stat_neg = helper.create_global_variable(
+        persistable=True, dtype='int64', shape=[1, num_thresholds + 1]
+    )
+
+    for var in [batch_stat_pos, batch_stat_neg, stat_pos, stat_neg]:
+        helper.set_variable_initializer(
+            var,
+            ConstantInitializer(value=0.0, force_cpu=False),
+        )
+
+    # "InsTagWeight": [ins_tag_weight]
+    # Batch AUC
+    helper.append_op(
+        type="auc",
+        inputs={
+            "Predict": [input],
+            "Label": [label],
+            "StatPos": [batch_stat_pos],
+            "StatNeg": [batch_stat_neg],
+        },
+        attrs={
+            "curve": curve,
+            "num_thresholds": num_thresholds,
+            "slide_steps": slide_steps,
+        },
+        outputs={
+            "AUC": [batch_auc_out],
+            "StatPosOut": [batch_stat_pos],
+            "StatNegOut": [batch_stat_neg],
+        },
+    )
+    # Global AUC
+    helper.append_op(
+        type="auc",
+        inputs={
+            "Predict": [input],
+            "Label": [label],
+            "StatPos": [stat_pos],
+            "StatNeg": [stat_neg],
+        },
+        attrs={
+            "curve": curve,
+            "num_thresholds": num_thresholds,
+            "slide_steps": 0,
+        },
+        outputs={
+            "AUC": [auc_out],
+            "StatPosOut": [stat_pos],
+            "StatNegOut": [stat_neg],
+        },
+    )
+    return (
+        auc_out,
+        batch_auc_out,
+        [batch_stat_pos, batch_stat_neg, stat_pos, stat_neg],
+    )
+
+
+def ctr_metric_bundle(input, label, ins_tag_weight=None):
+    """
+    ctr related metric layer
+
+    This function help compute the ctr related metrics: RMSE, MAE, predicted_ctr, q_value.
+    To compute the final values of these metrics, we should do following computations using
+    total instance number:
+    MAE = local_abserr / instance number
+    RMSE = sqrt(local_sqrerr / instance number)
+    predicted_ctr = local_prob / instance number
+    q = local_q / instance number
+    Note that if you are doing distribute job, you should all reduce these metrics and instance
+    number first
+
+    Args:
+        input(Tensor): A floating-point 2D Tensor, values are in the range
+                         [0, 1]. Each row is sorted in descending order. This
+                         input should be the output of topk. Typically, this
+                         Tensor indicates the probability of each label.
+        label(Tensor): A 2D int Tensor indicating the label of the training
+                         data. The height is batch size and width is always 1.
+        ins_tag_weight(Tensor): A 2D int Tensor indicating the ins_tag_weight of the training
+                         data. 1 means real data, 0 means fake data.
+                         A DenseTensor or Tensor with type float32,float64.
+
+    Returns:
+        local_sqrerr(Tensor): Local sum of squared error
+        local_abserr(Tensor): Local sum of abs error
+        local_prob(Tensor): Local sum of predicted ctr
+        local_q(Tensor): Local sum of q value
+        local_pos_num (Tensor): Local number of positive examples
+        local_ins_num (Tensor): Local number of instances
+
+    Examples:
+        .. code-block:: pycon
+            :name: example-1
+
+            >>> # doctest: +SKIP("This has diff in xdoctest env")
+            >>> import paddle
+            >>> paddle.enable_static()
+            >>> data = paddle.static.data(name="data", shape=[-1, 32], dtype="float32")
+            >>> label = paddle.static.data(name="label", shape=[-1, 1], dtype="int32")
+            >>> predict = paddle.nn.functional.sigmoid(paddle.static.nn.fc(x=data, size=1))
+            >>> auc_out = paddle.static.ctr_metric_bundle(input=predict, label=label)
+
+        .. code-block:: pycon
+            :name: example-2
+
+            >>> # doctest: +SKIP("This has diff in xdoctest env")
+            >>> import paddle
+            >>> paddle.enable_static()
+            >>> data = paddle.static.data(name="data", shape=[-1, 32], dtype="float32")
+            >>> label = paddle.static.data(name="label", shape=[-1, 1], dtype="int32")
+            >>> predict = paddle.nn.functional.sigmoid(paddle.static.nn.fc(x=data, size=1))
+            >>> ins_tag_weight = paddle.static.data(name='ins_tag_weight', shape=[-1, 1], dtype='int64')
+            >>> auc_out = paddle.static.ctr_metric_bundle(input=predict, label=label, ins_tag_weight=ins_tag_weight)
+    """
+    if ins_tag_weight is None:
+        ins_tag_weight = paddle.tensor.fill_constant(
+            shape=[1, 1], dtype="float32", value=1.0
+        )
+
+    assert input.shape == label.shape
+    helper = LayerHelper("ctr_metric_bundle", **locals())
+
+    local_abserr = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+    local_sqrerr = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+    local_prob = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+    local_q = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+    local_pos_num = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+    local_ins_num = helper.create_global_variable(
+        persistable=True, dtype='float32', shape=[1]
+    )
+
+    tmp_res_elesub = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[-1]
+    )
+    tmp_res_sigmoid = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[-1]
+    )
+    tmp_ones = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[-1]
+    )
+
+    batch_prob = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    batch_abserr = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    batch_sqrerr = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    batch_q = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    batch_pos_num = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    batch_ins_num = helper.create_global_variable(
+        persistable=False, dtype='float32', shape=[1]
+    )
+    for var in [
+        local_abserr,
+        batch_abserr,
+        local_sqrerr,
+        batch_sqrerr,
+        local_prob,
+        batch_prob,
+        local_q,
+        batch_q,
+        batch_pos_num,
+        batch_ins_num,
+        local_pos_num,
+        local_ins_num,
+    ]:
+        helper.set_variable_initializer(
+            var,
+            paddle.nn.initializer.ConstantInitializer(
+                value=0.0, force_cpu=True
+            ),
+        )
+
+    helper.append_op(
+        type="elementwise_sub",
+        inputs={"X": [input], "Y": [label]},
+        outputs={"Out": [tmp_res_elesub]},
+    )
+
+    helper.append_op(
+        type="squared_l2_norm",
+        inputs={"X": [tmp_res_elesub]},
+        outputs={"Out": [batch_sqrerr]},
+    )
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_sqrerr], "Y": [local_sqrerr]},
+        outputs={"Out": [local_sqrerr]},
+    )
+
+    helper.append_op(
+        type="l1_norm",
+        inputs={"X": [tmp_res_elesub]},
+        outputs={"Out": [batch_abserr]},
+    )
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_abserr], "Y": [local_abserr]},
+        outputs={"Out": [local_abserr]},
+    )
+
+    helper.append_op(
+        type="reduce_sum", inputs={"X": [input]}, outputs={"Out": [batch_prob]}
+    )
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_prob], "Y": [local_prob]},
+        outputs={"Out": [local_prob]},
+    )
+    helper.append_op(
+        type="sigmoid",
+        inputs={"X": [input]},
+        outputs={"Out": [tmp_res_sigmoid]},
+    )
+    helper.append_op(
+        type="reduce_sum",
+        inputs={"X": [tmp_res_sigmoid]},
+        outputs={"Out": [batch_q]},
+    )
+
+    helper.append_op(
+        type="reduce_sum",
+        inputs={"X": [label]},
+        outputs={"Out": [batch_pos_num]},
+    )
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_pos_num], "Y": [local_pos_num]},
+        outputs={"Out": [local_pos_num]},
+    )
+
+    helper.append_op(
+        type='fill_constant_batch_size_like',
+        inputs={"Input": label},
+        outputs={'Out': [tmp_ones]},
+        attrs={
+            'shape': [-1, 1],
+            'dtype': tmp_ones.dtype,
+            'value': 1.0,
+        },
+    )
+    helper.append_op(
+        type="reduce_sum",
+        inputs={"X": [tmp_ones]},
+        outputs={"Out": [batch_ins_num]},
+    )
+
+    # if data is fake, return 0
+    inputs_slice = {'Input': ins_tag_weight}
+    attrs = {'axes': [0]}
+    attrs['starts'] = [0]
+    attrs['ends'] = [1]
+    helper.append_op(
+        type="slice",
+        inputs=inputs_slice,
+        attrs=attrs,
+        outputs={"Out": ins_tag_weight},
+    )
+
+    axis = helper.kwargs.get('axis', 0)
+    helper.append_op(
+        type="elementwise_mul",
+        inputs={"X": [batch_ins_num], "Y": [ins_tag_weight]},
+        outputs={"Out": [batch_ins_num]},
+        attrs={'axis': axis},
+    )
+
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_ins_num], "Y": [local_ins_num]},
+        outputs={"Out": [local_ins_num]},
+    )
+
+    helper.append_op(
+        type="elementwise_mul",
+        inputs={"X": [batch_q], "Y": [ins_tag_weight]},
+        outputs={"Out": [batch_q]},
+        attrs={'axis': axis},
+    )
+    helper.append_op(
+        type="elementwise_add",
+        inputs={"X": [batch_q], "Y": [local_q]},
+        outputs={"Out": [local_q]},
+    )
+
+    return (
+        local_sqrerr,
+        local_abserr,
+        local_prob,
+        local_q,
+        local_pos_num,
+        local_ins_num,
+    )

python/paddle/static/nn/sequence_lod.py

@@ -0,0 +1,755 @@
+# 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 paddle
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.framework import in_dygraph_mode, in_pir_mode
+from paddle.base.layer_helper import LayerHelper
+from paddle.utils import deprecated
+
+__all__ = []
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_conv(
+    input,
+    num_filters,
+    filter_size=3,
+    filter_stride=1,
+    padding=True,
+    padding_start=None,
+    bias_attr=None,
+    param_attr=None,
+    act=None,
+    name=None,
+):
+    r"""
+
+    Note:
+        Only receives Tensor as input. If your input is Tensor, please use conv2d Op.(base.layers.** :ref:`api_paddle_nn_functional_conv2d` ).
+
+    This operator receives input sequences with variable length and other convolutional
+    configuration parameters(num_filters, filter_size) to apply the convolution operation.
+    It fills all-zero padding data on both sides of the sequence by default to ensure that
+    the output is the same length as the input. You can customize the padding behavior by
+    configuring the parameter :attr:`padding\_start` .
+
+    **Warning:** the parameter :attr:`padding` take no effect and will be deprecated in the future.
+
+    .. code-block:: text
+
+            Here we will illustrate the details of the padding operation:
+            For a mini-batch of 2 variable lengths sentences, containing 3, and 1 time-steps:
+            Assumed input (X) is a [4, N] float Tensor, and for the sake of simplicity, we assume N=2.
+            input.data = [[1, 1],
+                          [2, 2],
+                          [3, 3],
+                          [4, 4]]
+
+            This is to say that input (X) has 4 words and the dimension of each word
+            representation is 2.
+
+            * Case1:
+
+                If padding_start is -1 and filter_size is 3.
+                The length of padding data is calculated as follows:
+                up_pad_len = max(0, -padding_start) = 1
+                down_pad_len = max(0, filter_size + padding_start - 1) = 1
+
+                The output of the input sequence after padding is:
+                data_after_padding = [[0, 0, 1, 1, 2, 2],
+                                      [1, 1, 2, 2, 3, 3],
+                                      [2, 2, 3, 3, 0, 0],
+                                      [0, 0, 4, 4, 0, 0]]
+
+                It will be multiplied by the filter weight to get the final output.
+                Assume num_filters = 3
+                output.data = [[ 0.3234, -0.2334,  0.7433],
+                               [ 0.5646,  0.9464, -0.1223],
+                               [-0.1343,  0.5653,  0.4555],
+                               [ 0.9954, -0.1234, -0.1234]]
+                output.shape = [4, 3]     # 3 = num_filters
+                output.lod = [[0, 3, 4]]  # Remain the same
+
+
+    Args:
+        input (Tensor): Tensor with shape :math:`(M, K)`, where M is the total time-step of mini-batch
+            and K is hidden_size of input. Only lod_level of 1 is supported. The data type should be float32 or
+            float64.
+        num_filters (int): the number of filters.
+        filter_size (int): the height of filter. Specified filter width is not supported, the width is
+            hidden_size by default. Default: 3.
+        filter_stride (int, optional): stride of the filter. Currently only supports :attr:`stride` = 1.
+        padding (bool, optional): the parameter :attr:`padding` take no effect and will be discarded in the
+            future. Currently, it will always pad input to make sure the length of the output is
+            the same as input whether :attr:`padding` is set true or false. Because the length of
+            input sequence may be shorter than :attr:`filter\_size`, which will cause the convolution
+            result to not be computed correctly. These padding data will not be trainable or updated
+            while training. Default: True.
+        padding_start (int): It is used to indicate the start index for padding the input
+            sequence, which can be negative. The negative number means to pad
+            :attr:`|padding_start|` time-steps of all-zero data at the beginning of each instance.
+            The positive number means to skip :attr:`padding_start` time-steps of each instance,
+            and it will pad :math:`filter\_size + padding\_start - 1` time-steps of all-zero data
+            at the end of the sequence to ensure that the output is the same length as the input.
+            If set None, the same length :math:`\\frac{filter\_size}{2}` of data will be filled
+            on both sides of the sequence. If set 0, the length of :math:`filter\_size - 1` data
+            is padded at the end of each input sequence. Default: None.
+        bias_attr (ParamAttr): To specify the bias parameter property. Default: None, which means the
+            default bias parameter property is used. See usage for details in :ref:`api_paddle_ParamAttr` .
+        param_attr (ParamAttr): To specify the weight parameter property. Default: None, which means the
+            default weight parameter property is used. See usage for details in :ref:`api_paddle_ParamAttr` .
+        act (str): Activation to be applied to the output of this layer, such as tanh, softmax,
+            sigmoid, relu. For more information, please refer to :ref:`api_guide_activations_en` . Default: None.
+        name (str, optional): The default value is None.  Normally there is no need for user to set this property.
+            For more information, please refer to :ref:`api_guide_Name` .
+
+    Returns:
+        Tensor: Tensor with the same length as input. The data type is float32 or float64, which is same as input.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> x = paddle.static.data(name='x', shape=[-1, 10], dtype='float32', lod_level=1)
+            >>> x_conved = paddle.static.nn.sequence_conv(input=x, num_filters=2, filter_size=3, padding_start=-1)
+    """
+
+    assert not in_dygraph_mode(), (
+        "sequence layer is not supported in dygraph mode yet."
+    )
+    assert not in_pir_mode(), (
+        "sequence layer is not supported in pir mode, please set the environment variable FLAGS_enable_pir_api=0 to switch old static mode."
+    )
+    check_variable_and_dtype(
+        input, 'input', ['float32', 'float64'], 'sequence_conv'
+    )
+    helper = LayerHelper('sequence_conv', **locals())
+    dtype = helper.input_dtype()
+    filter_shape = [filter_size * input.shape[1], num_filters]
+    filter_param = helper.create_parameter(
+        attr=helper.param_attr, shape=filter_shape, dtype=dtype
+    )
+    pre_bias = helper.create_variable_for_type_inference(dtype)
+    if padding_start is None:
+        padding_start = -int(filter_size // 2)
+
+    helper.append_op(
+        type='sequence_conv',
+        inputs={
+            'X': [input],
+            'Filter': [filter_param],
+        },
+        outputs={"Out": pre_bias},
+        attrs={
+            'contextStride': filter_stride,
+            'contextStart': padding_start,
+            'contextLength': filter_size,
+        },
+    )
+    pre_act = helper.append_bias_op(pre_bias)
+    return helper.append_activation(pre_act)
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_softmax(input, use_cudnn=False, name=None):
+    r"""
+
+    Note:
+        The input type of the OP must be Tensor. For Tensor, use:** :ref:`api_paddle_nn_functional_softmax`
+
+    A LoD-tensor can be regarded as several sequences, and this op apply softmax algo on each sequence.
+    The shape of input Tensor can be :math:`[N, 1]` or :math:`[N]`, where :math:`N`
+    is the sum of the length of all sequences. Recommended usage: :math:`[N]`.
+
+    For i-th sequence in a mini-batch:
+
+    .. math::
+
+        Out(X[lod[i]:lod[i+1]], :) = \\frac{\exp(X[lod[i]:lod[i+1], :])}{\sum(\exp(X[lod[i]:lod[i+1], :]))}
+
+    For example, for a LoD-Tensor with 6 sequences ([3, 2, 4, 1, 2, 3] - sequence length list in order),
+    the lod in the runtime is [[0, 3, 5, 9, 10, 12, 15]],
+    then softmax will be computed among :math:`X[0:3,:],X[3:5,:],X[5:9,:],X[9:10,:],X[10:12,:],X[12:15,:]`,
+    and :math:`N` turns out to be 15.
+
+    .. code-block:: text
+
+        *Case 1:
+
+            Given:
+                input.data = [0.7, 1, 0.6,
+                              1.5, 1.1,
+                              1.2, 0.2, 0.6, 1.9,
+                              3.1,
+                              2.5, 0.8,
+                              0.1, 2.4, 1.3]
+                input.lod = [[0, 3, 5, 9, 10, 12, 15]]
+            then:
+                 output.data = [0.30724832, 0.41474187, 0.2780098,
+                                0.59868765, 0.40131235,
+                                0.2544242, 0.09359743, 0.13963096, 0.5123474,
+                                1.,
+                                0.84553474, 0.15446526,
+                                0.06995796, 0.69777346, 0.23226859]
+                 output.lod = [[0, 3, 5, 9, 10, 12, 15]]
+
+
+    Args:
+        input (Tensor):A Tensor with shape of  :math:`[N, 1]` or  :math:`[N]`, Recommended usage: :math:`[N]`.
+                         Supported data types: float32, float64.
+        use_cudnn (bool, optional): Use cudnn kernel or not. Effective only when the cudnn version of the paddle
+                                    library is installed and GPU is used for training or reasoning. Default: False.
+        name (str, optional): The default value is None. Normally there is no need for user to set this property.
+                              For more information, please refer to :ref:`api_guide_Name`
+
+    Returns:
+        Tensor: A LoD-Tensor which has the same shape and data type with input.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> x = paddle.static.data(name='x', shape=[7, 1], dtype='float32', lod_level=1)
+            >>> x_sequence_softmax_1 = paddle.static.nn.sequence_softmax(input=x)
+
+            >>> y = paddle.static.data(name='y', shape=[7], dtype='float32', lod_level=1)
+            >>> x_sequence_softmax_2 = paddle.static.nn.sequence_softmax(input=y)
+    """
+    assert not in_dygraph_mode(), (
+        "sequence layer is not supported in dygraph mode yet."
+    )
+    assert not in_pir_mode(), (
+        "sequence layer is not supported in pir mode, please set the environment variable FLAGS_enable_pir_api=0 to switch old static mode."
+    )
+    helper = LayerHelper('sequence_softmax', **locals())
+    check_variable_and_dtype(
+        input, 'input', ['float32', 'float64'], 'sequence_softmax'
+    )
+    dtype = helper.input_dtype()
+    softmax_out = helper.create_variable_for_type_inference(dtype)
+    helper.append_op(
+        type="sequence_softmax",
+        inputs={"X": input},
+        outputs={"Out": softmax_out},
+        attrs={"use_cudnn": use_cudnn},
+    )
+    return softmax_out
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_pool(input, pool_type, is_test=False, pad_value=0.0):
+    r"""
+
+    Note:
+        Only receives Tensor as input. If your input is Tensor, please use pool2d Op.(static.nn.** :ref:`api_paddle_nn_functional_avg_pool2d` or :ref:`api_paddle_nn_functional_max_pool2d` ).
+
+    This operator only supports Tensor as input. It will apply specified pooling
+    operation on the input Tensor. It pools features of all time-steps of each
+    sequence at the last lod_level using :attr:`pool_type` mentioned in the parameters,
+    such as sum, average, sqrt, etc.
+
+    It supports six pool_type:
+
+    - average: :math:`Out[i] = \\frac{\sum_i X_i}{N}`
+    - sum:     :math:`Out[i] = \sum_jX_{ij}`
+    - sqrt:    :math:`Out[i] = \\frac{\sum_jX_{ij}}{\sqrt{len(X_i)}}`
+    - max:     :math:`Out[i] = max(X_i)`
+    - last:    :math:`Out[i] = X_{N_i}`
+    - first:   :math:`Out[i]` = X_0
+
+    where :math:`N_i` is the length of i-th input sequence.
+
+    .. code-block:: text
+
+        Case 1:
+        input is a 1-level Tensor and pad_value = 0.0:
+            input.lod = [[0, 2, 5, 7, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        output is Tensor:
+            out.shape = [4, 1]
+            with condition out.shape[0] == len(x.lod[-1]) == 4
+
+        for different pool_type:
+            average: out.data = [[2.], [4.], [3.], [0.0]], where 2.=(1. + 3.)/2, 4.=(2. + 4. + 6.)/3, 3.=(5. + 1.)/2
+            sum    : out.data = [[4.], [12.], [6.], [0.0]], where 4.=1. + 3., 12.=2. + 4. + 6., 6.=5. + 1.
+            sqrt   : out.data = [[2.82], [6.93], [4.24], [0.0]], where 2.82=(1. + 3.)/sqrt(2), 6.93=(2. + 4. + 6.)/sqrt(3), 4.24=(5. + 1.)/sqrt(2)
+            max    : out.data = [[3.], [6.], [5.], [0.0]], where 3.=max(1., 3.), 6.=max(2., 4., 6.), 5.=max(5., 1.)
+            last   : out.data = [[3.], [6.], [1.], [0.0]], where 3.=last(1., 3.), 6.=last(2., 4., 6.), 1.=last(5., 1.)
+            first  : out.data = [[1.], [2.], [5.], [0.0]], where 1.=first(1., 3.), 2.=first(2., 4., 6.), 5.=first(5., 1.)
+
+            and all above [0.0] at last of out.data is padding data.
+
+        Case 2:
+        input is a 2-level Tensor containing 3 sequences with length info [2, 0, 3],
+        where 0 means empty sequence.
+        The first sequence contains 2 subsequence with length info [1, 2];
+        The last sequence contains 3 subsequence with length info [1, 0, 3].
+            input.lod = [[0, 2, 2, 5], [0, 1, 3, 4, 4, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        If pool_typ = sum, it will apply pooling on last lod_level [0, 1, 3, 4, 4, 7]. pad_value = 0.0
+        output is Tensor:
+            out.shape= [5, 1]
+            out.lod = [[0, 2, 2, 5]]
+            where out.shape[0] == len(x.lod[-1]) == 5
+            sum: out.data = [[1.], [5.], [4.], [0.0], [12.]]
+            where 1.=1., 5.=3. + 2., 4.=4., 0.0=pad_value, 12.=6. + 5. + 1.
+
+    Args:
+        input (variable): Tensor with lod_level no more than 2. The data type should be float32 or float64.
+        pool_type (str): The pooling type that supports average, sum, sqrt, max, last or first.
+        is_test (bool): Only works when :attr:`pool_type` is max. If set False, a temporary Tensor maxIndex is
+            created to record the index information corresponding to the maximum value, which is used for backward
+            gradient calculation in the training phase. Default: False.
+        pad_value (float): Used to pad the pooling result for empty input sequence. Default: 0.0
+
+    Returns:
+        Tensor: Tensor after pooling with data type float32 or float64.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> x = paddle.static.data(name='x', shape=[None, 10], dtype='float32', lod_level=1)
+            >>> avg_x = paddle.static.nn.sequence_pool(input=x, pool_type='average')
+            >>> sum_x = paddle.static.nn.sequence_pool(input=x, pool_type='sum')
+            >>> sqrt_x = paddle.static.nn.sequence_pool(input=x, pool_type='sqrt')
+            >>> max_x = paddle.static.nn.sequence_pool(input=x, pool_type='max')
+            >>> last_x = paddle.static.nn.sequence_pool(input=x, pool_type='last')
+            >>> first_x = paddle.static.nn.sequence_pool(input=x, pool_type='first')
+    """
+    assert not in_dygraph_mode(), (
+        "sequence layer is not supported in dygraph mode yet."
+    )
+    assert not in_pir_mode(), (
+        "sequence layer is not supported in pir mode, please set the environment variable FLAGS_enable_pir_api=0 to switch old static mode."
+    )
+
+    check_variable_and_dtype(
+        input, 'input', ['float32', 'float64'], 'sequence_pool'
+    )
+    helper = LayerHelper('sequence_pool', **locals())
+    dtype = helper.input_dtype()
+    pool_out = helper.create_variable_for_type_inference(dtype)
+    max_index = helper.create_variable_for_type_inference(dtype)
+
+    helper.append_op(
+        type="sequence_pool",
+        inputs={"X": input},
+        outputs={"Out": pool_out, "MaxIndex": max_index},
+        attrs={
+            "pooltype": pool_type.upper(),
+            "is_test": is_test,
+            "pad_value": pad_value,
+        },
+    )
+
+    # when pool_type is max, variable max_index is initialized,
+    # so we stop the gradient explicitly here
+    if pool_type == 'max':
+        max_index.stop_gradient = True
+
+    return pool_out
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_first_step(input):
+    """
+
+    Only supports Tensor as input. Given the input Tensor, it will
+    select first time-step feature of each sequence as output.
+
+    .. code-block:: text
+
+       Case 1:
+        input is 1-level Tensor:
+            input.lod = [[0, 2, 5, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        output is a Tensor:
+            out.shape = [3, 1]
+            out.shape[0] == len(x.lod[-1]) == 3
+            out.data = [[1.], [2.], [5.]], where 1.=first(1., 3.), 2.=first(2., 4., 6.), 5.=first(5., 1.)
+
+        Case 2:
+        input is a 2-level Tensor containing 3 sequences with length info [2, 0, 3],
+        where 0 means empty sequence.
+        The first sequence contains 2 subsequence with length info [1, 2];
+        The last sequence contains 3 subsequence with length info [1, 0, 3].
+            input.lod = [[0, 2, 2, 5], [0, 1, 3, 4, 4, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        It will apply pooling on last lod_level [0, 1, 3, 4, 4, 7]. pad_value = 0.0
+        output is a Tensor:
+            out.shape= [5, 1]
+            out.lod = [[0, 2, 2, 5]]
+            out.shape[0] == len(x.lod[-1]) == 5
+            out.data = [[1.], [3.], [4.], [0.0], [6.]]
+            where 1.=first(1.), 3.=first(3., 2.), 4.=first(4.), 0.0 = pad_value, 6.=first(6., 5., 1.)
+
+    Args:
+        input(Tensor): Tensor with lod_level no more than 2. The data type should be float32 or float64.
+
+    Returns:
+        Tensor: Tensor consist of the sequence's first step vector. The data type is float32 or float64.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> x = paddle.static.data(name='x', shape=[None, 10], dtype='float32', lod_level=1)
+            >>> x_first_step = paddle.static.nn.sequence_first_step(input=x)
+    """
+    check_variable_and_dtype(
+        input, 'input', ['float32', 'float64'], 'sequence_first_step'
+    )
+    return sequence_pool(input=input, pool_type="first")
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_last_step(input):
+    """
+
+    Only supports Tensor as input. Given the input Tensor, it will
+    select last time-step feature of each sequence as output.
+
+    .. code-block:: text
+
+        Case 1:
+        input is 1-level Tensor:
+            input.lod = [[0, 2, 5, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        output is a Tensor:
+            out.shape = [3, 1]
+            out.shape[0] == len(x.lod[-1]) == 3
+            out.data = [[3.], [6.], [1.]], where 3.=last(1., 3.), 6.=last(2., 4., 6.), 1.=last(5., 1.)
+
+        Case 2:
+        input is a 2-level Tensor containing 3 sequences with length info [2, 0, 3],
+        where 0 means empty sequence.
+        The first sequence contains 2 subsequence with length info [1, 2];
+        The last sequence contains 3 subsequence with length info [1, 0, 3].
+            input.lod = [[0, 2, 2, 5], [0, 1, 3, 4, 4, 7]]
+            input.data = [[1.], [3.], [2.], [4.], [6.], [5.], [1.]]
+            input.shape = [7, 1]
+
+        It will apply pooling on last lod_level [0, 1, 3, 4, 4, 7]. pad_value = 0.0
+        output is a Tensor:
+            out.shape= [5, 1]
+            out.lod = [[0, 2, 2, 5]]
+            out.shape[0] == len(x.lod[-1]) == 5
+            out.data = [[1.], [2.], [4.], [0.0], [1.]]
+            where 1.=last(1.), 2.=last(3., 2.), 4.=last(4.), 0.0 = pad_value, 1=last(6., 5., 1.)
+
+
+    Args:
+        input(Tensor): Tensor with lod_level no more than 2. The data type should be float32.
+
+    Returns:
+        Tensor: Tensor consist of the sequence's last step vector. The data type is float32.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> paddle.enable_static()
+
+            >>> x = paddle.static.data(name='x', shape=[None, 10], dtype='float32', lod_level=1)
+            >>> x_last_step = paddle.static.nn.sequence_last_step(input=x)
+    """
+    check_variable_and_dtype(
+        input, 'input', ['float32', 'float64'], 'sequence_last_step'
+    )
+    return sequence_pool(input=input, pool_type="last")
+
+
+@deprecated(
+    since="3.0.0",
+    level=1,
+    reason="This API will be deprecated in the future, because it's just for old statics mode.",
+)
+def sequence_expand(x, y, ref_level=-1, name=None):
+    r"""
+
+        Sequence Expand Layer. This layer will expand the input variable ``x`` \
+        according to specified level ``ref_level`` lod of ``y``. Please note that \
+        the lod level of ``x`` is at most 1. If the lod level of ``x`` is 1, than \
+        the size of lod of ``x`` must be equal to the length of ``ref_level`` lod \
+        of ``y``. If the lod level of ``x`` is 0, then the first dim of ``x`` should \
+        be equal to the size of ``ref_level`` of ``y``. The rank of **x** is at least 2. \
+        When rank of ``x`` is greater than 2, then it would be viewed as a 2-D tensor.
+
+    Note:
+
+        Please note that the input ``x`` should be Tensor or Tensor, \
+        and input ``y`` must be Tensor.
+
+    **Following examples will explain how sequence_expand works:**
+
+    .. code-block:: text
+
+        Case 1
+
+        Consider 2 sequences [a][b] and [c][d], now we want to expand them to [a][b], [a][b], [c][d] and [c][d].
+        Sequence [a][b] expand twice and [c][d] expands twice, so the lod which according to is [2, 2].
+
+        Input x is a 1-level Tensor:
+            x.lod  = [[2,        2]]    #lod based on length may be easier to understand
+            x.data = [[a], [b], [c], [d]]
+            x.dims = [4, 1]
+
+        input y is a Tensor:
+            y.lod = [[2,    2],    #the 0th level lod, according to this level
+                     [3, 3, 1, 1]] #the 1st level lod, it has nothing to do with this level
+
+        ref_level: 0
+
+        then output is a 1-level Tensor out:
+            out.lod =  [[2,        2,        2,        2]]    #lod based on offset
+            out.data = [[a], [b], [a], [b], [c], [d], [c], [d]]
+            out.dims = [8, 1]
+
+
+        Case 2
+
+        Consider 3 sequences [a], [b], [c], now we want to expand them to [a][a], [c][c][c].
+        It's obvious that the lod info of expanded sequences is [2, 0, 3].
+
+        x is a Tensor:
+            x.data = [[a], [b], [c]]
+            x.dims = [3, 1]
+
+        y is a Tensor:
+            y.lod = [[2, 0, 3]]
+
+        ref_level: -1
+
+        then output is a 1-level Tensor:
+            out.data = [[a], [a], [c], [c], [c]]
+            out.dims = [5, 1]
+
+    Args:
+        x (Tensor): The input variable which is a Tensor or Tensor, with the \
+            dims ``[M, K]``. The lod level is at most 1. The data type should be \
+            float32, float64, int32 or int64.
+        y (Tensor): The input variable which is a Tensor, the lod level is \
+            at least 1.
+        ref_level (int): Lod level of ``y`` to be referred by ``x``. If set to -1, \
+                         refer the last level of lod.
+        name(str, optional): For detailed information, please refer \
+            to :ref:`api_guide_Name`. Usually name is no need to set and \
+            None by default.
+
+    Returns:
+            Tensor, The expanded variable which is a Tensor, with dims ``[N, K]``. \
+            ``N`` depends on the lod info of ``x`` and ``y``. \
+            The data type is same as input.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP("env set will not work in ci check because import paddle in global_exec")
+            >>> # set env var before import paddle to disable pir mode, following example code use os module.
+            >>> import os
+            >>> os.environ['FLAGS_enable_pir_api'] = '0'
+            >>> import paddle
+            >>> from paddle import base
+            >>> paddle.enable_static()
+            >>> import numpy as np
+
+            >>> x = paddle.static.data(name='x', shape=[4, 1], dtype='float32')
+            >>> y = paddle.static.data(name='y', shape=[8, 1],
+            ...             dtype='float32', lod_level=1)
+            >>> out = paddle.static.nn.sequence_expand(x=x, y=y, ref_level=0)
+
+            >>> exe = paddle.static.Executor(base.CPUPlace())
+            >>> place = paddle.CPUPlace()
+
+            >>> np_data = np.array([[1], [2], [3], [4]]).astype('float32')
+            >>> x_lod_tensor = base.create_lod_tensor(np_data, [[2, 2]], place)
+            >>> print(x_lod_tensor)
+            - lod: {{0, 2, 4}}
+            - place: Place(cpu)
+            - shape: [4, 1]
+            - layout: NCHW
+            - dtype: float32
+            - data: [1 2 3 4]
+
+            >>> np_data = np.array([[1], [2], [3], [4], [5], [6], [7], [8]]).astype('float32')
+            >>> y_lod_tensor = base.create_lod_tensor(np_data, [[2, 2], [3,3,1,1]], place)
+            >>> print(y_lod_tensor)
+            - lod: {{0, 2, 4}{0, 3, 6, 7, 8}}
+            - place: Place(cpu)
+            - shape: [8, 1]
+            - layout: NCHW
+            - dtype: float32
+            - data: [1 2 3 4 5 6 7 8]
+
+            >>> out_main = exe.run(base.default_main_program(),
+            ...                 feed={'x': x_lod_tensor, 'y': y_lod_tensor},
+            ...                 fetch_list=[out], return_numpy=False)
+            >>> print(out_main[0])
+            - lod: {{0, 2, 4, 6, 8}}
+            - place: Place(cpu)
+            - shape: [8, 1]
+            - layout: NCHW
+            - dtype: float32
+            - data: [1 2 1 2 3 4 3 4]
+    """
+    assert not in_dygraph_mode(), (
+        "sequence layer is not supported in dygraph mode yet."
+    )
+    assert not in_pir_mode(), (
+        "sequence layer is not supported in pir mode, please set the environment variable FLAGS_enable_pir_api=0 to switch old static mode."
+    )
+    check_variable_and_dtype(
+        x, 'x', ['float32', 'float64', 'int32', 'int64'], 'sequence_expand'
+    )
+    helper = LayerHelper('sequence_expand', **locals())
+    dtype = helper.input_dtype(input_param_name='x')
+    tmp = helper.create_variable_for_type_inference(dtype)
+    helper.append_op(
+        type='sequence_expand',
+        inputs={'X': x, 'Y': y},
+        outputs={'Out': tmp},
+        attrs={'ref_level': ref_level},
+    )
+    return tmp
+
+
+@deprecated(
+    update_to="paddle.nn.functional.sequence_mask",
+    level=1,
+)
+def sequence_mask(x, maxlen=None, dtype='int64', name=None):
+    r"""
+    **SequenceMask Layer**
+
+    This layer outputs a mask according to the input :code:`x` and
+    :code:`maxlen` with data type of :code:`dtype`.
+
+    Supposing :code:`x` is a Tensor with shape [d_1, d_2, ..., d_n], the
+    :code:`y` is a mask with shape [d_1, d_2, ..., d_n, maxlen], where:
+
+    .. math::
+
+        y(i_1, i_2,..., i_n, j) = (j < x(i_1, i_2,..., i_n))
+
+    .. code-block:: text
+
+        Case:
+
+        Consider input:
+            x = [3, 1, 1, 0]    max_len = 4
+
+        then we get out:
+            mask = [[1, 1, 1, 0],
+                    [1, 0, 0, 0],
+                    [1, 0, 0, 0],
+                    [0, 0, 0, 0]]
+
+    Args:
+        x (Tensor): Input tensor of sequence_mask layer, \
+            whose elements are integers less than :code:`maxlen`. \
+            Tensor or Tensor with shape [d_1, d_2, ..., d_n].
+        maxlen (int, optional): Maximum length of the sequence. If :code:`maxlen` \
+                           is None, it would be replace with :math:`max(x)`.
+        dtype (np.dtype|paddle.dtype|str, optional): Data type of the output, \
+             ``int64`` by default.
+        name(str, optional): For detailed information, please refer \
+            to :ref:`api_guide_Name`. Usually name is no need to set and \
+            None by default.
+
+    Returns:
+            Tensor, The output sequence mask. Tensor with shape [d_1, d_2, ..., d_n, maxlen]
+            and data type of :code:`dtype`. The data type should be bool, float32, float64, int8,
+            int32 or int64.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> lengths = paddle.to_tensor([10, 9, 8])
+            >>> mask = paddle.nn.functional.sequence_mask(lengths)
+
+            >>> print(mask.numpy())
+            [[1 1 1 1 1 1 1 1 1 1]
+             [1 1 1 1 1 1 1 1 1 0]
+             [1 1 1 1 1 1 1 1 0 0]]
+
+    """
+
+    return paddle.nn.functional.sequence_mask(x, maxlen, dtype, name)

python/paddle/static/nn/static_pylayer.py

@@ -0,0 +1,604 @@
+# 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 paddle
+from paddle.base import core
+from paddle.base.backward import _append_grad_suffix_
+from paddle.base.framework import Variable, in_pir_mode
+from paddle.base.libpaddle.pir import build_pylayer_op, cf_yield
+from paddle.common_ops_import import LayerHelper, check_type, in_dygraph_mode
+from paddle.utils import flatten, map_structure
+
+# NOTE(MarioLulab): Borrowed from `python/paddle/static/nn/control_flow.py`
+from .control_flow import BlockGuard, copy_var_to_parent_block
+
+
+class StaticPyLayerBlockGuard(BlockGuard):
+    def __init__(self, block_manager):
+        check_type(
+            block_manager,
+            "block",
+            StaticPyLayerBlock,
+            "StaticPyLayerBlockGuard",
+        )
+        super().__init__(block_manager.helper.main_program)
+        self.block_manager = block_manager
+
+    def __enter__(self):
+        super().__enter__()
+        return self.block_manager
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.block_manager.complete()
+        return super().__exit__(exc_type, exc_val, exc_tb)
+
+
+class StaticPyLayerBlock:
+    def __init__(self, inputs, name=None, pylayer_context=None):
+        # used to specify the Variable type `Input` to `pylayer` op
+        self.fwd_inputs = [
+            each_input
+            for each_input in inputs
+            if isinstance(each_input, Variable)
+        ]  # filter non-Variable inputs
+
+        # used to specify the `Out` to `pylayer` op
+        self.fwd_outputs = []
+
+        self.context = pylayer_context
+
+        self.helper = LayerHelper("static_pylayer_block", name=name)
+        self.fwd_op_id = None
+        self._forward_block_id = None
+        self._backward_block_id = None
+        self.var_old_to_new = {}
+
+    def block(self, is_backward_block=False):
+        self.is_backward_block = is_backward_block
+        return StaticPyLayerBlockGuard(self)
+
+    @property
+    def forward_block_index(self):
+        return self._forward_block_id
+
+    @property
+    def backward_block_index(self):
+        return self._backward_block_id
+
+    @property
+    def fwd_op_index(self):
+        return self.fwd_op_id
+
+    def complete_forward_block(self):
+        inside_block = self.helper.main_program.current_block()
+        parent_block = self.helper.main_program.block(inside_block.parent_idx)
+        self._forward_block_id = inside_block.idx
+
+        step_scope = parent_block.create_var(
+            type=core.VarDesc.VarType.STEP_SCOPES
+        )
+
+        pylayer_op = parent_block.append_op(
+            type='pylayer',
+            inputs={
+                'Input': self.fwd_inputs,
+            },
+            outputs={"Out": self.fwd_outputs, "Scope": [step_scope]},
+            attrs={
+                'blocks': [inside_block],
+            },
+        )
+
+        self.fwd_op_id = pylayer_op.idx
+        self.helper.main_program._sync_with_cpp()
+
+    def complete_backward_block(self):
+        inside_block = self.helper.main_program.current_block()
+        parent_block = self.helper.main_program.block(inside_block.parent_idx)
+
+        self._backward_block_id = inside_block.idx
+        # Set OpRole to `backward`. The operators marked as `backward` are expected to be pruned in PruneBackward.
+        for op in inside_block.ops:
+            op_role_attr_name = (
+                core.op_proto_and_checker_maker.kOpRoleAttrName()
+            )
+            backward = core.op_proto_and_checker_maker.OpRole.Backward
+            op.desc._set_attr(op_role_attr_name, backward)
+        inside_block._set_forward_block_idx(self.forward_block_index)
+
+        # NOTE(MarioLulab): The reason of renaming the var name in the inside block is that
+        # we need to associating `inside_grads` and `outside_grads` at
+        # runtime `RunImpl` in pylayer op
+        _rename_var_recursively_(inside_block, self.var_old_to_new)
+
+        # update `blocks` attr by appending backward_block
+        forward_block_desc = parent_block.program.block(
+            self.forward_block_index
+        ).desc
+        backward_block_desc = inside_block.desc
+        parent_block.ops[self.fwd_op_index].desc.set_blocks_attr(
+            "blocks", [forward_block_desc, backward_block_desc]
+        )
+
+        # remove temporary vars created by `StaticPyLayerContext.saved_tensor`
+        if self.context:
+            for var in self.context.saved_vars:
+                if not inside_block.has_var(var.name):
+                    raise ValueError(
+                        f"{var.name} was saved in forward block but could not be found in backward block. Maybe {var.name} was renamed somewhere."
+                    )
+                inside_block._remove_var(var.name)
+
+        self.helper.main_program._sync_with_cpp()
+
+    def complete(self):
+        if not self.is_backward_block:
+            return self.complete_forward_block()
+        else:
+            return self.complete_backward_block()
+
+
+def _get_ctx_from_func_(func):
+    if func is None:
+        return None
+
+    fn_bind_args = getattr(func, "args", None)
+    if fn_bind_args is None:
+        return None
+
+    from paddle.jit.dy2static.py_layer import StaticPyLayerContext
+
+    fn_ctx = None
+    if len(fn_bind_args) > 0 and isinstance(
+        fn_bind_args[0], StaticPyLayerContext
+    ):
+        fn_ctx = fn_bind_args[0]
+
+    return fn_ctx
+
+
+def _rename_var_recursively_(cur_block, var_old_to_new):
+    """
+    Rename the var both the Variable instances and all ops' input and output arg names
+    in `cur_block` based on dict `var_old_to_new`.
+    Dict `var_old_to_new` should be the following format:
+    {
+        old_name_0 : new_name_0,
+        old_name_1 : new_name_1,
+        ...
+        old_name_n : new_name_n,
+    }
+    """
+
+    for old_var_name, new_var_name in var_old_to_new.items():
+        # NOTE(MarioLulab): The reason why not using `Block._rename_var`` is that `Block._rename_var` will raise ValueError, when `old_var_name` does not correspond to a Variable instance in Block.
+
+        if cur_block.has_var(old_var_name):
+            # `Block.desc._rename_var` can rename var in block and then rename var name in all ops
+            cur_block.desc._rename_var(
+                old_var_name.encode(), new_var_name.encode()
+            )
+        else:
+            # When cur_block does not have the var, `Block.desc._rename_var` can't rename var name in ops.
+            # In this case, we should traverse all ops and perform renaming manually.
+            for op in cur_block.ops:
+                op._rename_input(old_var_name, new_var_name)
+                op._rename_output(old_var_name, new_var_name)
+
+    # NOTE(MarioLulab): block attr type with the name of "blocks" or "sub_block" indicates
+    # the block might be executed. We should rename the var name in these blocks recursively
+    block_attr_names = ["blocks", "sub_block"]
+
+    for op in cur_block.ops:
+        for attr_name in op.all_attrs():
+            if attr_name not in block_attr_names:
+                continue
+
+            if op.attr_type(attr_name) == core.AttrType.BLOCK:
+                sub_block_id = op._block_attr_id(attr_name)
+                sub_block = cur_block.program.block(sub_block_id)
+                _rename_var_recursively_(sub_block, var_old_to_new)
+            elif op.attr_type(attr_name) == core.AttrType.BLOCKS:
+                sub_blocks_ids = op._blocks_attr_ids(attr_name)
+                for sub_block_id in sub_blocks_ids:
+                    sub_block = cur_block.program.block(sub_block_id)
+                    _rename_var_recursively_(sub_block, var_old_to_new)
+
+
+def copy_var_from_parent_block(parent_block_var, layer_helper):
+    if not isinstance(parent_block_var, Variable):
+        return parent_block_var
+    prog = layer_helper.main_program
+    current_block = prog.current_block()
+
+    if (
+        parent_block_var.type == core.VarDesc.VarType.DENSE_TENSOR_ARRAY
+        and current_block._find_var_recursive(parent_block_var.name)
+    ):
+        current_block_var = parent_block_var
+    else:
+        current_block_var = current_block.create_var(
+            dtype=parent_block_var.dtype,
+            shape=parent_block_var.shape,
+            type=parent_block_var.type,
+        )
+        paddle.assign(parent_block_var, current_block_var)
+    return current_block_var
+
+
+class PyLayerBackwardFunction:
+    _register_backward_funcs = []
+
+    def __init__(self, backward_function, hook_check_func):
+        if backward_function is None or not callable(backward_function):
+            raise TypeError('func must be a Python function')
+
+        self._func = backward_function
+
+        # Note: Used to verify the number of `Value` inputs to ``forward_fn`` the same as the
+        # number of `Value` outputs to ``backward_fn``, and the number of `Value` outputs to ``forward_fn``
+        # the same as the number of `Value` inputs to ``backward_fn``.
+        self._hook_check_func = hook_check_func
+
+        '''
+        Why record self here?
+           For increasing reference count of self.
+           It seems that to release Python object
+           whose reference count is 1 would cause
+           segmentation fault error in C++ side.
+           May be lack of Python GC in C++ side?
+        '''
+        PyLayerBackwardFunction._register_backward_funcs.append(self)
+
+    def __call__(self, *output_grads):
+        assert self._hook_check_func
+
+        input_grads = self._func(*output_grads)
+        if not isinstance(input_grads, (list, tuple)):
+            input_grads = (input_grads,)
+
+        self._hook_check_func(output_grads, input_grads)
+        input_grads = [
+            input_grad
+            for input_grad in flatten(input_grads)
+            if isinstance(input_grad, (paddle.pir.Value, type(None)))
+        ]
+
+        return input_grads
+
+
+def static_pylayer(forward_fn, inputs, backward_fn=None, name=None):
+    """
+    This API returns ``forward_fn(inputs)``, and two sub-block are created based on
+    the logic of ``forward_fn`` and ``backward_fn``, with the operator ``pylayer``
+    holding information about the two blocks.
+
+    ``forward_fn`` and ``backward_fn`` should return a nest structure of Variables.
+    A nest structure of Variables in PaddlePaddle is Variable(s), or tuple of Variables, or
+    list of Variables.
+
+    Note:
+        1. If ``backward_fn`` is not None, user needs to keep the number of `Variable` inputs to ``forward_fn`` the same as the
+        number of `Variable` outputs to ``backward_fn``, and the number of `Variable` outputs to ``forward_fn``
+        the same as the number of `Variable` inputs to ``backward_fn``.
+
+        2. If ``backward_fn`` is None, ``stop_gradient`` attr of all Variable in ``inputs`` is expected to be True.
+        Otherwise it might get unexpected results in backward propagation.
+
+        3. This API can only be used under static graph mode.
+
+    Args:
+        forward_fn (callable): A callable to be performed in forward propagation
+        inputs (list[Variable]): The list of input Variable to the ``forward_fn``
+        backward_fn (callable, optional): A callable to be performed in backward propagation. Default: None, which means no need to do backward propagation.
+        name (str, optional): The default value is ``None`` . Normally users
+            don't have to set this parameter. For more information, please
+            refer to :ref:`api_guide_Name` .
+
+    Returns:
+        Variable|list(Variable)|tuple(Variable): returns the output of ``forward_fn(inputs)``
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import numpy as np
+
+            >>> paddle.enable_static()
+
+            >>> def forward_fn(x):
+            ...     return paddle.exp(x)
+
+            >>> def backward_fn(dy):
+            ...     return 2 * paddle.exp(dy)
+
+            >>> main_program = paddle.static.Program()
+            >>> start_program = paddle.static.Program()
+
+            >>> place = paddle.CPUPlace()
+            >>> exe = paddle.static.Executor(place)
+            >>> with paddle.static.program_guard(main_program, start_program):
+            ...     data = paddle.static.data(name="X", shape=[None, 5], dtype="float32")
+            ...     data.stop_gradient = False
+            ...     ret = paddle.static.nn.static_pylayer(forward_fn, [data], backward_fn)
+            ...     data_grad = paddle.static.gradients([ret], data)[0]
+
+            >>> exe.run(start_program)
+            >>> x = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
+            >>> x, x_grad, y = exe.run(
+            ...     main_program,
+            ...     feed={"X": x},
+            ...     fetch_list=[data, data_grad, ret],
+            ... )
+
+            >>> print(x)
+            [[1. 2. 3. 4. 5.]]
+            >>> print(x_grad)
+            [[5.4365635 5.4365635 5.4365635 5.4365635 5.4365635]]
+            >>> print(y)
+            [[  2.7182817   7.389056   20.085537   54.59815   148.41316  ]]
+    """
+    assert in_dygraph_mode() is False, (
+        "please use PyLayer instead of static_pylayer in dygraph mode"
+    )
+
+    assert isinstance(inputs, list)
+    if backward_fn is None:
+        for input_var in inputs:
+            if input_var.stop_gradient is False:
+                raise ValueError(
+                    f"``stop_gradient`` attr of all inputs to ``forward_fn`` are expected to be True, when ``backward_fn == None``, but {input_var.name}.stop_gradient got {input_var.stop_gradient}"
+                )
+
+    # judge if in dy2st or not, by checking binding args of `forward_fn` and `backward_fn`
+    fwd_fn_ctx = _get_ctx_from_func_(forward_fn)
+    bwd_fn_ctx = _get_ctx_from_func_(backward_fn)
+    static_pylayer_context = (
+        fwd_fn_ctx if fwd_fn_ctx and (fwd_fn_ctx == bwd_fn_ctx) else None
+    )
+
+    if in_pir_mode():
+        fwd_inputs = [
+            inp for inp in flatten(inputs) if isinstance(inp, paddle.pir.Value)
+        ]
+        pylayer_op = build_pylayer_op(fwd_inputs)
+        outputs = None
+        if forward_fn is not None:
+            if not callable(forward_fn):
+                raise ValueError("`forward_fn` should be callable")
+            with pylayer_op.forward_block():
+                outputs = forward_fn(*inputs)
+
+            if outputs is None:
+                return None
+
+            fwd_outputs = [
+                out
+                for out in flatten(outputs)
+                if isinstance(out, paddle.pir.Value)
+            ]
+
+            with pylayer_op.forward_block():
+                if fwd_outputs is not None:
+                    cf_yield(flatten(fwd_outputs))
+            pylayer_op.update_output()
+        if backward_fn is not None:
+            if not callable(backward_fn):
+                raise ValueError("`backward_fn` should be callable")
+
+            def hook_inputs_outputs_check_function(output_grads, input_grads):
+                # 1. Verify the number of `Value` inputs to ``forward_fn`` the same as the
+                # number of `Value` outputs to ``backward_fn``
+                forward_inputs = [
+                    x
+                    for x in flatten(inputs)
+                    if isinstance(x, paddle.pir.Value)
+                ]
+                input_grads = [
+                    x
+                    for x in flatten(input_grads)
+                    if isinstance(x, (paddle.pir.Value, type(None)))
+                ]
+                if len(input_grads) != len(forward_inputs):
+                    raise ValueError(
+                        f"The number of input grads should be equal to the number of inputs, but got {len(input_grads)} and {len(forward_inputs)}."
+                    )
+                for inp_grad, fwd_input in zip(input_grads, forward_inputs):
+                    # NOTE: inp_grad will be None if fwd_input.stop_gradients=True
+                    if inp_grad is None:
+                        continue
+                    assert inp_grad.dtype == fwd_input.dtype, (
+                        f"dtype of inp_grad({inp_grad.dtype}) and fwd_input({fwd_input.dtype}) should be the same"
+                    )
+                    assert inp_grad.shape == fwd_input.shape, (
+                        f"shape of inp_grad({inp_grad.shape}) and fwd_input({fwd_input.shape}) should be the same"
+                    )
+                    if fwd_input.is_dist():
+                        # NOTE: placements may be not the same, so do not check it.
+                        assert inp_grad.is_dist(), (
+                            "fwd_input and inp_grad should both be distributed"
+                        )
+                        assert (
+                            fwd_input.dist_attr().process_mesh
+                            == inp_grad.dist_attr().process_mesh
+                        ), (
+                            f"process_mesh of fwd_input({fwd_input.dist_attr().process_mesh}) and inp_grad({inp_grad.dist_attr().process_mesh}) should be the same"
+                        )
+                    else:
+                        assert inp_grad.type() == fwd_input.type(), (
+                            f"type of inp_grad({inp_grad.type()}) and fwd_input({fwd_input.type()}) should be the same"
+                        )
+
+                # 2. Verify the number of `Value` outputs to ``forward_fn``
+                # the same as the number of `Value` inputs to ``backward_fn``
+                forward_outputs = [
+                    x
+                    for x in flatten(fwd_outputs)
+                    if isinstance(x, paddle.pir.Value)
+                ]
+                if len(output_grads) != len(forward_outputs):
+                    raise ValueError(
+                        f"The number of output grads should be equal to the number of outputs, but got {len(output_grads)} and {len(fwd_outputs)}."
+                    )
+                for out_grad, fwd_output in zip(output_grads, forward_outputs):
+                    if out_grad is None:
+                        continue
+                    assert out_grad.dtype == fwd_output.dtype, (
+                        f"dtype of out_grad({out_grad.dtype}) and fwd_output({fwd_output.dtype}) should be the same"
+                    )
+                    assert out_grad.shape == fwd_output.shape, (
+                        f"shape of out_grad({out_grad.shape}) and fwd_output({fwd_output.shape}) should be the same"
+                    )
+                    if fwd_output.is_dist():
+                        # NOTE: placements may be not the same, so do not check it.
+                        assert out_grad.is_dist(), (
+                            "fwd_output and out_grad should both be distributed"
+                        )
+                        assert (
+                            fwd_output.dist_attr().process_mesh
+                            == out_grad.dist_attr().process_mesh
+                        ), (
+                            f"process_mesh of fwd_output({fwd_output.dist_attr().process_mesh}) and out_grad({out_grad.dist_attr().process_mesh}) should be the same"
+                        )
+                    else:
+                        assert out_grad.type() == fwd_output.type(), (
+                            f"type of out_grad({out_grad.type}) and fwd_output({fwd_output.type}) should be the same"
+                        )
+
+            bwd_fn = PyLayerBackwardFunction(
+                backward_fn, hook_check_func=hook_inputs_outputs_check_function
+            )
+            pylayer_op.register_backward_function(bwd_fn)
+
+        # NOTE: Replace pir.Value of `outputs` with pylayer_op.result, because value of `outputs` which is inside pylayer block can't be reference outside the block.
+        op_result_idx = 0
+        outputs = flatten(outputs)
+        for i in range(len(outputs)):
+            if isinstance(outputs[i], paddle.pir.Value):
+                outputs[i] = pylayer_op.results()[op_result_idx]
+                op_result_idx += 1
+        return outputs[0] if len(outputs) == 1 else outputs
+
+    check_type(name, "name", (str, type(None)), "base.layers.static_pylayer")
+    helper = LayerHelper('static_pylayer', **locals())
+    copy_to_parent_func = lambda var: copy_var_to_parent_block(var, helper)
+
+    assert forward_fn is not None and callable(forward_fn)
+    pylayer_block_manager = StaticPyLayerBlock(
+        inputs, pylayer_context=static_pylayer_context
+    )
+    with pylayer_block_manager.block(is_backward_block=False) as mgr:
+        origin_output = forward_fn(*inputs)
+        if origin_output is not None:
+            output = map_structure(copy_to_parent_func, origin_output)
+            mgr.fwd_outputs = [
+                x for x in flatten(output) if isinstance(x, Variable)
+            ]
+        else:
+            mgr.fwd_outputs = []
+
+    current_block = helper.main_program.current_block()
+    current_block._sync_with_cpp()
+    if backward_fn is not None:
+        assert callable(backward_fn)
+        if origin_output is None:
+            output = []
+
+        # **Create the backward input** from the output of the op to build the
+        # backward block, and then delete it.
+        grad_var_ins = []
+        for fwd_var in pylayer_block_manager.fwd_outputs:
+            fwd_var_name = fwd_var.name
+            bwd_var_name = _append_grad_suffix_(fwd_var_name)
+            if not current_block.desc.has_var_recursive(fwd_var_name.encode()):
+                raise ValueError(
+                    f"Grad var {bwd_var_name} , we can't find its related forward var {fwd_var_name}"
+                )
+
+            var = current_block.create_var(
+                dtype=fwd_var.dtype,
+                shape=fwd_var.shape,
+                type=fwd_var.type,
+                name=bwd_var_name,
+            )
+
+            grad_var_ins.append(var)
+
+        copy_from_parent_func = lambda var: copy_var_from_parent_block(
+            var, helper
+        )
+        assert isinstance(grad_var_ins, list)
+        with pylayer_block_manager.block(is_backward_block=True) as mgr:
+            # Step1. Copy var from parent block
+            inside_block_inputs = map_structure(
+                copy_from_parent_func, grad_var_ins
+            )
+
+            # Step2. Do backward propagation
+            grad_origin_output = backward_fn(*inside_block_inputs)
+
+            if grad_origin_output is not None:
+                # Step3. Check the number of inputs to ``forward_fn`` the
+                # same as the number of outputs to ``backward_fn``
+                flat_grad_origin = flatten(grad_origin_output)
+
+                # NOTE(MarioLulab): ``current_block`` was defined outside
+                forward_input_names = current_block.ops[
+                    pylayer_block_manager.fwd_op_index
+                ].desc.input_arg_names()
+                assert len(forward_input_names) == len(flat_grad_origin), (
+                    f"needs to keep the number of inputs to ``forward_fn`` the same as the number of outputs to ``backward_fn``, \
+                    but got {len(forward_input_names)} and {len(flat_grad_origin)}"
+                )
+
+                # Step4. Rename var name with suffix of "@GRAD"
+                for bwd_output, fwd_input_name in zip(
+                    flat_grad_origin, forward_input_names
+                ):
+                    # NOTE(MarioLulab): Because `flat_grad_origin` are the Variables inside the backward block, which one by one corresponds
+                    # to the gradients of the inputs to the forward function, we need to establish a link between `flat_grad_origin`,
+                    # and the Variable outside the backward block which represent the gradient of the input ot the forward function.
+                    # The approach we have taken is renaming `flat_grad_origin` by forward input name with suffix of "@GRAD", and aligning
+                    # the order of `Out@GRAD` in `pylayer_grad` op with `flat_grad_origin`. And in the runtime `RunImpl` in `pylayer_grad` op,
+                    # we will find inside_grad with the name of forward input name with suffix of "@GRAD" in the scope, and assign `inside_grads`
+                    # to `outside_grads`.
+                    #
+                    # Example:
+                    # after run the code below to create forward and backward block:
+                    #
+                    #    out = forward_fn(x, y)                 # create forward block
+                    #    x_grad, y_grad = backward_fn(out_grad) # create backward block
+                    #
+                    # x.name is "X", y.name is "Y", and out.name is "tmp_0", but x_grad.name is "_generate_0", y_grad.name is "_generate_1".
+                    # we rename x_grad by "X@GRAD", and y_grad by "Y@GRAD" inside backward block.
+                    # One thing to keep in mind is that we assume there were no Variable naming "X@GRAD" inside backward block before performing rename operation.
+                    # TODO(MarioLulab): We will validate the assumption above is whether a strong hypothesis or not.
+
+                    # attach old var name into new
+                    if isinstance(bwd_output, Variable):
+                        bwd_out_new = _append_grad_suffix_(
+                            fwd_input_name
+                        )  # "X" => "X@GRAD"
+                        mgr.var_old_to_new[bwd_output.name] = (
+                            bwd_out_new  # e.g. "tmp_0.mean_0": "X@GRAD"
+                        )
+
+        # **Delete the backward input**
+        for bwd_var in grad_var_ins:
+            current_block._remove_var(bwd_var.name)
+
+    if origin_output is None:
+        return None
+
+    return output

python/paddle/static/pir_io.py

@@ -0,0 +1,946 @@
+#   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 errno
+import inspect
+import logging
+import os
+import pickle
+import sys
+import warnings
+
+import numpy as np
+
+import paddle
+from paddle import pir
+from paddle.autograd.backward_utils import (
+    ValueSet,
+    get_real_op_inputs,
+    get_real_op_outputs,
+    some_in_set,
+)
+from paddle.base import (
+    core,
+    default_main_program,
+)
+from paddle.base.executor import Executor, global_scope
+from paddle.base.framework import (
+    dygraph_not_support,
+    static_only,
+)
+from paddle.base.log_helper import get_logger
+from paddle.framework.io_utils import (
+    _pack_loaded_dict,
+    _pickle_loads_mac,
+    _unpack_saved_dict,
+)
+
+from .io_utils import (
+    _check_args,
+    _check_vars,
+    _get_valid_program,
+    _normalize_path_prefix,
+    _safe_load_pickle,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+def get_pir_parameters(program):
+    """
+    Get parameters and optimizer variables from program.
+        Args:
+            program(Program): The program to get parameters and optimizer variables.
+    """
+    params = []
+    opts = []
+    for var in program.list_vars():
+        if (
+            var.is_parameter
+            or var.get_defining_op().name() == "builtin.parameter"
+        ):
+            params.append(var)
+        elif var.persistable and var.get_defining_op().name() == "pd_op.data":
+            opts.append(var)
+    return params, opts
+
+
+def get_pir_feed_and_fetch(program):
+    feed_name_list = []
+    fetch_targets = []
+    for op in program.global_block().ops:
+        if op.name() == "pd_op.data" or op.name() == "pd_op.feed":
+            feed_name_list.append(op.attrs()["name"])
+        if op.name() == "pd_op.fetch":
+            fetch_targets.extend(op.operands_source())
+    return feed_name_list, fetch_targets
+
+
+def set_var(name, ndarray):
+    t = global_scope().find_var(name).get_tensor()
+    p = t._place()
+    if p.is_cpu_place():
+        place = paddle.base.CPUPlace()
+    elif p.is_cuda_pinned_place():
+        place = paddle.base.CUDAPinnedPlace()
+    elif p.is_xpu_place():
+        p = paddle.base.core.Place()
+        p.set_place(t._place())
+        place = paddle.base.XPUPlace(p.xpu_device_id())
+    elif p.is_custom_place():
+        p = paddle.base.core.Place()
+        p.set_place(t._place())
+        place = paddle.base.CustomPlace(
+            paddle.device.get_device().split(':')[0], p.custom_device_id()
+        )
+    else:
+        p = paddle.base.core.Place()
+        p.set_place(t._place())
+        place = paddle.base.CUDAPlace(p.gpu_device_id())
+
+    t.set(ndarray, place)
+
+
+def append_pir_feed_ops(program, feed_vars):
+    """
+    Append feed ops to the program.
+    Args:
+        program(Program): Specify a program you want to append fetch op.
+        feed_vars(Value | list[Value]): Values should be feed.
+    Returns:
+        modify program
+    """
+    for i, var in enumerate(feed_vars):
+        orig_op = var.get_defining_op()
+        if orig_op.name() != 'pd_op.feed' and orig_op.name() != 'pd_op.data':
+            value = paddle._pir_ops.data(
+                "feed_name_" + str(i),
+                var.shape,
+                var.dtype,
+                paddle.base.core.Place(),
+            )
+            var.replace_all_uses_with(value)
+            value.get_defining_op().move_before(orig_op)
+
+    for i, var in enumerate(feed_vars):
+        orig_op = var.get_defining_op()
+        if orig_op.name() != 'pd_op.feed' and orig_op.name() != 'pd_op.data':
+            orig_op.get_parent_block().remove_op(orig_op)
+
+
+def append_pir_fetch_ops(program, fetch_name_var_maps):
+    """
+    Append fetch ops to the program.
+    Args:
+        program(Program): Specify a program you want to append fetch op.
+        fetch_vars(Tensor | list[Tensor]): Values returned by inference.
+    Returns:
+        modify program
+    """
+    for i, (var, name) in enumerate(fetch_name_var_maps):
+        out = paddle._pir_ops.fetch(var, name, i)
+        out.persistable = True
+
+
+def pir_prune_with_input(program, feed_vars, target_vars):
+    """
+    Prune a program according to feed_vars and target_vars.
+    Args:
+        program(Program): Specify a program you want to prune.
+        feed_vars(Tensor | list[Tensor]): Values needed by inference.
+        target_vars(Tensor | list[Tensor]): Values returned by inference.
+    Returns
+        modify program
+    """
+    if not isinstance(program, paddle.static.Program):
+        raise TypeError(
+            f"program type must be `paddle.static.Program`, but received `{type(program)}`"
+        )
+
+    total_ops = program.global_block().ops
+    intersection_op_flags = [True] * len(total_ops)
+
+    # from output to input
+    target_vars_ = ValueSet(target_vars)
+    for i, op in reversed(list(enumerate(total_ops))):
+        if some_in_set(get_real_op_outputs(op), target_vars_):
+            for operand in get_real_op_inputs(op):
+                target_vars_.add(operand)
+        else:
+            intersection_op_flags[i] = False
+
+    for i, op in reversed(list(enumerate(total_ops))):
+        if not intersection_op_flags[i]:
+            if some_in_set(get_real_op_outputs(op), ValueSet(feed_vars)):
+                raise ValueError(
+                    f"The feed_var create by: '{op.name()}' is not involved in the target_vars calculation"
+                    f"Please remove it from feed_vars ."
+                )
+            program.global_block().remove_op(op)
+
+
+def _inference_optimize(program, prune_read_op=True):
+    """
+    This method will create a new program and do following adjustments on it:
+    1. Remove all reader variables and their creator ops if exist.
+
+    2. Remove the :code:`read_op` if exists.
+
+    3. change the :code:`is_test`
+    attribute of operators to :code:`True`. All the :code:`Parameter`
+    information will be lost.
+
+    Args:
+        prune_read_op(bool): remove the read ops that are added by py_reader
+                             for cpp inference library
+
+    Notes: This API is a very low level API. Use
+    :code:`Program.clone(for_test=True)` instead.
+
+    Returns:
+        Program: The new program.
+    """
+
+    # remove all readers and the read_op if exist
+    if prune_read_op:
+        pass
+
+    # change all `is_test` attributes to True
+    for block in program.blocks:
+        for op in block.ops:
+            if op.has_attr("is_test"):
+                op.set_bool_attr("is_test", True)
+            if op.name() == "pd_op.batch_norm":
+                # Remove the output ReserveSpace of batch_norm if exists.
+                pass
+
+
+def normalize_pir_program(program, feed_vars, fetch_vars, **kwargs):
+    """
+
+    Normalize/Optimize a program according to feed_vars and fetch_vars.
+
+    Args:
+        program(Program): Specify a program you want to optimize.
+        feed_vars(Tensor | list[Tensor]): Values needed by inference.
+        fetch_vars(Tensor | list[Tensor]): Values returned by inference.
+        kwargs: Supported keys including ``skip_prune_program``.
+            - skip_prune_program(bool): whether to skip pruning program. Defaults to False.
+
+    Returns:
+        Program: Normalized/Optimized program.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> paddle.enable_static()
+
+            >>> path_prefix = "./infer_model"
+
+            # User defined network, here a softmax regression example
+            >>> image = paddle.static.data(name='img', shape=[None, 28, 28], dtype='float32')
+            >>> label = paddle.static.data(name='label', shape=[None, 1], dtype='int64')
+            >>> predict = paddle.static.nn.fc(image, 10, activation='softmax')
+
+            >>> loss = paddle.nn.functional.cross_entropy(predict, label)
+
+            >>> exe = paddle.static.Executor(paddle.CPUPlace())
+            >>> exe.run(paddle.static.default_startup_program())
+
+            # normalize main program.
+            >>> program = paddle.static.default_main_program()
+            >>> normalized_program = paddle.static.normalize_program(program, [image], [predict])
+
+    """
+    if not isinstance(program, paddle.static.Program):
+        raise TypeError(
+            f"program type must be `paddle.static.Program`, but received `{type(program)}`"
+        )
+    if not isinstance(feed_vars, list):
+        feed_vars = [feed_vars]
+    if not all(isinstance(v, pir.Value) for v in feed_vars):
+        raise TypeError("feed_vars type must be a Value or a list of Value.")
+    if not isinstance(fetch_vars, list):
+        fetch_vars = [fetch_vars]
+    if not all(isinstance(v, pir.Value) for v in fetch_vars):
+        raise TypeError("fetch_vars type must be a Value or a list of Value.")
+
+    if len(program.global_block().ops) == 0:
+        raise ValueError(
+            "program must not be empty. at least one operator is required!"
+        )
+
+    # remind users to set auc_states to 0 if auc op were found.
+    for op in program.global_block().ops:
+        if op.name() == 'pd_op.auc':
+            warnings.warn(
+                "Be sure that you have set auc states to 0 before saving inference model."
+            )
+            break
+
+    # serialize program
+    value_map = paddle.pir.IrMapping()
+    copy_program = program.clone(value_map)
+    global_block = copy_program.global_block()
+    clone_feed_vars = [value_map.look_up(v) for v in feed_vars]
+    clone_fetch_vars = [value_map.look_up(v) for v in fetch_vars]
+
+    for op in global_block.ops:
+        # can not delete feed op because it's output used by other op.
+        if op.name() == "pd_op.fetch":
+            global_block.remove_op(op)
+
+    skip_prune_program = kwargs.get('skip_prune_program', False)
+    # if feed var is not connect with target_vars, it will be delete.
+    if not skip_prune_program:
+        pir_prune_with_input(copy_program, clone_feed_vars, clone_fetch_vars)
+    _inference_optimize(copy_program, prune_read_op=True)
+
+    fetch_vars_tuple = []
+    for i, var in enumerate(clone_fetch_vars):
+        scale_op = var.get_defining_op()
+        orig_var = var
+        if scale_op.name() == "pd_op.scale":
+            full_op = scale_op.operand_source(1).get_defining_op()
+            if full_op.has_attr("value") and full_op.attrs()['value'] == 1.0:
+                orig_var = scale_op.operand_source(0)
+        if orig_var.has_name:
+            fetch_vars_tuple.append((orig_var, orig_var.name))
+        else:
+            fetch_vars_tuple.append((var, "fetch_name_" + str(i)))
+    with paddle.static.program_guard(copy_program):
+        append_pir_feed_ops(copy_program, clone_feed_vars)
+        append_pir_fetch_ops(copy_program, fetch_vars_tuple)
+
+    return copy_program
+
+
+@dygraph_not_support
+def save_vars_pir(
+    dirname,
+    main_program=None,
+    vars=None,
+    filename=None,
+):
+    """
+    Save specific variables in the `Program` to files.
+
+    There are two ways to specify the variables to be saved: set variables in
+    a list and assign it to the `vars`, or use the `predicate` function to select
+    variables that make `predicate(variable) == True`. The first way has a higher priority.
+
+    The `dirname` is used to specify the folder where to save variables.
+    If you prefer to save variables in separate files in the `dirname` folder,
+    do not set `filename`. If you prefer to save all variables in a single file,
+    use `filename` to specify it.
+
+    Args:
+        dirname(str, optional): The folder to save variables.
+                            When you need to save the parameter to the memory, set it to None.
+        main_program(Program, optional): The program whose variables will be saved.
+                                    If it is None, the default main program will
+                                    be used automatically.
+                                    Default: None
+        vars(list[Variable], optional): The list contains all variables to be saved.
+                                        Default: None
+        filename(str, optional): If you prefer to save all variables in a single file,
+                                 use `filename` to specify it. Otherwise, let `filename` be None.
+                                 Default: None
+
+    Returns:
+        str: When saving parameters to a file, returns None.
+             When saving parameters to memory, returns a binary string containing parameters.
+    """
+
+    save_to_memory = False
+    if dirname is None and filename is None:
+        save_to_memory = True
+
+    main_program = _get_valid_program(main_program)
+
+    if vars is None:
+        param, opt = get_pir_parameters(main_program)
+        vars_list = param + opt
+        return save_vars_pir(
+            main_program=main_program,
+            dirname=dirname,
+            vars=[var for var in vars_list if var.persistable],
+            filename=filename,
+        )
+    else:
+        params_var_name = "saved_params"
+        # give warning when there is no var in model
+        if len(list(vars)) == 0:
+            warnings.warn(
+                "no variable in your model, please ensure there are any variables in your model to save"
+            )
+            return None
+
+        save_var_map = {}
+        for v in vars:
+            var = global_scope().find_var(v.name)
+            # TODO(chenzhiyang): deal with RAW type and sparse
+            if filename is None and save_to_memory is False:
+                save_file_path = os.path.join(os.path.normpath(dirname), v.name)
+                core.save_func(
+                    var.get_tensor(), v.name, save_file_path, True, False
+                )
+            else:
+                save_var_map[v.name] = var.get_tensor()
+
+        if filename is not None or save_to_memory:
+            save_var_list = []
+            save_var_names = []
+            for name in sorted(save_var_map.keys()):
+                save_var_list.append(save_var_map[name])
+                save_var_names.append(name)
+
+            save_path = ''
+            if save_to_memory is False:
+                save_path = os.path.join(os.path.normpath(dirname), filename)
+            core.save_combine_func(
+                save_var_list,
+                save_var_names,
+                save_path,
+                True,
+                False,
+                save_to_memory,
+            )
+
+        if save_to_memory:
+            return global_scope().find_var(params_var_name).get_bytes()
+
+
+def load_vars_pir(
+    executor,
+    dirname,
+    main_program=None,
+    vars=None,
+    filename=None,
+):
+    """
+    :api_attr: PIR Static Graph
+
+    This API loads variables from files by C++ function.
+
+    There are two ways to specify the variables to be loaded: the first way, set
+    variables in a list and assign it to the `vars`; the second way, use the
+    `predicate` function to select variables that make `predicate(variable) == True`.
+    The first way has a higher priority.
+
+    The `dirname` is used to specify the folder where to load variables.
+    If variables were saved in separate files in the folder `dirname`,
+    set `filename` None. If all variables were saved in a single file,
+    use `filename` to specify it.
+
+    Args:
+        executor(Executor): The executor to create variables in scope.
+        dirname(str): The folder where to load the variables.
+        main_program(Program, optional): The program whose variables will be loaded.
+                                    If it is None, the default main program will
+                                    be used automatically.
+                                    Default: None
+        vars(list[Variable], optional): The list that contains all variables to be loaded.
+                                   Default: None
+        filename(str, optional): The file which saved all required variables. If variables
+                                were saved in separate files, set it to be None.
+                                Default: None
+
+    Returns:
+        None
+    """
+    assert executor is None or isinstance(executor, Executor)
+
+    vars_from_memory = False
+    if dirname is not None:
+        dirname = os.path.normpath(dirname)
+    # TODO(chenzhiyang): vars_from_memory
+
+    if filename == '':
+        filename = None
+
+    if vars is None:
+        if main_program is None:
+            main_program = default_main_program()
+
+        if not isinstance(main_program, paddle.static.Program):
+            raise TypeError(
+                f"The type of input main_program is invalid, expected type is paddle.static.Program, but received {type(main_program)}"
+            )
+        param, opt = get_pir_parameters(main_program)
+        vars = param + opt
+        paddle.base.libpaddle.pir.create_loaded_parameter(
+            vars, global_scope(), executor._default_executor
+        )
+        load_vars_pir(
+            executor,
+            dirname=dirname,
+            main_program=main_program,
+            vars=[var for var in vars if var.persistable],
+            filename=filename,
+        )
+    else:
+        if main_program is None:
+            main_program = default_main_program()
+
+        if not isinstance(main_program, paddle.static.Program):
+            raise TypeError(
+                f"The type of input main_program is invalid, expected type is paddle.static.Program, but received {type(main_program)}"
+            )
+
+        # TODO(chenzhiyang):save origin param shape, check vars
+        load_var_map = {}
+
+        for v in vars:
+            var = global_scope().find_var(v.name)
+            assert isinstance(var, paddle.base.libpaddle.Variable)
+            if filename is None:
+                if dirname is None:
+                    raise ValueError(
+                        "The directory path and params cannot be None at the same time."
+                    )
+                file_path = os.path.join(dirname, v.name)
+                core.load_func(
+                    file_path,
+                    -1,
+                    [],
+                    False,
+                    var.get_tensor(),
+                    executor._default_executor.get_place(),
+                )
+            else:
+                load_var_map[v.name] = var
+
+        if filename is not None:
+            load_var_list = []
+            load_var_names = []
+            for name in sorted(load_var_map.keys()):
+                load_var_list.append(load_var_map[name].get_tensor())
+                load_var_names.append(name)
+
+            if vars_from_memory is False:
+                filename = os.path.join(dirname, filename)
+
+            core.load_combine_func(
+                filename,
+                load_var_names,
+                load_var_list,
+                False,
+                executor._default_executor.get_place(),
+            )
+            for name, var in zip(load_var_names, load_var_list):
+                set_var(name, np.array(var))
+
+
+@static_only
+def save_pir(program, model_path, protocol=4, **configs):
+    """
+    This function saves parameters, optimizer information and network description to model_path.
+
+    The parameters contain all the trainable Tensor, and save to a file with suffix ".pdparams".
+    The optimizer information contains all the Tensor used by optimizer. For Adam optimizer, contains beta1, beta2, momentum etc. All the information will be saved to a file with suffix ".pdopt". (If the optimizer has no Tensor to save (like SGD), the file will not be generated).
+    The network description is the description of the program. It's only used for deployment. The description will be saved to a file with a suffix ".pdmodel".
+
+    Args:
+        program(Program) : The program to be saved.
+        model_path(str): The file prefix to save the program. The format is "dirname/file_prefix". If file_prefix is an empty str, an exception will be raised.
+        protocol(int, optional): The protocol version of pickle module must be greater than 1 and less than 5.
+                                 Default: 4
+        configs(dict, optional) : Optional keyword arguments.
+
+    Returns:
+        None
+    """
+
+    base_name = os.path.basename(model_path)
+    assert base_name != "", (
+        "The input model_path MUST be format of dirname/filename [dirname\\filename in Windows system], but received model_path is empty string."
+    )
+    if 'pickle_protocol' in configs:
+        protocol = configs['pickle_protocol']
+        warnings.warn(
+            "'pickle_protocol' is a deprecated argument. Please use 'protocol' instead."
+        )
+
+    if not isinstance(protocol, int):
+        raise ValueError(
+            f"The 'protocol' MUST be `int`, but received {type(protocol)}"
+        )
+
+    if protocol < 2 or protocol > 4:
+        raise ValueError(
+            f"Expected 1<'protocol'<5, but received protocol={protocol}"
+        )
+
+    dir_name = os.path.dirname(model_path)
+    if dir_name and not os.path.exists(dir_name):
+        os.makedirs(dir_name)
+
+    def get_tensor(var):
+        t = global_scope().find_var(var.name).get_tensor()
+        return np.array(t)
+
+    # get parameters and optimizer variables
+    parameter_list, optimizer_param_list = get_pir_parameters(program)
+    param_dict = {
+        var.name: get_tensor(var) for var in parameter_list if var.persistable
+    }
+    opt_dict = {
+        var.name: get_tensor(var)
+        for var in optimizer_param_list
+        if var.persistable
+    }
+
+    # save parameters
+    param_dict = _unpack_saved_dict(param_dict, protocol)
+
+    # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
+    if sys.platform == 'darwin' and sys.version_info.major == 3:
+        pickle_bytes = pickle.dumps(param_dict, protocol=protocol)
+        with open(model_path + ".pdparams", 'wb') as f:
+            max_bytes = 2**30
+            f.writelines(
+                pickle_bytes[i : i + max_bytes]
+                for i in range(0, len(pickle_bytes), max_bytes)
+            )
+    else:
+        with open(model_path + ".pdparams", 'wb') as f:
+            pickle.dump(param_dict, f, protocol=protocol)
+
+    # save optimizer parameters
+    with open(model_path + ".pdopt", 'wb') as f:
+        pickle.dump(opt_dict, f, protocol=protocol)
+
+    # save program
+    paddle.core.serialize_pir_program(program, model_path + ".json")
+
+
+@static_only
+def load_pir(program, model_prefix, executor=None, var_list=None):
+    """
+    :api_attr: PIR Static Graph
+
+    This function gets parameters and optimizer information from program, and then gets corresponding value from file.
+    An exception will be thrown if shape or dtype of the parameters does not match.
+
+    This function can also load model file saved with [ save_params, save_persistables, save_vars ].
+    var_list can not be None when loading a single model file
+    ( filename is not None when save_params, save_persistables or save_vars is called ).
+
+    Args:
+        program(Program): The program to be loaded
+        model_prefix(str): The file prefix to store the program
+        executor(Executor, optional): The executor used for initializing the parameter
+                                      when startup program is not run.
+        var_list(list|tuple, optional): The Tensor list/tuple to load a single model file saved with
+                                  [ save_params, save_persistables, save_vars ].
+                                  Default: None
+
+    Returns:
+        None
+    """
+
+    assert executor is None or isinstance(executor, Executor)
+
+    parameter_file_name = model_prefix + ".pdparams"
+
+    # TODO(chenzhiyang):if not os.path.exists(parameter_file_name): load_vars
+
+    parameter_list, optimizer_param_list = get_pir_parameters(program)
+
+    with open(parameter_file_name, 'rb') as f:
+        # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
+        if sys.platform == 'darwin' and sys.version_info.major == 3:
+            load_dict = _pickle_loads_mac(parameter_file_name, f)
+        else:
+            load_dict = _safe_load_pickle(f, encoding='latin1')
+        load_dict = _pack_loaded_dict(load_dict)
+    for var in parameter_list:
+        if var.persistable:
+            assert var.name in load_dict, (
+                f"Can not find [{var.name}] in model file [{parameter_file_name}]"
+            )
+            set_var(var.name, load_dict[var.name])
+
+    if len(optimizer_param_list) > 0:
+        opt_file_name = model_prefix + ".pdopt"
+        assert os.path.exists(opt_file_name), (
+            f"Optimizer file [{opt_file_name}] not exits"
+        )
+
+        if executor:
+            paddle.base.libpaddle.pir.create_loaded_parameter(
+                optimizer_param_list, global_scope(), executor._default_executor
+            )
+
+        with open(opt_file_name, 'rb') as f:
+            load_dict = _safe_load_pickle(f, encoding='latin1')
+        for var in optimizer_param_list:
+            if var.persistable:
+                assert var.name in load_dict, (
+                    f"Can not find [{var.name}] in model file [{opt_file_name}]"
+                )
+                set_var(var.name, load_dict[var.name])
+
+
+@static_only
+def save_inference_model_pir(
+    path_prefix, feed_vars, fetch_vars, executor, **kwargs
+):
+    """
+    Save current model and its parameters to given path. i.e.
+    Given ``path_prefix = "PATH/modelname"``, after invoking
+    ``save_inference_model(path_prefix, feed_vars, fetch_vars, executor)``,
+    you will find two files named ``modelname.pdmodel`` and ``modelname.pdiparams``
+    under ``PATH``, which represent your model and parameters respectively.
+
+    Args:
+        path_prefix(str): Directory path to save model + model name without suffix.
+        feed_vars(Tensor | list[Tensor]): Variables needed by inference.
+        fetch_vars(Tensor | list[Tensor]): Variables returned by inference.
+        executor(Executor): The executor that saves the inference model. You can refer
+                            to :ref:`api_guide_executor_en` for more details.
+        kwargs: Supported keys including 'program' and "clip_extra". Attention please, kwargs is used for backward compatibility mainly.
+
+            - program(Program): specify a program if you don't want to use default main program.
+
+            - clip_extra(bool): the flag indicating whether to clip extra information for every operator. Default: True.
+
+            - legacy_format(bool): whether to save inference model in legacy format. Default: False.
+
+    Returns:
+        None
+    """
+    # check path_prefix, set model_path and params_path
+    path_prefix = _normalize_path_prefix(path_prefix)
+    try:
+        # mkdir may conflict if pserver and trainer are running on the same machine
+        dirname = os.path.dirname(path_prefix)
+        os.makedirs(dirname)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+    model_path = path_prefix + ".json"
+    params_path = path_prefix + ".pdiparams"
+    if os.path.isdir(model_path):
+        raise ValueError(f"'{model_path}' is an existing directory.")
+    if os.path.isdir(params_path):
+        raise ValueError(f"'{params_path}' is an existing directory.")
+
+    # verify feed_vars
+    _check_vars('feed_vars', feed_vars)
+    # verify fetch_vars
+    _check_vars('fetch_vars', fetch_vars)
+
+    program = _get_valid_program(kwargs.get('program', None))
+    # serialize and save program
+    program = normalize_pir_program(
+        program,
+        feed_vars,
+        fetch_vars,
+        skip_prune_program=kwargs.get('skip_prune_program', False),
+    )
+
+    readable = kwargs.get('readable', False)
+    trainable = kwargs.get('trainable', True)
+    paddle.core.serialize_pir_program(
+        program,
+        (
+            os.path.join(os.path.dirname(model_path), "__model__.json")
+            if kwargs.get('separate_parameters', False)
+            else model_path
+        ),
+        True,
+        readable,
+        trainable,
+    )
+
+    # serialize and save params
+    save_dirname = os.path.dirname(params_path)
+    params_filename = os.path.basename(params_path)
+    save_vars_pir(
+        dirname=save_dirname,
+        main_program=program,
+        filename=(
+            None
+            if kwargs.get('separate_parameters', False)
+            else params_filename
+        ),
+    )
+
+
+@static_only
+def load_inference_model_pir(path_prefix, executor, **kwargs):
+    """
+
+    Load inference model from a given path. By this API, you can get the model
+    structure(Inference Program) and model parameters.
+
+    Args:
+        path_prefix(str | None): One of the following:
+          - Directory path to save model + model name without suffix.
+          - Set to None when reading the model from memory.
+        executor(Executor): The executor to run for loading inference model.
+                            See :ref:`api_guide_executor_en` for more details about it.
+        kwargs: Supported keys including 'model_filename', 'params_filename'. Attention please, kwargs is used for backward compatibility mainly.
+
+            - model_filename(str): specify model_filename if you don't want to use default name.
+
+            - params_filename(str): specify params_filename if you don't want to use default name.
+
+    Returns:
+        list: The return of this API is a list with three elements:
+        (program, feed_target_names, fetch_targets). The `program` is a
+        ``Program`` (refer to :ref:`api_guide_Program_en`), which is used for inference.
+        The `feed_target_names` is a list of ``str``, which contains names of variables
+        that need to feed data in the inference program. The `fetch_targets` is a list of
+        ``Variable`` (refer to :ref:`api_guide_Program_en`). It contains variables from which
+        we can get inference results.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import numpy as np
+
+            >>> paddle.enable_static()
+
+            # Build the model
+            >>> startup_prog = paddle.static.default_startup_program()
+            >>> main_prog = paddle.static.default_main_program()
+            >>> with paddle.static.program_guard(main_prog, startup_prog):
+            ...     image = paddle.static.data(name="img", shape=[64, 784])
+            ...     w = paddle.create_parameter(shape=[784, 200], dtype='float32')
+            ...     b = paddle.create_parameter(shape=[200], dtype='float32')
+            ...     hidden_w = paddle.matmul(x=image, y=w)
+            ...     hidden_b = paddle.add(hidden_w, b)
+            >>> exe = paddle.static.Executor(paddle.CPUPlace())
+            >>> exe.run(startup_prog)
+
+            # Save the inference model
+            >>> path_prefix = "./infer_model"
+            >>> paddle.static.save_inference_model(path_prefix, [image], [hidden_b], exe)
+
+            >>> [inference_program, feed_target_names, fetch_targets] = paddle.static.load_inference_model(path_prefix, exe)
+            >>> tensor_img = np.array(np.random.random((64, 784)), dtype=np.float32)
+            >>> results = exe.run(
+            ...     inference_program,
+            ...     feed={feed_target_names[0]: tensor_img},
+            ...     fetch_list=fetch_targets,
+            ... )
+
+            # In this example, the inference program was saved in file
+            # "./infer_model.pdmodel" and parameters were saved in file
+            # " ./infer_model.pdiparams".
+            # By the inference program, feed_target_names and
+            # fetch_targets, we can use an executor to run the inference
+            # program to get the inference result.
+    """
+    # check kwargs
+    supported_args = ('model_filename', 'params_filename')
+    deprecated_args = ('pserver_endpoints',)
+    caller = inspect.currentframe().f_code.co_name
+    _check_args(caller, kwargs, supported_args, deprecated_args)
+
+    # load from memory
+    if path_prefix is None:
+        _logger.warning(
+            "Load inference model from memory is deprecated. Please specify path_prefix."
+        )
+        model_filename = kwargs.get('model_filename', None)
+        params_filename = kwargs.get('params_filename', None)
+        if params_filename is None:
+            raise ValueError(
+                "params_filename cannot be None when path_prefix is None."
+            )
+
+        # deserialize bytes to program
+        program = paddle.static.Program()
+        paddle.base.core.deserialize_pir_program(model_filename, program)
+
+        params, opts = get_pir_parameters(program)
+        vars = params + opts
+        vars = [var for var in vars if var.persistable]
+        if len(vars) > 0:
+            load_vars_pir(
+                # load from memory, dirname is None
+                executor,
+                dirname=None,
+                main_program=program,
+                filename=params_filename,
+            )
+    # load from file
+    else:
+        # check and norm path_prefix
+        path_prefix = _normalize_path_prefix(path_prefix)
+        dir_path = os.path.dirname(path_prefix)
+        if not os.path.isdir(dir_path):
+            raise ValueError(f"There is no directory named {dir_path}")
+        # set model_path and params_path in new way,
+        # path_prefix represents a file path without suffix in this case.
+        if not kwargs:
+            model_path = path_prefix + ".json"
+            params_path = path_prefix + ".pdiparams"
+        # set model_path and params_path in old way for compatible,
+        # path_prefix represents a directory path.
+        else:
+            model_filename = kwargs.get('model_filename', None)
+            params_filename = kwargs.get('params_filename', None)
+            # set model_path
+            if model_filename is None:
+                model_path = os.path.join(path_prefix, "__model__")
+            else:
+                model_path = os.path.join(path_prefix, model_filename + ".json")
+
+                if not os.path.exists(model_path):
+                    model_path = os.path.join(path_prefix, model_filename)
+            # set params_path
+            if params_filename is None:
+                params_path = os.path.join(path_prefix, "")
+            else:
+                params_path = os.path.join(
+                    path_prefix, params_filename + ".pdiparams"
+                )
+                if not os.path.exists(params_path):
+                    params_path = os.path.join(path_prefix, params_filename)
+            _logger.warning(
+                "The old way to load inference model is deprecated. Please specify path_prefix."
+                f" model path: {model_path}, params path: {params_path}"
+            )
+
+        # deserialize bytes to program
+        program = paddle.static.Program()
+        paddle.base.core.deserialize_pir_program(model_path, program)
+        # load parameters
+        params, opts = get_pir_parameters(program)
+        vars = params + opts
+        vars = [var for var in vars if var.persistable]
+        if len(vars) > 0:
+            load_dirname = os.path.dirname(params_path)
+            params_filename = os.path.basename(params_path)
+
+            load_vars_pir(
+                executor,
+                dirname=load_dirname,
+                main_program=program,
+                filename=params_filename,
+            )
+
+    feed_names, fetch_targets = get_pir_feed_and_fetch(program)
+    return [program, feed_names, fetch_targets]

python/paddle/static/python_op.py

@@ -0,0 +1,390 @@
+# Copyright (c) 2025 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 inspect
+import sys
+import types
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from functools import cached_property, partial, wraps
+from typing import Any, Generic, TypeVar, overload
+
+from typing_extensions import ParamSpec
+
+import paddle
+from paddle import _C_ops
+
+HAS_VAR_ARGS_OR_KWARGS: int = inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
+
+
+P1 = ParamSpec("P1")
+R1 = TypeVar("R1")
+
+
+class MissingArgument:
+    def __init__(self, fn: Callable[P1, R1], name: str):
+        self.fn = fn
+        self.name = name
+
+    def __repr__(self):
+        return f"<Required parameter '{self.name}' for function {self.fn.__name__}>"
+
+
+def extract_default(fn: Callable[P1, R1], parameter: inspect.Parameter):
+    if parameter.kind is inspect.Parameter.VAR_POSITIONAL:
+        return ()
+    elif parameter.kind is inspect.Parameter.VAR_KEYWORD:
+        return {}
+    elif parameter.default is inspect.Parameter.empty:
+        return MissingArgument(fn, parameter.name)
+    return parameter.default
+
+
+def get_fn_defaults_params(fn: Callable[P1, R1]) -> tuple:
+    fn_defaults_params = [
+        extract_default(fn, param)
+        for param in inspect.signature(fn).parameters.values()
+    ]
+    for i, default in enumerate(fn_defaults_params):
+        if not isinstance(default, MissingArgument):
+            fn_defaults_params = fn_defaults_params[i:]
+            break
+    return tuple(fn_defaults_params)
+
+
+def eliminate_positional_or_keyword_only(
+    fn: Callable[P1, R1],
+) -> Callable[P1, R1]:
+    assert isinstance(fn, types.FunctionType), "Only support regular function"
+    code = fn.__code__
+    co_flags: int = code.co_flags & ~HAS_VAR_ARGS_OR_KWARGS
+
+    argcount = (
+        code.co_argcount
+        + code.co_kwonlyargcount
+        + bool(code.co_flags & inspect.CO_VARARGS)
+        + bool(code.co_flags & inspect.CO_VARKEYWORDS)
+    )
+    if sys.version_info >= (3, 11):
+        new_code = types.CodeType(
+            argcount,  # co_argcount
+            0,  # posonlyargcount, eliminated
+            0,  # kwonlyargcount, eliminated
+            code.co_nlocals,
+            code.co_stacksize,
+            co_flags,
+            code.co_code,
+            code.co_consts,
+            code.co_names,
+            code.co_varnames,
+            code.co_filename,
+            code.co_name,
+            code.co_qualname,
+            code.co_firstlineno,
+            code.co_linetable,
+            code.co_exceptiontable,
+            code.co_freevars,
+            code.co_cellvars,
+        )
+    else:
+        new_code = types.CodeType(
+            argcount,  # co_argcount
+            0,  # posonlyargcount, eliminated
+            0,  # kwonlyargcount, eliminated
+            code.co_nlocals,
+            code.co_stacksize,
+            co_flags,
+            code.co_code,
+            code.co_consts,
+            code.co_names,
+            code.co_varnames,
+            code.co_filename,
+            code.co_name,
+            code.co_firstlineno,
+            code.co_linetable,
+            code.co_freevars,
+            code.co_cellvars,
+        )
+
+    fn_defaults_params = get_fn_defaults_params(fn)
+    new_fn = types.FunctionType(
+        new_code,
+        fn.__globals__,
+        fn.__name__,
+        fn_defaults_params,
+        fn.__closure__,
+    )
+    new_fn.__name__ = fn.__name__
+    new_fn.__doc__ = fn.__doc__
+    new_fn.__annotations__ = fn.__annotations__
+    new_fn.__kwdefaults__ = None  # already merged into defaults
+    return new_fn
+
+
+@dataclass
+class FunctionPack(Generic[P1, R1]):
+    fn: Callable[P1, R1]
+    infer_meta: Callable[..., Any]
+
+    def id(self) -> int:
+        return id(self.fn)
+
+
+class ConstantParams:
+    def __init__(self, params: dict[str, Any]):
+        self.params = params
+
+    def __hash__(self):
+        return custom_hash(self.params)
+
+    def __eq__(self, other):
+        if not isinstance(other, ConstantParams):
+            return False
+        return self.params == other.params
+
+
+@dataclass
+class OriginalFunctionPack(FunctionPack[P1, R1]):
+    def __post_init__(self):
+        self._specialized_fns: dict[ConstantParams, FunctionPack[P1, R1]] = {}
+
+    @cached_property
+    def fn_eliminated(self) -> Callable[P1, R1]:
+        return eliminate_positional_or_keyword_only(self.fn)
+
+    @cached_property
+    def infer_meta_eliminated(self) -> Callable[..., Any]:
+        return eliminate_positional_or_keyword_only(self.infer_meta)
+
+    def get_bound_args(self, /, *args: P1.args, **kwargs: P1.kwargs):
+        sig = inspect.signature(self.fn)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        return bound_args.arguments
+
+    def separate_mutable_and_const_params(
+        self, /, *args: P1.args, **kwargs: P1.kwargs
+    ) -> tuple[dict[str, paddle.pir.Value], dict[str, Any]]:
+        params = self.get_bound_args(*args, **kwargs)
+
+        mutable_params = {}
+        const_params = {}
+
+        # TODO: Support container types like list, dict, tuple
+        for k, v in params.items():
+            if isinstance(v, paddle.pir.Value):
+                mutable_params[k] = v
+            else:
+                const_params[k] = v
+
+        return mutable_params, const_params
+
+    def specialize(self, const_params: dict[str, Any]) -> FunctionPack[P1, R1]:
+        const_params_wrapper = ConstantParams(const_params)
+        if const_params_wrapper in self._specialized_fns:
+            return self._specialized_fns[const_params_wrapper]
+
+        specialized_fn = partial(self.fn_eliminated, **const_params)
+        specialized_infer_meta = partial(
+            self.infer_meta_eliminated, **const_params
+        )
+        specialized_fn_pack = FunctionPack(
+            specialized_fn, specialized_infer_meta
+        )
+        self._specialized_fns[const_params_wrapper] = specialized_fn_pack
+        return specialized_fn_pack
+
+
+class FunctionRegistry:
+    def __init__(self):
+        self._registry: dict[str, OriginalFunctionPack[Any, Any]] = {}
+
+    def register(
+        self,
+        name: str,
+        fn: Callable[P1, R1],
+        infer_meta: Callable[..., Any],
+    ):
+        if name not in self._registry:
+            self._registry[name] = OriginalFunctionPack(fn, infer_meta)
+            return self._registry[name]
+        fn_pack = self._registry[name]
+        if fn is not fn_pack.fn or infer_meta is not fn_pack.infer_meta:
+            raise ValueError(
+                f"Function '{name}' is already registered with a different implementation."
+            )
+        return fn_pack
+
+    def get(self, name: str) -> OriginalFunctionPack[Any, Any]:
+        if name not in self._registry:
+            raise KeyError(f"Function '{name}' is not registered.")
+        return self._registry[name]
+
+
+FUNCTION_REGISTRY = FunctionRegistry()
+
+
+def bind_constants(fn, infer_meta, *args, **kwargs):
+    sig = inspect.signature(fn)
+    bound_args = sig.bind(*args, **kwargs)
+    bound_args.apply_defaults()
+    params = bound_args.arguments
+
+    mutable_params = {}
+    const_params = {}
+
+    for k, v in params.items():
+        if isinstance(v, paddle.pir.Value):
+            mutable_params[k] = v
+        else:
+            const_params[k] = v
+
+    mutable_arg_names = list(mutable_params.keys())
+    fn = eliminate_positional_or_keyword_only(fn)
+    infer_meta = eliminate_positional_or_keyword_only(infer_meta)
+    return (
+        mutable_arg_names,
+        partial(fn, **const_params),
+        partial(infer_meta, **const_params),
+        list(mutable_params.values()),
+        const_params,
+    )
+
+
+def run_in_dynamic_mode(fn):
+    def dynamic_mode_fn(*args, **kwargs):
+        with paddle.base.dygraph.base.guard():
+            return fn(*args, **kwargs)
+
+    return dynamic_mode_fn
+
+
+def custom_hash(obj):
+    # Compute a hash for various types of objects, including unhashable ones.
+    # This may not be collision-free. For example, hash(-1) is same as hash(-2).
+    # We use dict to resolve collisions in ConstantParams.
+
+    # Handle basic types
+    if isinstance(obj, (int, float, str, bool, bytes)):
+        return hash(obj)
+
+    # Handle sequences (like list, tuple, set, frozenset)
+    if isinstance(obj, (Sequence, frozenset, set)):
+        type_id_map = {list: 1, tuple: 2, frozenset: 3, set: 4}
+        type_id = type_id_map.get(type(obj), 0)
+        return hash((type_id, *tuple(custom_hash(item) for item in obj)))
+
+    # Handle mappings (like dict)
+    if isinstance(obj, Mapping):
+        type_id = 5
+        items_hashed = tuple(
+            sorted((custom_hash(k), custom_hash(v)) for k, v in obj.items())
+        )
+        return hash((type_id, *items_hashed))
+
+    # Fallback: try to use the built-in hash, or use id() if unhashable
+    try:
+        return hash(obj)
+    except TypeError:
+        return id(obj)
+
+
+@overload
+def register_op(
+    fn: Callable[P1, R1],
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+) -> Callable[P1, R1]: ...
+
+
+@overload
+def register_op(
+    fn: None = None,
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+) -> Callable[[Callable[P1, R1]], Callable[P1, R1]]: ...
+
+
+def register_op(
+    fn: Callable[P1, R1] | None = None,
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+):
+    if input_names is None:
+        raise ValueError("Currently, input_names must be provided.")
+    if output_names is None:
+        raise ValueError("Currently, output_names must be provided.")
+    if infer_meta is None:
+        raise ValueError("Currently, infer_meta must be provided.")
+
+    def _register_op(
+        real_fn: Callable[P1, R1],
+    ) -> Callable[P1, R1]:
+        op_name = name or real_fn.__name__
+
+        @paddle.jit.marker.unified
+        @wraps(real_fn)
+        def wrapped_fn(*args: P1.args, **kwargs: P1.kwargs) -> R1:
+            if paddle.in_dynamic_mode():
+                return real_fn(*args, **kwargs)
+
+            fn_pack = FUNCTION_REGISTRY.register(op_name, real_fn, infer_meta)
+            mutable_params, const_params = (
+                fn_pack.separate_mutable_and_const_params(*args, **kwargs)
+            )
+            specialized_fn_pack = fn_pack.specialize(const_params)
+
+            assert len(mutable_params) == len(input_names), (
+                f"Number of mutable arguments ({len(mutable_params)}) does not match "
+                f"the number of input names ({len(input_names)})."
+            )
+
+            out = _C_ops._run_python_op(
+                *mutable_params.values(),
+                name=f"{op_name}_{specialized_fn_pack.id()}",
+                input_names=input_names,
+                output_names=output_names,
+                attrs={
+                    "infer_meta_fn_ptr": specialized_fn_pack.infer_meta,
+                    "fn_ptr": run_in_dynamic_mode(specialized_fn_pack.fn),
+                },
+                inplace_map=inplace_map or {},
+            )
+
+            return out[0] if len(output_names) == 1 else out
+
+        return wrapped_fn
+
+    # Handle @register_op(...)
+    if fn is None:
+        return _register_op
+    # Handle @register_op
+    return _register_op(fn)

python/paddle/static/quantization/__init__.py

@@ -0,0 +1,45 @@
+# 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.
+
+from .post_training_quantization import (  # noqa: F401
+    PostTrainingQuantization,
+    PostTrainingQuantizationProgram,
+    WeightQuantization,
+)
+from .quant2_int8_onednn_pass import (  # noqa: F401
+    Quant2Int8MkldnnPass,
+    Quant2Int8OnednnPass,
+)
+from .quant_int8_onednn_pass import (  # noqa: F401
+    QuantInt8MkldnnPass,
+    QuantInt8OnednnPass,
+)
+from .quanter import (  # noqa: F401
+    convert,
+    quant_aware,
+)
+from .quantization_pass import (  # noqa: F401
+    AddQuantDequantForInferencePass,
+    AddQuantDequantPass,
+    AddQuantDequantPassV2,
+    ConvertToInt8Pass,
+    OutScaleForInferencePass,
+    OutScaleForTrainingPass,
+    QuantizationFreezePass,
+    QuantizationTransformPass,
+    QuantizationTransformPassV2,
+    QuantWeightPass,
+    ReplaceFakeQuantDequantPass,
+    TransformForMobilePass,
+)

python/paddle/static/quantization/adaround.py

@@ -0,0 +1,377 @@
+#   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 logging
+import sys
+import time
+
+import numpy as np
+
+import paddle
+from paddle import static
+
+from ..log_helper import get_logger
+from .utils import (
+    _channelwise_quant_axis1_ops,
+    bias_correction_w,
+    calculate_quant_cos_error,
+    dequant_tensor,
+    load_variable_data,
+    quant_tensor,
+    set_variable_data,
+    stable_sigmoid,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+GAMMA = -0.1
+ZETA = 1.1
+
+
+def compute_soft_rounding(alpha_v):
+    return paddle.clip(
+        paddle.nn.functional.sigmoid(alpha_v) * (ZETA - GAMMA) + GAMMA,
+        min=0,
+        max=1,
+    )
+
+
+def compute_soft_rounding_np(alpha_v):
+    return np.clip(
+        stable_sigmoid(alpha_v) * (ZETA - GAMMA) + GAMMA, a_min=0, a_max=1
+    )
+
+
+class AdaRoundLoss:
+    def __init__(self, reg_param=0.01, default_beta_range=(20, 2)):
+        self.default_reg_param = reg_param
+        self.default_beta_range = default_beta_range
+
+    def compute_recon_loss(self, ada_quantized_output, orig_output):
+        square_cost = paddle.nn.functional.square_error_cost(
+            ada_quantized_output, orig_output
+        )
+        recon_loss = paddle.mean(paddle.sum(square_cost, axis=-1))
+        return recon_loss
+
+    def compute_round_loss(self, alpha_v, warm_start, beta):
+        def round_loss_fn():
+            # compute rectified sigmoid of parameter 'alpha' which maps it between zero and one
+            h_v = compute_soft_rounding(alpha_v)
+
+            # calculate regularization term - which ensures parameter to converge to exactly zeros and ones
+            # at the end of optimization
+            reg_term = paddle.sum(
+                -paddle.pow(paddle.abs(2 * h_v - 1), beta) + 1
+            )
+
+            # calculate the rounding loss
+            round_loss = self.default_reg_param * reg_term
+
+            return round_loss
+
+        round_loss = static.nn.cond(
+            warm_start,
+            lambda: paddle.full(shape=[1], dtype='float32', fill_value=0.0),
+            round_loss_fn,
+        )
+
+        return round_loss
+
+    def compute_beta(self, max_iter, cur_iter, warm_start):
+        #  Start and stop beta for annealing of rounding loss (start_beta, end_beta)
+        start_beta, end_beta = self.default_beta_range
+
+        # iteration at end of warm start period, which is 20% of max iterations
+        warm_start_end_iter = warm_start * max_iter
+
+        # compute relative iteration of current iteration
+        rel_iter = (cur_iter - warm_start_end_iter) / (
+            max_iter - warm_start_end_iter
+        )
+        beta = end_beta + 0.5 * (start_beta - end_beta) * (
+            1 + np.cos(rel_iter * np.pi)
+        )
+
+        return beta
+
+
+class AdaRound:
+    def __init__(
+        self,
+        scale,
+        weight_tensor,
+        scope=None,
+        weight_var_name=None,
+        weight_op_type=None,
+        is_train=True,
+        num_iterations=1000,
+    ):
+        self.is_train = is_train
+        self.num_iterations = num_iterations
+        self.warm_start = 0.1
+        self.weight_bits = 8
+        self.offset = 0.0  # zero-point offset
+        self.adaround_loss = AdaRoundLoss()
+        self.ori_weight_tensor = weight_tensor
+        self.scale = scale
+        self.scope = scope
+        self.quant_axis = 0
+        if weight_op_type in _channelwise_quant_axis1_ops:
+            self.quant_axis = 1
+        self.weight_var_name = weight_var_name
+        self.alpha_name = weight_var_name + ".alpha"
+        self.initialize_alpha(weight_tensor.copy(), scale, weight_var_name)
+
+    def initialize_alpha(self, tensor, scale, var_name):
+        """
+        Initializes alpha parameter, same shape as the weight tensor
+        """
+        tensor_scale = quant_tensor(tensor, scale, quant_axis=self.quant_axis)
+        tensor_floor = np.floor(tensor_scale)
+        tensor = tensor_scale - tensor_floor
+        alpha = -np.log((ZETA - GAMMA) / (tensor - GAMMA) - 1)
+        self.alpha_v = paddle.create_parameter(
+            shape=alpha.shape,
+            dtype="float32",
+            name=var_name + ".alpha",
+            default_initializer=paddle.nn.initializer.Assign(alpha),
+        )
+
+    def _calculate_output_with_adarounded_weights(
+        self, program, place, exe, data, fp32_fetch_list, weight_tensor_dequant
+    ):
+        set_variable_data(
+            self.scope, place, self.weight_var_name, weight_tensor_dequant
+        )
+
+        adaround_out_tensor = exe.run(
+            program=program,
+            feed=data,
+            fetch_list=[fp32_fetch_list],
+            return_numpy=True,
+            scope=self.scope,
+        )
+        return adaround_out_tensor
+
+    def _calculate_quant_weight(self):
+        np_alpha = load_variable_data(self.scope, self.alpha_name)
+        h_alpha = compute_soft_rounding_np(np_alpha)
+
+        # Scale the tensor
+        tensor_scale = quant_tensor(
+            self.ori_weight_tensor.copy(),
+            self.scale,
+            quant_axis=self.quant_axis,
+        )
+
+        weight_tensor = np.floor(tensor_scale)
+
+        # Adaround the tensor
+        weight_tensor_quant = np.add(weight_tensor, h_alpha)
+        return weight_tensor_quant
+
+    def _calculate_adarounded_weights(self):
+        weight_tensor_quant = self._calculate_quant_weight()
+
+        # Dequantize the tensor
+        weight_tensor_dequant = dequant_tensor(
+            weight_tensor_quant + self.offset,
+            self.scale,
+            quant_axis=self.quant_axis,
+        )
+        return weight_tensor_dequant
+
+    def update_final_weights(self):
+        weight_tensor_quant = self._calculate_quant_weight()
+        return weight_tensor_quant
+
+    def get_loss(self, beta, warm_start, adaround_out_tensor, orig_out_tensor):
+        round_loss = self.adaround_loss.compute_round_loss(
+            self.alpha_v, warm_start, beta
+        )
+        recon_loss = self.adaround_loss.compute_recon_loss(
+            adaround_out_tensor, orig_out_tensor
+        )
+        loss = round_loss + recon_loss
+        losses = {
+            'loss': loss,
+            'round_loss': round_loss,
+            'recon_loss': recon_loss,
+        }
+        return losses
+
+    def update_beta_warm(self, cur_iteration):
+        warm_start = cur_iteration < self.num_iterations * self.warm_start
+        beta = self.adaround_loss.compute_beta(
+            self.num_iterations, cur_iteration, self.warm_start
+        )
+        return beta, warm_start
+
+
+def run_adaround(
+    data_loader,
+    fp32_program,
+    fetch_list,
+    exe,
+    scope,
+    place,
+    quantized_op_pairs,
+    weight_op_pairs,
+    scale_dict,
+    num_iterations=1000,
+    lr=0.001,
+    bias_correction=False,
+    fast_mode=True,
+):
+    fetch_op_name = fetch_list[0].name
+    final_weight_tensor_quant_dict = {}
+    for weight_var_name, quant_op_out_name in quantized_op_pairs.items():
+        _logger.info(f'Start adaround op: {weight_var_name}')
+        weight_op_type = weight_op_pairs[weight_var_name]
+        # get scale and weight tensor
+        weight_var_tensor = load_variable_data(scope, weight_var_name)
+        scale = scale_dict[weight_var_name]
+        fp32_fetch_list = None
+        for _op in fp32_program.global_block().ops:
+            if _op.type == "fetch":
+                _op._rename_input(fetch_op_name, quant_op_out_name)
+                fp32_fetch_list = fp32_program.global_block().var(
+                    quant_op_out_name
+                )
+                fetch_op_name = quant_op_out_name
+
+        # build adaround program
+        startup_program = static.Program()
+        train_program = static.Program()
+        with (
+            static.program_guard(train_program, startup_program),
+            paddle.utils.unique_name.guard(),
+        ):
+            # initialize adaround
+            adaround = AdaRound(
+                scale,
+                weight_var_tensor,
+                scope=scope,
+                weight_var_name=weight_var_name,
+                weight_op_type=weight_op_type,
+                num_iterations=num_iterations,
+            )
+            orig_out_tensor = static.data(
+                name='orig_out_tensor',
+                shape=(-1, *fp32_fetch_list.shape),
+                dtype='float32',
+            )
+            adaround_out_tensor = static.data(
+                name='adaround_out_tensor',
+                shape=(-1, *fp32_fetch_list.shape),
+                dtype='float32',
+            )
+            beta_tensor = static.data(
+                name='beta', shape=[-1, 1], dtype='float32'
+            )
+            warm_start_tensor = static.data(
+                name='warm_start', shape=[-1, 1], dtype='bool'
+            )
+
+            train_fetches_loss = adaround.get_loss(
+                beta_tensor,
+                warm_start_tensor,
+                adaround_out_tensor,
+                orig_out_tensor,
+            )
+            optimizer = paddle.optimizer.Adam(learning_rate=lr)
+            loss = train_fetches_loss['loss']
+            optimizer.minimize(loss)
+        exe.run(startup_program)
+
+        start_time = time.time()
+        prev_start_time = start_time
+        for i, data in enumerate(data_loader()):
+            prev_start_time = start_time
+            start_time = time.time()
+            # run fp32 model
+            np_orig_out_tensor = exe.run(
+                program=fp32_program,
+                feed=data,
+                fetch_list=[fp32_fetch_list],
+                return_numpy=True,
+                scope=scope,
+            )
+
+            adaround_weight_tensor_dequant = (
+                adaround._calculate_adarounded_weights()
+            )
+            np_adaround_out_tensor = (
+                adaround._calculate_output_with_adarounded_weights(
+                    fp32_program,
+                    place,
+                    exe,
+                    data,
+                    fp32_fetch_list,
+                    adaround_weight_tensor_dequant,
+                )
+            )
+
+            # If the cosine distance of the two tensor is small, skip training
+            cos_error = calculate_quant_cos_error(
+                np_orig_out_tensor[0], np_adaround_out_tensor[0]
+            )
+            if fast_mode and cos_error > 0.99:
+                _logger.info("The cosine error is small, skip training.")
+                break
+            beta, warm_start = adaround.update_beta_warm(i)
+            feed_dict = {
+                'orig_out_tensor': np_orig_out_tensor[0],
+                'adaround_out_tensor': np_adaround_out_tensor[0],
+                'beta': beta,
+                'warm_start': warm_start,
+            }
+            out = exe.run(
+                train_program,
+                feed=feed_dict,
+                fetch_list=[v.name for v in train_fetches_loss.values()],
+                return_numpy=True,
+            )
+            _logger.info(
+                f"Iter {i:d}, lr {lr:.5f}, loss {np.mean(out[0]):.5f}, loss_round {np.mean(out[1]):.5f}, loss_recon {np.mean(out[2]):.5f}, time {start_time - prev_start_time:.5f}s"
+            )
+            sys.stdout.flush()
+            if i == num_iterations:
+                break
+        final_weight_tensor_quant_dict[weight_var_name] = (
+            adaround.update_final_weights()
+        )
+
+        if bias_correction:
+            final_weight_tensor_quant_dict[weight_var_name] = bias_correction_w(
+                weight_var_tensor,
+                final_weight_tensor_quant_dict[weight_var_name],
+                scale,
+                adaround.quant_axis,
+                weight_bits=adaround.weight_bits,
+            )
+
+        del adaround
+
+    # update adarounded calibrated weights
+    for weight_var_name in quantized_op_pairs.keys():
+        set_variable_data(
+            scope,
+            place,
+            weight_var_name,
+            final_weight_tensor_quant_dict[weight_var_name],
+        )

python/paddle/static/quantization/cal_kl_threshold.py

@@ -0,0 +1,144 @@
+#   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 logging
+import math
+
+import numpy as np
+
+from ..log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+def expand_quantized_bins(quantized_bins, reference_bins):
+    '''
+    Expand hist bins.
+    '''
+    expanded_quantized_bins = [0] * len(reference_bins)
+    num_merged_bins = int(len(reference_bins) / len(quantized_bins))
+    j_start = 0
+    j_end = num_merged_bins
+    for idx in range(len(quantized_bins)):
+        zero_count = reference_bins[j_start:j_end].count(0)
+        num_merged_bins = j_end - j_start
+        if zero_count == num_merged_bins:
+            avg_bin_ele = 0
+        else:
+            avg_bin_ele = quantized_bins[idx] / (
+                num_merged_bins - zero_count + 0.0
+            )
+        for idx1 in range(j_start, j_end):
+            expanded_quantized_bins[idx1] = (
+                0 if reference_bins[idx1] == 0 else avg_bin_ele
+            )
+        j_start += num_merged_bins
+        j_end += num_merged_bins
+        if (idx + 1) == len(quantized_bins) - 1:
+            j_end = len(reference_bins)
+    return expanded_quantized_bins
+
+
+def safe_entropy(reference_distr_P, P_sum, candidate_distr_Q, Q_sum):
+    '''
+    Calculate the entropy.
+    '''
+    assert len(reference_distr_P) == len(candidate_distr_Q)
+    tmp_sum1 = 0
+    tmp_sum2 = 0
+    for idx in range(len(reference_distr_P)):
+        p_idx = reference_distr_P[idx]
+        q_idx = candidate_distr_Q[idx]
+        if p_idx == 0:
+            tmp_sum1 += 0
+            tmp_sum2 += 0
+        else:
+            if q_idx == 0:
+                _logger.error(
+                    "Fatal error!, idx = "
+                    + str(idx)
+                    + " qindex = 0! p_idx = "
+                    + str(p_idx)
+                )
+            tmp_sum1 += p_idx * (math.log(Q_sum * p_idx))
+            tmp_sum2 += p_idx * (math.log(P_sum * q_idx))
+    return (tmp_sum1 - tmp_sum2) / P_sum
+
+
+def cal_kl_threshold(hist, bin_width, bits):
+    '''
+    Using the KL-divergence method to get the more precise threshold.
+
+    Args:
+        hist(List): The hist of the tensor.
+        bin_width(float): The bin width for the hist.
+        bits(int): The quantization bits.
+    '''
+    assert hist.ndim == 1
+    hist_bins = hist.shape[0]
+    starting_iter = int((hist_bins - 1) * 0.5)
+    quant_range = 2 ** (bits - 1) - 1
+
+    P_sum = np.sum(np.array(hist).ravel())
+    min_kl_divergence = 0
+    min_kl_index = 0
+    kl_inited = False
+
+    for i in range(starting_iter, hist_bins):
+        reference_distr_P = hist[0:i].tolist()
+        outliers_count = sum(hist[i:])
+        if reference_distr_P[i - 1] == 0:
+            continue
+        reference_distr_P[i - 1] += outliers_count
+        reference_distr_bins = reference_distr_P[:]
+        candidate_distr_Q = hist[0:i].tolist()
+        num_merged_bins = int(i / quant_range)
+        candidate_distr_Q_quantized = [0] * quant_range
+        j_start = 0
+        j_end = num_merged_bins
+        for idx in range(quant_range):
+            candidate_distr_Q_quantized[idx] = sum(
+                candidate_distr_Q[j_start:j_end]
+            )
+            j_start += num_merged_bins
+            j_end += num_merged_bins
+            if (idx + 1) == quant_range - 1:
+                j_end = i
+        candidate_distr_Q = expand_quantized_bins(
+            candidate_distr_Q_quantized, reference_distr_bins
+        )
+        Q_sum = sum(candidate_distr_Q)
+        kl_divergence = safe_entropy(
+            reference_distr_P, P_sum, candidate_distr_Q, Q_sum
+        )
+        if not kl_inited:
+            min_kl_divergence = kl_divergence
+            min_kl_index = i
+            kl_inited = True
+        elif kl_divergence < min_kl_divergence:
+            min_kl_divergence = kl_divergence
+            min_kl_index = i
+        else:
+            pass
+    if min_kl_index == 0:
+        while starting_iter > 0:
+            if hist[starting_iter] == 0:
+                starting_iter -= 1
+                continue
+            else:
+                break
+        min_kl_index = starting_iter
+    return (min_kl_index + 0.5) * bin_width

python/paddle/static/quantization/quant2_int8_onednn_pass.py

@@ -0,0 +1,736 @@
+#   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 numpy as np
+
+from paddle.utils import deprecated
+
+from ...base.framework import IrGraph
+from ...framework import _get_paddle_place, core
+
+OpRole = core.op_proto_and_checker_maker.OpRole
+
+
+class Quant2Int8OnednnPass:
+    """
+    Transform a quant model IrGraph into MKL-DNN supported INT8 IrGraph.
+    The pass consists of the following transformations:
+        1. gather scale values from fake quantize/dequantize operators,
+        2. extract FP32 inference model graph from the quant graph, i.e.
+            a.  remove fake quantize/dequantize operators,
+            b.  dequantize conv2d and mul's weights,
+        3. optimize the FP32 graph using standard FP32 optimization fuses
+            (e.g. `conv2d`+`bn` -> `conv2d`),
+        4. quantize the optimized FP32 graph using standard INT8v2 quantization
+            passes (`cpu_quantize_pass`, `cpu_quantize_squash_pass`).
+    """
+
+    def __init__(
+        self,
+        _ops_to_quantize,
+        _op_ids_to_skip=None,
+        _scope=None,
+        _place=None,
+        _core=None,
+        _debug=False,
+    ):
+        self._scope = _scope
+        self._place = _get_paddle_place(_place)
+        self._core = _core
+        self._debug = _debug
+        self._fake_quantize_types = [
+            'fake_quantize_moving_average_abs_max',
+            'fake_quantize_range_abs_max',
+        ]
+        self._fake_dequantize_types = [
+            'fake_dequantize_max_abs',
+            'fake_channel_wise_dequantize_max_abs',
+        ]
+        self._fake_quantize_dequantize_types = [
+            'fake_quantize_dequantize_abs_max',
+            'fake_quantize_dequantize_moving_average_abs_max',
+            'fake_channel_wise_quantize_dequantize_abs_max',
+        ]
+        self._ops_to_quantize = _ops_to_quantize
+        self._op_ids_to_skip = (
+            _op_ids_to_skip if _op_ids_to_skip is not None else {-1}
+        )
+        self._scale_immutable_ops = [
+            'transpose2',
+            'reshape2',
+            'pool2d',
+            'slice',
+            'shape',
+            'nearest_interp',
+            'nearest_interp_v2',
+            'split',
+        ]
+        self._scale_ops = ['scale']
+        self._conv_ops = ['conv2d', 'depthwise_conv2d']
+        self._pool_ops = ['pool2d']
+        self._mul_ops = ['mul']
+        self._fc_ops = ['fc']
+        self._relu_ops = ['relu', 'relu6']
+        self._matmul_ops = ['matmul', 'matmul_v2']
+        self._gru_ops = ['fusion_gru', 'multi_gru']
+        self._lstm_ops = ['fusion_lstm']
+        self._weight_thresholds = {}
+        # Collect the Input and Output scales from Fake quant models
+        self._var_quant_scales = {}
+        self._max_range = {}
+        self._s8_max = 127
+        self._pass_idx = 0
+        self._pass_group = 'int8'
+
+    def apply(self, graph):
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+
+        self._reset_pass_idx_and_group('int8')
+        graph = self._label_skip_quantized_op(graph)
+        graph = self._gather_weight_thresholds_from_fake(graph)
+        graph = self._gather_input_scales_from_fake(graph)
+        graph = self._gather_output_scales_from_attr(graph)
+        graph = self._remove_fake_ops(graph)
+        graph = self._dequantize_weights(graph)
+        graph = self._optimize_fp32_graph(graph)
+        graph = self._compute_weight_scales(graph)
+        # This function causes nondeterministic quantization behavior
+        # graph = self._update_relu_output_scales(graph)
+        graph = self._propagate_scales(graph)
+        graph = self._quantize_fp32_graph(graph)
+        graph = self._cleanup(graph)
+        return graph
+
+    def prepare_and_optimize_fp32(self, graph):
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+
+        self._reset_pass_idx_and_group('fp32')
+        graph = self._optimize_fp32_graph(graph)
+        graph = self._cleanup(graph)
+        return graph
+
+    def _reset_pass_idx_and_group(self, group):
+        self._pass_idx = 0
+        self._pass_group = group
+
+    def _convert_scale2tensor(self, scale):
+        tensor = core.DenseTensor()
+        tensor.set(scale, core.CPUPlace())
+        return tensor
+
+    def _is_quantizing_all_ops(self):
+        return len(self._ops_to_quantize) == 0
+
+    def _is_any_of_op_types_in_graph(self, op_types, graph):
+        return any(op.name() in op_types for op in graph.all_op_nodes())
+
+    def _is_any_of_op_types_quantized(self, op_types, graph):
+        return self._is_any_of_op_types_in_graph(op_types, graph) and (
+            self._is_quantizing_all_ops()
+            or any(op_type in self._ops_to_quantize for op_type in op_types)
+        )
+
+    def _is_conv_quantized(self, graph):
+        return self._is_any_of_op_types_quantized(self._conv_ops, graph)
+
+    def _is_fc_quantized(self, graph):
+        return self._is_any_of_op_types_quantized(self._fc_ops, graph)
+
+    def _label_skip_quantized_op(self, graph):
+        """
+        For some ops(conv2d, depthwise_conv2d, mul, matmul), find and label
+        the skip quantized ops. cpu_quantize_placement_pass will use the
+        label to identify it.
+        For static models, the skip quantized ops have `skip_quant` attr.
+        Therefore, it only needs to find and label the skip quantized ops for
+        dygraph models, in which the quantized ops don't have `quantization_type`
+        attr.
+        """
+        target_ops = self._conv_ops + self._mul_ops + self._matmul_ops
+        for op_node in graph.all_op_nodes():
+            if op_node.name() in target_ops and not op_node.op().has_attr(
+                "quantization_type"
+            ):
+                is_quantized_op = True
+                for var_node in op_node.inputs:
+                    for front_op_node in var_node.inputs:
+                        if "quantize" not in front_op_node.name():
+                            is_quantized_op = False
+                if not is_quantized_op:
+                    op_node.op()._set_attr("skip_quant", True)
+        return graph
+
+    def _add_scale_for_vars(self, var_names, use_unsigned_int, lod_tensor):
+        """
+        Save quantization scales for variables. Do not overwrite.
+        """
+        scales = self._var_quant_scales
+        for var_name in var_names:
+            if var_name not in scales:
+                scales[var_name] = (use_unsigned_int, lod_tensor)
+
+    def _gather_input_scales_from_fake(self, graph):
+        # fake_quantize_dequantize_abs_max doesn't have scale value
+        fake_ops = ['fake_quantize_dequantize_moving_average_abs_max']
+        fake_ops.extend(self._fake_quantize_types)
+
+        for op in graph.all_op_nodes():
+            if op.name() in fake_ops:
+                bit_length = op.op().attr("bit_length")
+                assert bit_length == 8, (
+                    f'Unsupported number quantization bits ({bit_length}). Only 8 is supported now.'
+                )
+
+                input_name = op.input("X")[0]
+                scale_name = op.input("InScale")[0]
+                output_name = op.output("Out")[0]
+                # Gather new weight scales after folding batchnorm in convolution
+                scale = np.array(
+                    1.0 / self._load_param(self._scope, scale_name)[0]
+                ).astype(np.float64)
+                scale[scale == np.inf] = 0.0
+                lod_tensor = self._convert_scale2tensor(scale)
+                use_unsigned_int = False
+                self._add_scale_for_vars(
+                    [input_name, output_name], use_unsigned_int, lod_tensor
+                )
+
+        return graph
+
+    def _gather_weight_thresholds_from_fake(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._fake_dequantize_types:
+                input_name = op.input("X")[0]
+                if op.op().has_attr("max_range"):
+                    _max_range = np.array(op.op().attr("max_range")).astype(
+                        np.float64
+                    )
+                    self._weight_thresholds[input_name] = np.array(
+                        self._s8_max * self._s8_max / _max_range
+                    ).astype(np.float64)
+                else:
+                    scale_name = op.input("Scales")[0]
+                    self._weight_thresholds[input_name] = np.array(
+                        self._load_param(self._scope, scale_name)
+                    ).astype(np.float64)
+
+        return graph
+
+    def _gather_output_scales_from_attr(self, graph):
+        for op in graph.all_op_nodes():
+            if op.op().has_attr("out_threshold"):
+                attr_scale = op.op().attr("out_threshold")
+                if attr_scale == 0.0:
+                    continue
+                scale = np.array(1.0 / attr_scale).astype(np.float64)
+                scale[scale == np.inf] = 0.0
+                scale_lod_tensor = self._convert_scale2tensor(scale)
+                use_unsigned_int = False
+                for output_name in op.op().outputs():
+                    for out_var_name in op.op().output(output_name):
+                        self._add_scale_for_vars(
+                            [out_var_name], use_unsigned_int, scale_lod_tensor
+                        )
+
+        return graph
+
+    def _propagate_scales(self, graph):
+        def _update_scale_op_in_scale(op, input, output):
+            unsigned, tensor = self._var_quant_scales[output]
+            scale = np.array(tensor) * op.op().attr("scale")
+            new_tensor = self._convert_scale2tensor(scale.astype(np.float64))
+            self._var_quant_scales[input] = (unsigned, new_tensor)
+
+        def _update_scales(graph):
+            waiting_for_scale = set()
+            for op in graph.all_op_nodes():
+                if op.name() in self._scale_immutable_ops:
+                    if op.name() == 'slice' or op.name() == 'shape':
+                        input_name = op.input("Input")[0]
+                    else:
+                        input_name = op.input("X")[0]
+                    output_name = op.output("Out")[0]
+                    tensor_names = [input_name, output_name]
+
+                    if all(
+                        name not in self._var_quant_scales
+                        for name in tensor_names
+                    ):
+                        waiting_for_scale.update(tensor_names)
+                        continue
+                    elif input_name in self._var_quant_scales:
+                        self._var_quant_scales[output_name] = (
+                            self._var_quant_scales[input_name]
+                        )
+                    elif output_name in self._var_quant_scales:
+                        self._var_quant_scales[input_name] = (
+                            self._var_quant_scales[output_name]
+                        )
+
+                elif op.name() == 'concat':
+                    output_name = op.output("Out")[0]
+                    if output_name in self._var_quant_scales:
+                        input_names = op.input("X")
+                        for input_name in input_names:
+                            self._var_quant_scales[input_name] = (
+                                self._var_quant_scales[output_name]
+                            )
+                elif op.name() in self._scale_ops:
+                    input_name = op.input("X")[0]
+                    output_name = op.output("Out")[0]
+                    if output_name in self._var_quant_scales:
+                        _update_scale_op_in_scale(op, input_name, output_name)
+            return waiting_for_scale
+
+        waiting_for_scale = _update_scales(graph)
+        waiting_for_scale_prev = set()
+
+        while (
+            len(waiting_for_scale) != 0
+            and waiting_for_scale != waiting_for_scale_prev
+        ):
+            waiting_for_scale_prev = waiting_for_scale
+            waiting_for_scale = _update_scales(graph)
+
+        return graph
+
+    def _load_param(self, scope, param_name):
+        return np.array(scope.find_var(param_name).get_tensor())
+
+    def _remove_fake_ops(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._fake_quantize_types:
+                self._remove_fake_quantize(graph, op)
+            elif op.name() in self._fake_dequantize_types:
+                self._remove_fake_dequantize(graph, op)
+            elif op.name() in self._fake_quantize_dequantize_types:
+                self._remove_fake_dequantize(graph, op)
+
+        return graph
+
+    def _remove_fake_quantize(self, graph, op):
+        fake_quant_in = graph._find_node_by_name(op.inputs, op.input("X")[0])
+        fake_quant_in_scale = graph._find_node_by_name(
+            op.inputs, op.input("InScale")[0]
+        )
+        fake_quant_out = graph._find_node_by_name(
+            op.outputs, op.output("Out")[0]
+        )
+        fake_quant_out_scale = graph._find_node_by_name(
+            op.outputs, op.output("OutScale")[0]
+        )
+
+        next_ops = fake_quant_out.outputs
+        for next_op in next_ops:
+            self._swap_inputs(next_op, fake_quant_out, fake_quant_in)
+            graph.link_to(fake_quant_in, next_op)
+        graph.safe_remove_nodes(
+            {op, fake_quant_in_scale, fake_quant_out, fake_quant_out_scale}
+        )
+
+        return graph
+
+    def _remove_fake_dequantize(self, graph, op):
+        fake_dequant_in = graph._find_node_by_name(op.inputs, op.input("X")[0])
+        fake_dequant_out = graph._find_node_by_name(
+            op.outputs, op.output("Out")[0]
+        )
+
+        next_ops = fake_dequant_out.outputs
+        for next_op in next_ops:
+            self._swap_inputs(next_op, fake_dequant_out, fake_dequant_in)
+            graph.link_to(fake_dequant_in, next_op)
+        graph.safe_remove_nodes({op, fake_dequant_out})
+
+        return graph
+
+    def _swap_inputs(self, op, old_input, new_input):
+        for input_name in op.op().input_names():
+            if old_input.name() in op.input(input_name):
+                op.op().set_input(
+                    input_name,
+                    [
+                        new_input.name() if x == old_input.name() else x
+                        for x in op.input(input_name)
+                    ],
+                )
+
+    def _dequantize_weights(self, graph):
+        def _is_int8_weights(op_node, weight_name):
+            weight_var_name = op_node.input(weight_name)[0]
+            if self._scope.find_var(weight_var_name) is None:
+                return False
+            weight = self._load_param(self._scope, weight_var_name)
+            return np.all(np.mod(weight, 1) == 0)
+
+        mul_and_matmul_ops = self._mul_ops + self._matmul_ops
+        for op in graph.all_op_nodes():
+            if op.name() in self._conv_ops and _is_int8_weights(op, "Filter"):
+                self._dequantize_op_weights(graph, op, "Filter", "Output")
+            elif op.name() in mul_and_matmul_ops and _is_int8_weights(op, "Y"):
+                self._dequantize_op_weights(graph, op, "Y", "Out")
+
+        return graph
+
+    def _dequantize_op_weights(self, graph, op_node, weight_name, output_name):
+        weight_var_name = op_node.input(weight_name)[0]
+        output_var_name = op_node.output(output_name)[0]
+        # Convert int8 range weights to fp32 range weights
+        scales = self._weight_thresholds[output_var_name]
+        weight = self._load_param(self._scope, weight_var_name)
+        if scales.size == 1 or scales.size == weight.shape[0]:
+            w_fp32 = np.multiply(np.divide(weight, self._s8_max).T, scales.T).T
+        elif len(weight.shape) > 1 and scales.size == weight.shape[1]:
+            w_fp32 = np.multiply(np.divide(weight, self._s8_max), scales)
+        else:
+            raise ValueError(
+                f"The size of weight scales vector ({scales.size}) does not match the dimensions ({weight.shape}) of the weights tensor {weight_var_name}."
+            )
+        w_fp32 = w_fp32.reshape(weight.shape).astype(np.float32)
+        self._restore_var(weight_var_name, w_fp32)
+
+    def _restore_var(self, name, array):
+        tensor = self._scope.find_var(name).get_tensor()
+        tensor.set(array, self._place)
+
+    def _update_activations(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._conv_ops and not op.op().has_attr(
+                "fuse_activation"
+            ):
+                activation = ""
+                if op.op().has_attr("fuse_relu") and op.op().attr("fuse_relu"):
+                    activation = "relu"
+                op.set_attr("fuse_activation", activation)
+        return graph
+
+    def _remove_ctrl_vars(self, graph):
+        remove_ctr_vars = set()
+        for node in graph.all_var_nodes():
+            if node.is_ctrl_var():
+                remove_ctr_vars.add(node)
+        graph.safe_remove_nodes(remove_ctr_vars)
+        return graph
+
+    def _optimize_fp32_graph(self, graph):
+        graph = self._update_activations(graph)
+        graph = self._remove_ctrl_vars(graph)
+        graph = self._apply_pass(
+            graph, 'onednn_placement_pass', ['onednn_enabled_op_types'], [set()]
+        )
+        # remove dropout ops
+        graph = self._apply_pass(graph, 'simplify_with_basic_ops_pass')
+        graph = self._apply_pass(graph, 'layer_norm_fuse_pass')
+        graph = self._apply_pass(graph, 'attention_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'seqconv_eltadd_relu_fuse_pass')
+        graph = self._apply_pass(graph, 'fc_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'mul_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'fc_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'mul_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'multi_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'multi_gru_seq_fuse_pass')
+        graph = self._apply_pass(graph, 'seq_concat_fc_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_squeeze2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_reshape2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_flatten2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'matmul_v2_scale_fuse_pass')
+        graph = self._apply_pass(graph, 'squared_mat_sub_fuse_pass')
+        graph = self._apply_pass(graph, 'is_test_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_v2_to_mul_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_v2_to_matmul_pass')
+        graph = self._apply_pass(graph, 'matmul_scale_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_to_mul_pass')
+        graph = self._apply_pass(graph, 'repeated_fc_relu_fuse_pass')
+        graph = self._apply_pass(graph, 'depthwise_conv_onednn_pass')
+        graph = self._apply_pass(graph, 'conv_bn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_eltwiseadd_bn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_affine_channel_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_transpose_bn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'conv_transpose_eltwiseadd_bn_fuse_pass'
+        )
+        graph = self._apply_pass(graph, 'conv_bias_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_transpose_bias_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_elementwise_add_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_activation_onednn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'fc_fuse_pass', ['use_gpu', 'use_fc_padding'], [False, False]
+        )
+        graph = self._apply_pass(graph, 'repeated_fc_relu_fuse_pass')
+        if self._is_fc_quantized(graph):
+            # Disabled due to topology-dependent speed-up
+            graph = self._apply_pass(graph, 'fc_onednn_pass')
+            graph = self._apply_pass(graph, 'fc_act_onednn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'matmul_transpose_reshape_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(
+            graph, 'matmul_elementwise_add_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(graph, 'matmul_activation_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'batch_norm_act_fuse_pass')
+        graph = self._apply_pass(graph, 'softplus_activation_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'scale_matmul_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'reshape_transpose_matmul_onednn_fuse_pass'
+        )
+        # the following pass should be the last one since it will work on all fused ops.
+        graph = self._apply_pass(graph, 'runtime_context_cache_pass')
+        return graph
+
+    def _apply_pass(self, graph, pass_name, attrs=None, attr_values=None):
+        ir_pass = core.get_pass(pass_name)
+        cpp_graph = graph.graph
+        if not cpp_graph.has('__param_scope__'):
+            cpp_graph.set_not_owned('__param_scope__', self._scope)
+        if attrs:
+            assert attr_values and len(attrs) == len(attr_values), (
+                "Different number of pass attributes and their values."
+            )
+            for attr, value in zip(attrs, attr_values):
+                ir_pass.set(attr, value)
+        ir_pass.apply(cpp_graph)
+        if self._debug:
+            graph.draw(
+                '.',
+                f'{self._pass_group}_{self._pass_idx}_{pass_name}',
+                graph.all_op_nodes(),
+            )
+        self._remove_unused_var_nodes(graph)
+        self._pass_idx += 1
+        return graph
+
+    def _cleanup(self, graph):
+        graph = self._remove_unused_var_nodes(graph)
+        graph = self._set_op_role_forward(graph)
+        return graph
+
+    def _remove_unused_var_nodes(self, graph):
+        all_used_vars = set()
+        ops = graph.all_op_nodes()
+        for op_node in ops:
+            for input_node in op_node.inputs:
+                all_used_vars.add(input_node)
+            for output_node in op_node.outputs:
+                all_used_vars.add(output_node)
+
+        all_used_vars = {n.node for n in all_used_vars}
+        all_unused_vars = set(
+            filter(
+                lambda node: node.node not in all_used_vars,
+                graph.all_var_nodes(),
+            )
+        )
+        graph.safe_remove_nodes(all_unused_vars)
+        return graph
+
+    def _set_op_role_forward(self, graph):
+        ops = graph.all_op_nodes()
+        for op in ops:
+            op.set_attr("op_role", OpRole.Forward)
+        return graph
+
+    def _compute_weight_scales(self, graph):
+        def _compute_var_scales(ops, w_name, axis):
+            for op in graph.all_op_nodes():
+                if op.op().type() in ops:
+                    weight_var_name = op.input(w_name)[0]
+                    weights = np.array(
+                        self._load_param(self._scope, weight_var_name)
+                    )
+                    scales = 1.0 / np.amax(
+                        np.abs(weights.reshape(weights.shape[0], -1)).astype(
+                            np.float64
+                        ),
+                        axis=axis,
+                    )
+                    scales[scales == np.inf] = 0.0
+
+                    lod_tensor = self._convert_scale2tensor(scales)
+                    use_unsigned_int = False
+                    self._var_quant_scales[weight_var_name] = (
+                        use_unsigned_int,
+                        lod_tensor,
+                    )
+
+        def _compute_single_gru_weight_scales(wx_var_name, wh_var_name):
+            wx = np.array(self._load_param(self._scope, wx_var_name))
+            wh = np.array(self._load_param(self._scope, wh_var_name))
+            OC = wh.shape[0]
+            scale_ur = 1.0 / np.max(
+                np.abs(
+                    np.concatenate(
+                        [
+                            wx[:, : 2 * OC],
+                            wh.flatten()[: 2 * OC * OC].reshape(OC, 2 * OC),
+                        ],
+                        axis=0,
+                    )
+                ),
+                axis=0,
+            )
+            scale_o = 1.0 / np.max(
+                np.abs(
+                    np.concatenate(
+                        [
+                            wx[:, 2 * OC :],
+                            wh.flatten()[2 * OC * OC :].reshape(OC, OC),
+                        ],
+                        axis=0,
+                    )
+                ),
+                axis=0,
+            )
+
+            gru_weights_scale = np.concatenate([scale_ur, scale_o]).astype(
+                'float'
+            )
+
+            return self._convert_scale2tensor(gru_weights_scale)
+
+        def _compute_gru_weight_scales(wx_name, wh_name):
+            for op in graph.all_op_nodes():
+                if op.op().type() in self._gru_ops:
+                    assert len(op.input(wx_name)) == len(op.input(wh_name)), (
+                        f'Mismatch in number of weights inputs ({len(op.input(wx_name))} for WeightX vs. {len(op.input(wh_name))} for WeightH).'
+                    )
+                    for i, wx_var_name in enumerate(op.input(wx_name)):
+                        wh_var_name = op.input(wh_name)[i]
+                        use_unsigned_int = False
+                        lod_tensor = _compute_single_gru_weight_scales(
+                            wx_var_name, wh_var_name
+                        )
+                        self._var_quant_scales[wx_var_name] = (
+                            use_unsigned_int,
+                            lod_tensor,
+                        )
+
+        def _compute_single_lstm_weight_scales(wx_var_name, wh_var_name):
+            wx = np.array(self._load_param(self._scope, wx_var_name))
+            wh = np.array(self._load_param(self._scope, wh_var_name))
+
+            lstm_weights_scale = 1.0 / np.max(
+                np.abs(np.concatenate([wx[:, :], wh[:, :]], axis=0)), axis=0
+            )
+            lstm_weights_scale = lstm_weights_scale.astype('float')
+
+            return self._convert_scale2tensor(lstm_weights_scale)
+
+        def _compute_lstm_weight_scales(wx_name, wh_name):
+            for op in graph.all_op_nodes():
+                if op.op().type() in self._lstm_ops:
+                    assert len(op.input(wx_name)) == len(op.input(wh_name)), (
+                        f'Mismatch in number of weights inputs ({len(op.input(wx_name))} for WeightX vs. {len(op.input(wh_name))} for WeightH).'
+                    )
+                    for i, wx_var_name in enumerate(op.input(wx_name)):
+                        wh_var_name = op.input(wh_name)[i]
+                        use_unsigned_int = False
+                        lod_tensor = _compute_single_lstm_weight_scales(
+                            wx_var_name, wh_var_name
+                        )
+                        self._var_quant_scales[wx_var_name] = (
+                            use_unsigned_int,
+                            lod_tensor,
+                        )
+
+        _compute_var_scales(self._conv_ops, "Filter", axis=1)
+        _compute_var_scales(self._fc_ops, "W", axis=0)
+        _compute_var_scales(self._gru_ops, "WeightH", axis=0)
+        _compute_var_scales(self._lstm_ops, "WeightH", axis=0)
+        _compute_gru_weight_scales("WeightX", "WeightH")
+        _compute_lstm_weight_scales("WeightX", "WeightH")
+        return graph
+
+    def _update_relu_output_scales(self, graph):
+        def _set_unsigned_scale(graph, ops, op_out_name, predicate):
+            '''
+            Sets the type of an output scale of a passed op type(s) to 'unsigned int8' if the
+            predicate applied on op passes. Typically, the predicate checks if op's
+            activation is set to relu.
+            '''
+            for op in graph.all_op_nodes():
+                if op.name() in ops:
+                    out_name = op.output(op_out_name)[0]
+                    if out_name in self._var_quant_scales and predicate(
+                        op.op()
+                    ):
+                        is_unsigned, tensor = self._var_quant_scales[out_name]
+                        if is_unsigned is False:
+                            # If the variable is signed, it means that the scales for this var
+                            # were computed for signed data, so the scale must be multiplied by 2
+                            # to fill the entire range of uint8
+                            scale = np.array(tensor) * 2
+                            tensor = self._convert_scale2tensor(
+                                scale.astype(np.float64)
+                            )
+                        self._var_quant_scales[out_name] = (True, tensor)
+            return graph
+
+        def conv_predicate(op):
+            return op.attr("fuse_activation") in self._relu_ops
+
+        graph = _set_unsigned_scale(
+            graph, self._conv_ops, "Output", conv_predicate
+        )
+
+        def fc_predicate(op):
+            return op.attr("activation_type") in self._relu_ops
+
+        graph = _set_unsigned_scale(graph, self._fc_ops, "Out", fc_predicate)
+
+        graph = _set_unsigned_scale(
+            graph, self._relu_ops, 'Out', lambda op: True
+        )
+
+        return graph
+
+    def _get_data_layout(self, graph):
+        return 'NHWC' if self._is_conv_quantized(graph) else 'NCHW'
+
+    def _quantize_fp32_graph(self, graph):
+        graph = self._apply_pass(graph, 'scale_matmul_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'reshape_transpose_matmul_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(
+            graph,
+            'cpu_quantize_placement_pass',
+            ['quantize_enabled_op_types'],
+            [self._ops_to_quantize],
+        )
+        graph = self._apply_pass(
+            graph,
+            'cpu_quantize_pass',
+            ['quant_var_scales', 'data_layout'],
+            [self._var_quant_scales, self._get_data_layout(graph)],
+        )
+        graph = self._apply_pass(graph, 'cpu_quantize_squash_pass')
+        graph = self._apply_pass(graph, 'int8_scale_calculation_onednn_pass')
+        graph = self._apply_pass(graph, 'params_quantization_onednn_pass')
+        return graph
+
+
+class Quant2Int8MkldnnPass(Quant2Int8OnednnPass):
+    @deprecated(
+        since="3.1.0",
+        update_to="paddle.static.quantization.Quant2Int8OnednnPass",
+        level=1,
+        reason="Quant2Int8MkldnnPass will be removed in future",
+    )
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

python/paddle/static/quantization/quant_config.py

@@ -0,0 +1,287 @@
+#   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.
+
+
+# A dict of operators that contain weights and support quantization,
+# including operator names, actual input and output names.
+SUPPORT_WEIGHT_QUANTIZATION_OP_DICT = {
+    "conv2d": [["Input", "Filter"], ["Output"]],
+    "depthwise_conv2d": [["Input", "Filter"], ["Output"]],
+    "conv2d_transpose": [["Input", "Filter"], ["Output"]],
+    "mul": [["X", "Y"], ["Out"]],
+    "matmul": [["X", "Y"], ["Out"]],
+    "matmul_v2": [["X", "Y"], ["Out"]],
+}
+
+# A dict of operators that supports quantization and has only activation inputs,
+# including operator names, actual input and output names.
+SUPPORT_ACT_QUANTIZATION_OP_DICT = {
+    "mul": [["X", "Y"], ["Out"]],
+    "matmul": [["X", "Y"], ["Out"]],
+    "matmul_v2": [["X", "Y"], ["Out"]],
+    "pool2d": [["X"], ["Out"]],
+    "elementwise_add": [["X", "Y"], ["Out"]],
+    "concat": [["X"], ["Out"]],
+    "softmax": [["X"], ["Out"]],
+    "argmax": [["X"], ["Out"]],
+    "transpose": [["X"], ["Out"]],
+    "equal": [["X", "Y"], ["Out"]],
+    "gather": [["X"], ["Out"]],
+    "greater_equal": [["X", "Y"], ["Out"]],
+    "greater_than": [["X", "Y"], ["Out"]],
+    "less_equal": [["X", "Y"], ["Out"]],
+    "less_than": [["X", "Y"], ["Out"]],
+    "mean": [["X"], ["Out"]],
+    "not_equal": [["X", "Y"], ["Out"]],
+    "reshape": [["X"], ["Out"]],
+    "reshape2": [["X"], ["Out"]],
+    "transpose2": [["X"], ["Out"]],
+    "nearest_interp": [["X"], ["Out"]],
+    "trilinear_interp": [["X"], ["Out"]],
+    "slice": [["Input"], ["Out"]],
+    "squeeze": [["X"], ["Out"]],
+    "elementwise_sub": [["X", "Y"], ["Out"]],
+    "relu": [["X"], ["Out"]],
+    "relu6": [["X"], ["Out"]],
+    "leaky_relu": [["X"], ["Out"]],
+    "prelu": [["X", "Alpha"], ["Out"]],
+    "tanh": [["X"], ["Out"]],
+    "swish": [["X"], ["Out"]],
+    "dropout": [["X"], ["Out"]],
+    "batch_norm": [["X"], ["Y"]],
+    "layer_norm": [["X"], ["Y"]],
+    "sigmoid": [["X"], ["Out"]],
+    "elementwise_mul": [["X", "Y"], ["Out"]],
+    "elementwise_pow": [["X", "Y"], ["Out"]],
+    "hard_swish": [["X"], ["Out"]],
+    "hard_sigmoid": [["X"], ["Out"]],
+    "gru": [["Input", "Weight"], ["Hidden"]],
+    "lstm": [["Input", "Weight"], ["Hidden"]],
+    "pad2d": [["X"], ["Out"]],
+    "pad3d": [["X"], ["Out"]],
+    "flatten": [["X"], ["Out"]],
+    "flatten2": [["X"], ["Out"]],
+    "unsqueeze2": [["X"], ["Out"]],
+    "flatten_contiguous_range": [["X"], ["Out"]],
+    "split": [["X"], ["Out"]],
+    "squeeze2": [["X"], ["Out"]],
+    "nearest_interp_v2": [["X"], ["Out"]],
+    "bilinear_interp": [["X"], ["Out"]],
+    "bilinear_interp_v2": [["X"], ["Out"]],
+    "fill_constant_batch_size_like": [["Input"], ["Out"]],
+    "arg_max": [["X"], ["Out"]],
+    "abs": [["X"], ["Out"]],
+    "assign": [["X"], ["Out"]],
+    "cast": [["X"], ["Out"]],
+    "clip": [["X"], ["Out"]],
+    "box_coder": [["PriorBox"], ["OutputBox"]],
+    "crop": [["X"], ["Out"]],
+    "cumsum": [["X"], ["Out"]],
+    "expand_v2": [["X"], ["Out"]],
+    "fill_any_like": [["X"], ["Out"]],
+    "fill_constant": [[], ["Out"]],
+    "gelu": [["X"], ["Out"]],
+    "instance_norm": [["X"], ["Y"]],
+    "lookup_table": [["W", "Ids"], ["Out"]],
+    "lookup_table_v2": [["W", "Ids"], ["Out"]],
+    "norm": [["X"], ["Norm"]],
+    "p_norm": [["X"], ["Out"]],
+    "pow": [["X"], ["Out"]],
+    "reduce_mean": [["X"], ["Out"]],
+    "stack": [["X"], ["Y"]],
+    "top_k_v2": [["X"], ["Out", "Indices"]],
+    "logical_and": [["X", "Y"], ["Out"]],
+    "logical_not": [["X"], ["Out"]],
+    "meshgrid": [["X"], ["Out"]],
+    "roi_align": [["X", "ROIs"], ["Out"]],
+    "strided_slice": [["Input"], ["Out"]],
+    "where": [["Condition", "X", "Y"], ["Out"]],
+    "grid_sampler": [["X", "Grid"], ["Output"]],
+    "tile": [["X"], ["Out"]],
+    "group_norm": [["X"], ["Y", "Mean", "Variance"]],
+    "reduce_sum": [["X"], ["Out"]],
+    "square": [["X"], ["Out"]],
+    "softplus": [["X"], ["Out"]],
+    "shuffle_channel": [["X"], ["Out"]],
+    "reduce_max": [["X"], ["Out"]],
+    "scale": [["X"], ["Out"]],
+}
+
+# A full dict of operators that supports quantization,
+# including operator names, actual input and output names.
+SUPPORT_QUANTIZATION_OP_DICT = SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.copy()
+SUPPORT_QUANTIZATION_OP_DICT.update(SUPPORT_ACT_QUANTIZATION_OP_DICT)
+
+
+class BaseQuantizer:
+    """
+    Basic quantization configuration class, which configures some hyperparameters
+    required for quantization, including the list of op types to be quantized,
+    quantization bit number for weight and activation and the range of quantization values.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -128
+        self._quant_max = 127
+
+    @property
+    def weight_quant_operation_types(self):
+        """
+        Operation type list which should support weight quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        base_weight_op_type_list = list(
+            SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.keys()
+        )
+        if self._quantizable_op_type:
+            weight_list = []
+            for _op_type in self._quantizable_op_type:
+                if _op_type in base_weight_op_type_list:
+                    weight_list.append(_op_type)
+            return weight_list
+        else:
+            return base_weight_op_type_list
+
+    @property
+    def activation_quant_operation_types(self):
+        """
+        Operation type list which should support activation quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        base_act_op_type_list = list(SUPPORT_ACT_QUANTIZATION_OP_DICT.keys())
+        act_quant_op_list = []
+        if self._quantizable_op_type:
+            for _op_type in self._quantizable_op_type:
+                if _op_type in base_act_op_type_list:
+                    act_quant_op_list.append(_op_type)
+        else:
+            act_quant_op_list = [
+                'mul',
+                'matmul',
+                'matmul_v2',
+            ]
+        return act_quant_op_list
+
+    @property
+    def observer_operation_types(self):
+        """
+        Operation type list for observer in quantization. These nodes only count the
+        calibration boundary scale and do not participate in the fake quantization.
+        In order to facilitate the deployment of the prediction engine, quant
+        and dequant nodes will be inserted after these ops when exporting the model.
+        """
+        return list(SUPPORT_ACT_QUANTIZATION_OP_DICT.keys())
+
+
+class TensorRTQuantizer(BaseQuantizer):
+    """
+    TensorRT quantization configuration class.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        super().__init__()
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -128
+        self._quant_max = 127
+
+    @property
+    def activation_quant_operation_types(self):
+        """
+        Operation type list which should support activation quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        return [
+            "pool2d",
+            "elementwise_add",
+            "elementwise_sub",
+            "elementwise_mul",
+            "elementwise_pow",
+            "concat",
+            "softmax",
+            "argmax",
+            "mean",
+            "relu",
+            "relu6",
+            "leaky_relu",
+            "tanh",
+            "swish",
+            "softplus",
+            "gelu",
+            "hard_sigmoid",
+            "hard_swish",
+            "sigmoid",
+            "layer_norm",
+            "matmul_v2",
+            "split",
+            "bilinear_interp",
+            "nearest_interp",
+            "trilinear_interp",
+            "nearest_interp_v2",
+            "bilinear_interp",
+            "bilinear_interp_v2",
+            "clip",
+            "pow",
+            "reduce_mean",
+            "reduce_sum",
+            "reduce_max",
+        ]
+
+
+class ARMCPUQuantizer(BaseQuantizer):
+    """
+    ARM CPU with Paddle Lite quantization configuration class.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        super().__init__()
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -127
+        self._quant_max = 127

python/paddle/static/quantization/quant_int8_onednn_pass.py

@@ -0,0 +1,302 @@
+#   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 numpy as np
+
+from paddle.utils import deprecated
+
+from ...base.framework import IrGraph
+from ...framework import _get_paddle_place
+
+
+class QuantInt8OnednnPass:
+    """
+    Convert QuantizationFreezePass generated IrGraph to MKL-DNN supported INT8
+    IrGraph. Following transformations did in this pass:
+        1. Convert int8 range weights with float32 data type, which are generated by
+           the QuantizationFreezePass, to float32 range weights with float32 data type
+           by using the corresponding scales. This conversion is because MKL-DNN INT8
+           conv2d kernel and mul kernel now only support float32 weights input, hence
+           weights quantization will happen inside the conv2d and mul INT8 kernel.
+        2. Create the new conv2d or mul op with the converted weights and link its output
+           to fake_dequantize_abs_max op's output and set conv2d's attribute "force_fp32
+           _output" as true
+        3. Transform fake_quantize_xx op to quantize op
+        4. Remove fake_dequantize_abs_max op
+    """
+
+    def __init__(self, _scope=None, _place=None):
+        r"""
+        Args:
+            scope(static.Scope): scope is used to initialize the new parameters.
+            place(static.CPUPlace|str): place is used to initialize the new parameters.
+            When it is string, it can be only 'cpu'.
+
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # The original graph will be rewrite.
+                >>> import paddle
+                >>> from paddle import static
+                >>> from paddle.static.quantization import QuantInt8OnednnPass
+                >>> from paddle.framework import IrGraph
+                >>> from paddle.framework import core
+
+                >>> graph = IrGraph(core.Graph(static.Program().desc), for_test=False)
+                >>> place = paddle.CPUPlace()
+                >>> onednn_pass = QuantInt8OnednnPass(static.global_scope(), place)
+                >>> onednn_pass.apply(graph)
+        """
+
+        self._scope = _scope
+        self._place = _get_paddle_place(_place)
+
+        self._quantize_type = [
+            'fake_quantize_moving_average_abs_max',
+            'fake_quantize_range_abs_max',
+        ]
+        self._dequantize_type = ['fake_dequantize_max_abs']
+        self._quantize_dequantize_type = [
+            'fake_quantize_dequantize_moving_average_abs_max'
+        ]
+
+        self._quantizable_ops = ['conv2d', 'depthwise_conv2d', 'mul']
+        self._conv_ops = ['conv2d', 'depthwise_conv2d']
+        self._pool_ops = ['pool2d']
+
+        self._in_scale = {}
+        self._max_range = {}
+        self._new_output = {}
+        self._s8_max = 127
+
+    def apply(self, graph):
+        """
+        Quantize the graph for running MKL-DNN INT8 inference. According
+        to activation quantization type, the graph will transform fake
+        quantize ops to quantize ops and remove the fake dequantize ops.
+
+        Args:
+            graph(IrGraph): the applied graph.
+        """
+
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+        ops = graph.all_op_nodes()
+
+        persistable_vars = [p.name() for p in graph.all_persistable_nodes()]
+        # Collect the _in_scales and _max_range to calculate the new scales for MKL-DNN
+        # INT8 conv2d and mul
+        for op_node in ops:
+            if op_node.name() in self._dequantize_type:
+                input_name = op_node.input("X")[0]
+                scale_name = op_node.input("Scale")[0]
+                self._in_scale[input_name] = self._load_param(
+                    self._scope, scale_name
+                )[0]
+                self._max_range[input_name] = op_node.op().attr("max_range")
+                self._new_output[input_name] = op_node.output("Out")[0]
+
+            if op_node.name() in self._quantize_dequantize_type:
+                inputs = op_node.op().input_names()
+                attrs = op_node.op().attr_names()
+                input_name = op_node.input("X")[0]
+                scale_name = op_node.input("InScale")[0]
+                self._in_scale[input_name] = self._load_param(
+                    self._scope, scale_name
+                )[0]
+                #  self._max_range[input_name] = op_node.op().attr("max_range")
+                self._new_output[input_name] = op_node.output("Out")[0]
+
+        for op_node in ops:
+            if op_node.name() in self._quantizable_ops:
+                if op_node.name() in self._conv_ops:
+                    self._transform_to_conv_onednn(graph, op_node)
+                elif op_node.name() in self._pool_ops:
+                    self._transform_to_pool_onednn(graph, op_node)
+                else:
+                    self._transform_to_mul_onednn(graph, op_node)
+            elif op_node.name() in self._quantize_type:
+                self._transform_to_quantize_onednn(graph, op_node)
+            elif op_node.name() in self._dequantize_type:
+                self._remove_fake_dequantize_op(graph, op_node)
+            self._remove_unused_var_nodes(graph)
+        return graph
+
+    def _transform_to_pool_onednn(self, graph, op):
+        output_name = op.output("Out")[0]
+        input_name = op.input("X")[0]
+
+    def _transform_to_conv_onednn(self, graph, op_node):
+        weight_name = op_node.input("Filter")[0]
+        output_name = op_node.output("Output")[0]
+        # Convert int8 range weights to fp32 range weights
+        weight = self._load_param(self._scope, weight_name)
+        w_fp32 = np.divide(
+            np.multiply(weight, self._s8_max), self._max_range[output_name]
+        )
+        w_fp32 = w_fp32.reshape(weight.shape)
+        self._restore_var(weight_name, w_fp32)
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("Input")[0]
+        )
+        weight_var_node = graph._find_node_by_name(op_node.inputs, weight_name)
+
+        # Set fake_dequantize_abs_max's output as new output of conv2d
+        output_var_node = graph._find_node_by_name(
+            graph.all_var_nodes(), self._new_output[output_name]
+        )
+        attrs = {
+            name: op_node.op().attr(name) for name in op_node.op().attr_names()
+        }
+
+        conv_op_node = graph.create_op_node(
+            op_type='fused_conv2d',
+            attrs=attrs,
+            inputs={'Input': input_var_node, 'Filter': weight_var_node},
+            outputs={'Output': output_var_node},
+        )
+
+        # Based on the Quant's scales to calculate the scales of MKL-DNN INT8 conv2d
+        scale_in = self._s8_max / self._in_scale[output_name]
+        scale_w = []
+        scale_w = [self._max_range[output_name] / self._s8_max]
+
+        conv_op_node.set_attr("Scale_weights", scale_w)
+        conv_op_node.set_attr("Scale_in", scale_in)
+        conv_op_node.set_attr("Scale_out", 1.0)
+        conv_op_node.set_attr("use_onednn", 1)
+        conv_op_node.set_attr("force_fp32_output", 1)
+        graph.link_to(input_var_node, conv_op_node)
+        graph.link_to(weight_var_node, conv_op_node)
+        graph.link_to(conv_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _transform_to_mul_onednn(self, graph, op_node):
+        # For MKL-DNN INT8 mul, input Y should be the weights
+        weight_name = op_node.input("Y")[0]
+        output_name = op_node.output("Out")[0]
+        # Convert int8 range weights to fp32 range weights
+        weight = self._load_param(self._scope, weight_name)
+        w_fp32 = np.divide(
+            np.multiply(weight, self._s8_max), self._max_range[output_name]
+        )
+        w_fp32 = w_fp32.reshape(weight.shape)
+        self._restore_var(weight_name, w_fp32)
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        weight_var_node = graph._find_node_by_name(op_node.inputs, weight_name)
+
+        # Set fake_dequantize_abs_max's output as new output of mul
+        output_var_node = graph._find_node_by_name(
+            graph.all_var_nodes(), self._new_output[output_name]
+        )
+        attrs = {
+            name: op_node.op().attr(name) for name in op_node.op().attr_names()
+        }
+
+        mul_op_node = graph.create_op_node(
+            op_type='mul',
+            attrs=attrs,
+            inputs={'X': input_var_node, 'Y': weight_var_node},
+            outputs={'Out': output_var_node},
+        )
+
+        # Based on the Quant's scales to calculate MKL-DNN INT8 mul's scales
+        scale_in = self._s8_max / self._in_scale[output_name]
+        scale_w = []
+        scale_w = [self._max_range[output_name] / self._s8_max]
+
+        mul_op_node.set_attr("scale_y", scale_w)
+        mul_op_node.set_attr("scale_x", scale_in)
+        mul_op_node.set_attr("scale_out", 1.0)
+        mul_op_node.set_attr("use_onednn", 1)
+        mul_op_node.set_attr("force_fp32_output", 1)
+        graph.link_to(input_var_node, mul_op_node)
+        graph.link_to(weight_var_node, mul_op_node)
+        graph.link_to(mul_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _transform_to_quantize_onednn(self, graph, op_node):
+        """
+        Transform fake_quantize_xx op to quantize onednn op in the graph.
+        """
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        output_var_node = graph._find_node_by_name(
+            op_node.outputs, op_node.output("Out")[0]
+        )
+        scale_in = (
+            self._s8_max
+            / self._load_param(self._scope, op_node.input("InScale")[0])[0]
+        )
+        quant_op_node = graph.create_op_node(
+            op_type='quantize',
+            attrs={
+                'data_format': 'ONEDNNLAYOUT',
+                'use_onednn': 1,
+                'Scale': scale_in,
+                'is_negative_input': 1,
+            },
+            inputs={'Input': input_var_node},
+            outputs={'Output': output_var_node},
+        )
+        graph.link_to(input_var_node, quant_op_node)
+        graph.link_to(quant_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _remove_fake_dequantize_op(self, graph, op_node):
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        graph.safe_remove_nodes(op_node)
+
+    def _load_param(self, scope, param_name):
+        return np.array(scope.find_var(param_name).get_tensor())
+
+    def _restore_var(self, name, array):
+        tensor = self._scope.find_var(name).get_tensor()
+        tensor.set(array, self._place)
+
+    def _remove_unused_var_nodes(self, graph):
+        all_used_vars = set()
+        ops = graph.all_op_nodes()
+        for op_node in ops:
+            for input_node in op_node.inputs:
+                all_used_vars.add(input_node)
+            for output_node in op_node.outputs:
+                all_used_vars.add(output_node)
+
+        all_used_vars = {n.node for n in all_used_vars}
+        all_unused_vars = set(
+            filter(
+                lambda node: node.node not in all_used_vars,
+                graph.all_var_nodes(),
+            )
+        )
+        graph.safe_remove_nodes(all_unused_vars)
+
+
+class QuantInt8MkldnnPass(QuantInt8OnednnPass):
+    @deprecated(
+        since="3.1.0",
+        update_to="paddle.static.quantization.QuantInt8OnednnPass",
+        level=1,
+        reason="QuantInt8MkldnnPass will be removed in future",
+    )
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

python/paddle/static/quantization/quanter.py

@@ -0,0 +1,534 @@
+# 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 copy
+import json
+import logging
+import os
+
+import paddle
+
+from ...base.framework import IrGraph, core
+from ..log_helper import get_logger
+from .quantization_pass import (
+    AddQuantDequantForResidual,
+    AddQuantDequantPass,
+    ConvertToInt8Pass,
+    OutScaleForInferencePass,
+    OutScaleForTrainingPass,
+    QuantizationFreezePass,
+    QuantizationTransformPass,
+)
+
+_logger = get_logger(__name__, level=logging.INFO)
+
+from . import quant_config
+from .post_training_quantization import PostTrainingQuantizationProgram
+from .quantization_pass import (
+    AddQuantDequantForInferencePass,
+    AddQuantDequantPassV2,
+    QuantizationTransformPassV2,
+    QuantWeightPass,
+)
+
+WEIGHT_QUANTIZATION_TYPES = [
+    'abs_max',
+    'channel_wise_abs_max',
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+WEIGHT_QUANTIZATION_TYPES_TENSORRT = ['channel_wise_abs_max']
+
+ACTIVATION_QUANTIZATION_TYPES = [
+    'abs_max',
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+
+ACTIVATION_QUANTIZATION_TYPES_TENSORRT = [
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+
+VALID_DTYPES = ['int8']
+
+TRANSFORM_PASS_OP_TYPES = list(
+    quant_config.SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.keys()
+)
+QUANT_DEQUANT_PASS_OP_TYPES = list(
+    quant_config.SUPPORT_ACT_QUANTIZATION_OP_DICT.keys()
+)
+
+TENSORRT_OP_TYPES = [
+    'mul',
+    'conv2d',
+    'pool2d',
+    'depthwise_conv2d',
+    'elementwise_add',
+    'leaky_relu',
+]
+
+VARS_MAPPING_TABLE = './mapping_table_for_saving_inference_model'
+
+_quant_config_default = {
+    # weight quantize type, default is 'channel_wise_abs_max'
+    'weight_quantize_type': 'channel_wise_abs_max',
+    # activation quantize type, default is 'moving_average_abs_max'
+    'activation_quantize_type': 'moving_average_abs_max',
+    # weight quantize bit num, default is 8
+    'weight_bits': 8,
+    # activation quantize bit num, default is 8
+    'activation_bits': 8,
+    # ops of name_scope in not_quant_pattern list, will not be quantized
+    'not_quant_pattern': ['skip_quant'],
+    # ops of type in quantize_op_types, will be quantized
+    'quantize_op_types': ['conv2d', 'depthwise_conv2d', 'mul'],
+    # data type after quantization, such as 'uint8', 'int8', etc. default is 'int8'
+    'dtype': 'int8',
+    # window size for 'range_abs_max' quantization. default is 10000
+    'window_size': 10000,
+    # The decay coefficient of moving average, default is 0.9
+    'moving_rate': 0.9,
+    # if True, 'quantize_op_types' will be TENSORRT_OP_TYPES
+    'for_tensorrt': False,
+    # if True, 'quantize_op_types' will be TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES
+    'is_full_quantize': False,
+    # if True, use onnx format to quant.
+    'onnx_format': True,
+    # quant post to get initial scale for quant_aware
+    'quant_post_first': False,
+    # whether scale can be train
+    'scale_trainable': True,
+}
+
+
+def load_dict():
+    with open(VARS_MAPPING_TABLE, 'r') as file:
+        data = file.read()
+        data = json.loads(data)
+        return data
+
+
+def save_dict(table):
+    with open(VARS_MAPPING_TABLE, 'w') as file:
+        file.write(json.dumps(table))
+
+
+def _parse_configs(user_config):
+    """
+    check if user's configs are valid.
+    Args:
+        user_config(dict): user's config.
+    Return:
+        configs(dict): final configs will be used.
+    """
+
+    configs = copy.deepcopy(_quant_config_default)
+    configs.update(user_config)
+
+    assert isinstance(configs['for_tensorrt'], bool) and isinstance(
+        configs['is_full_quantize'], bool
+    ), "'for_tensorrt' and 'is_full_quantize' must both be bool'"
+
+    # check if configs is valid
+    if configs['for_tensorrt']:
+        weight_types = WEIGHT_QUANTIZATION_TYPES_TENSORRT
+        activation_types = ACTIVATION_QUANTIZATION_TYPES_TENSORRT
+        platform = 'TensorRT'
+    else:
+        weight_types = WEIGHT_QUANTIZATION_TYPES
+        activation_types = WEIGHT_QUANTIZATION_TYPES
+        platform = 'PaddleLite'
+    assert configs['weight_quantize_type'] in weight_types, (
+        "Unknown weight_quantize_type: {}. {} only supports {} ".format(
+            configs['weight_quantize_type'], platform, weight_types
+        )
+    )
+
+    assert configs['activation_quantize_type'] in activation_types, (
+        "Unknown activation_quantize_type: {}. {} only supports {}".format(
+            configs['activation_quantize_type'], platform, activation_types
+        )
+    )
+
+    assert isinstance(configs['weight_bits'], int), (
+        "weight_bits must be int value."
+    )
+
+    assert configs['weight_bits'] >= 1 and configs['weight_bits'] <= 16, (
+        "weight_bits should be between 1 and 16."
+    )
+
+    assert isinstance(configs['activation_bits'], int), (
+        "activation_bits must be int value."
+    )
+
+    assert (
+        configs['activation_bits'] >= 1 and configs['activation_bits'] <= 16
+    ), "activation_bits should be between 1 and 16."
+
+    assert isinstance(configs['not_quant_pattern'], (list, str)), (
+        "not_quant_pattern must be list or str"
+    )
+
+    assert isinstance(configs['quantize_op_types'], list), (
+        "quantize_op_types must be a list"
+    )
+
+    if configs['for_tensorrt']:
+        configs['quantize_op_types'] = TENSORRT_OP_TYPES
+    elif configs['is_full_quantize']:
+        configs['quantize_op_types'] = (
+            TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES
+        )
+    else:
+        for op_type in configs['quantize_op_types']:
+            assert (op_type in QUANT_DEQUANT_PASS_OP_TYPES) or (
+                op_type in TRANSFORM_PASS_OP_TYPES
+            ), (
+                f"{op_type} is not support, \
+                        now support op types are {TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES}"
+            )
+
+    assert isinstance(configs['dtype'], str), "dtype must be a str."
+
+    assert configs['dtype'] in VALID_DTYPES, "dtype can only be " + " ".join(
+        VALID_DTYPES
+    )
+
+    assert isinstance(configs['window_size'], int), (
+        "window_size must be int value, window size for 'range_abs_max' quantization, default is 10000."
+    )
+
+    assert isinstance(configs['moving_rate'], float), (
+        "moving_rate must be float value, The decay coefficient of moving average, default is 0.9."
+    )
+
+    return configs
+
+
+def quant_aware(
+    program,
+    place,
+    config=None,
+    scope=None,
+    for_test=False,
+    weight_quantize_func=None,
+    act_quantize_func=None,
+    weight_preprocess_func=None,
+    act_preprocess_func=None,
+    optimizer_func=None,
+    executor=None,
+    return_program=False,
+    calib_config={},
+    draw_graph=False,
+    return_scale_dict=False,
+    scale_dict=None,
+    model_type=None,
+    pattern_ops=None,
+):
+    """Add quantization  and dequantization operators to "program"
+    for quantization training or testing.
+    Args:
+        program(paddle.static.Program): training or testing ``program``.
+        place(paddle.CPUPlace or paddle.CUDAPlace): This parameter represents
+            the executor run on which device.
+        config(dict, optional): configs for quantization. if None, will use default config.
+            Default: None.
+        scope(paddle.static.Scope): Scope records the mapping between variable names and variables,
+            similar to brackets in programming languages. Usually users can use
+            `paddle.static.global_scope <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_.
+            When ``None`` will use `paddle.static.global_scope() <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_ .
+            Default: ``None``.
+        for_test(bool): If the 'program' parameter is a test program, this parameter should be set to ``True``.
+            Otherwise, set to ``False``.Default: False
+        weight_quantize_func(function): Function that defines how to quantize weight. Using this
+                can quickly test if user's quantization method works or not. In this function, user should
+                both define quantization function and dequantization function, that is, the function's input
+                is non-quantized weight and function returns dequantized weight. If None, will use
+                quantization op defined by 'weight_quantize_type'.
+                Default is None.
+        act_quantize_func(function): Function that defines how to quantize activation. Using this
+                can quickly test if user's quantization method works or not. In this function, user should
+                both define quantization and dequantization process, that is, the function's input
+                is non-quantized activation and function returns dequantized activation. If None, will use
+                quantization op defined by 'activation_quantize_type'.
+                Default is None.
+        weight_preprocess_func(function): Function that defines how to preprocess weight before quantization. Using this
+                can quickly test if user's preprocess method works or not. The function's input
+                is non-quantized weight and function returns processed weight to be quantized. If None, the weight will
+                be quantized directly.
+                Default is None.
+        act_preprocess_func(function): Function that defines how to preprocess activation before quantization. Using this
+                can quickly test if user's preprocess method works or not. The function's input
+                is non-quantized activation and function returns processed activation to be quantized. If None, the activation will
+                be quantized directly.
+                Default is None.
+        optimizer_func(function): Function return a optimizer. When 'is_test' is False and user want to use self-defined
+            quantization function and preprocess function, this function must be set. Default is None.
+        exe(paddle.static.Executor): If user want to use self-defined quantization function and preprocess function, exe must be set for
+                initialization. Default is None.
+        return_program(bool): If user want return value is a Program rather than Compiled Program, This argument should be set True.
+                Default is False.
+        draw_graph(bool): whether to draw graph when quantization is initialized. In order to prevent cycle,
+                the ERNIE model needs to be set to True. Default is False.
+        return_scale_dict(bool): If user want to return scale dict, model_type and pattern_ops, this argument should be set True.
+                Default is False.
+        scale_dict(dict): Use scale dict to initialize scales in program. Default is None.
+        model_type(str): Model type can be 'transformer' or 'non-transformer'. If model type is transformer, patterns will be analyzed.
+                Default is None.
+        pattern_ops(dict): Pattern_ops contain pattern name and corresponding ops. Default is None.
+    Returns:
+        paddle.static.CompiledProgram | paddle.static.Program: Program with quantization and dequantization ``operators``
+    """
+
+    scope = paddle.static.global_scope() if not scope else scope
+    if config is None:
+        config = _quant_config_default
+    else:
+        assert isinstance(config, dict), "config must be dict"
+        config = _parse_configs(config)
+    _logger.info(f"quant_aware config {config}")
+
+    skip_tensor_list = []
+    same_scale_tensor_list = []
+
+    is_test = True if for_test else not config['scale_trainable']
+    if config['quant_post_first'] and for_test:
+        if 'quantizable_op_type' not in calib_config:
+            calib_config['quantizable_op_type'] = config['quantize_op_types']
+        exe = paddle.static.Executor() if executor is None else executor
+        post_training_quantization = PostTrainingQuantizationProgram(
+            exe,
+            program,
+            freeze_model=False,
+            skip_tensor_list=skip_tensor_list,
+            same_scale_tensor_list=same_scale_tensor_list,
+            batch_nums=10,
+            scale_dict=scale_dict,
+            return_graph=True,
+            **calib_config,
+        )
+        main_graph = post_training_quantization.quantize()
+        scale_dict = post_training_quantization._scale_dict
+        sub_graphs = list(main_graph.all_sub_graphs())
+    else:
+        main_graph = IrGraph(core.Graph(program.desc), for_test=for_test)
+        sub_graphs = list(main_graph.all_sub_graphs())
+        transform_pass_ops = []
+        quant_dequant_ops = []
+        if config.get('quant_config'):
+            transform_pass_ops = config[
+                'quant_config'
+            ].weight_quant_operation_types
+            quant_dequant_ops = config[
+                'quant_config'
+            ].activation_quant_operation_types
+        else:
+            for op_type in config['quantize_op_types']:
+                if op_type in TRANSFORM_PASS_OP_TYPES:
+                    transform_pass_ops.append(op_type)
+                elif op_type in QUANT_DEQUANT_PASS_OP_TYPES:
+                    quant_dequant_ops.append(op_type)
+        if len(transform_pass_ops) > 0:
+            transform_func = (
+                QuantizationTransformPassV2
+                if config['onnx_format']
+                else QuantizationTransformPass
+            )
+            transform_pass = transform_func(
+                scope=scope,
+                place=place,
+                weight_bits=config['weight_bits'],
+                activation_bits=config['activation_bits'],
+                activation_quantize_type=config['activation_quantize_type'],
+                weight_quantize_type=config['weight_quantize_type'],
+                window_size=config['window_size'],
+                moving_rate=config['moving_rate'],
+                quantizable_op_type=transform_pass_ops,
+                skip_pattern=config['not_quant_pattern'],
+                weight_quantize_func=weight_quantize_func,
+                act_quantize_func=act_quantize_func,
+                weight_preprocess_func=weight_preprocess_func,
+                act_preprocess_func=act_preprocess_func,
+                optimizer_func=optimizer_func,
+                executor=executor,
+                is_test=is_test,
+            )
+
+            for sub_graph in sub_graphs:
+                transform_pass.apply(sub_graph)
+
+            residual_pass = AddQuantDequantForResidual(
+                scope=scope,
+                place=place,
+                quant_bits=config['activation_bits'],
+                is_test=is_test,
+            )
+
+            for subgraph in sub_graphs:
+                residual_pass.apply(sub_graph)
+
+        if len(quant_dequant_ops) > 0:
+            qdq_func = (
+                AddQuantDequantPassV2
+                if config['onnx_format']
+                else AddQuantDequantPass
+            )
+            quant_dequant_pass = qdq_func(
+                scope=scope,
+                place=place,
+                moving_rate=config['moving_rate'],
+                quant_bits=config['activation_bits'],
+                skip_pattern=config['not_quant_pattern'],
+                quantizable_op_type=quant_dequant_ops,
+                is_test=is_test,
+            )
+
+            for sub_graph in sub_graphs:
+                quant_dequant_pass.apply(sub_graph)
+
+    out_scale_training_pass = OutScaleForTrainingPass(
+        scope=scope,
+        place=place,
+        moving_rate=config['moving_rate'],
+        is_test=is_test,
+        scale_dict=scale_dict,
+    )
+
+    for sub_graph in sub_graphs:
+        out_scale_training_pass.apply(sub_graph)
+
+    if (
+        (weight_preprocess_func is not None or act_preprocess_func is not None)
+        and not for_test
+        and not config['onnx_format']
+    ):
+        _logger.info(
+            "When a preprocess_func is used in quant_aware, Need to save a mapping table to match variable names in the convert phase."
+        )
+        _logger.info(f"The mapping table is saved as '{VARS_MAPPING_TABLE}'.")
+        for sub_graph in sub_graphs:
+            save_dict(sub_graph.out_node_mapping_table)
+
+    # TODO: remove it.
+    if draw_graph:
+        main_graph.draw('./', 'graph.pdf')
+
+    if for_test or return_program:
+        quant_program = main_graph.to_program()
+    else:
+        quant_program = paddle.static.CompiledProgram(main_graph.graph)
+
+    if return_scale_dict:
+        return quant_program, scale_dict, model_type, pattern_ops
+    else:
+        return quant_program
+
+
+def convert(program, place, config=None, scope=None, save_int8=False):
+    """
+    convert quantized and well-trained ``program`` to final  quantized
+    ``program``that can be used to  save ``inference model``.
+
+    Args:
+        program(paddle.static.Program): quantized and well-trained ``test program``.
+        place(paddle.CPUPlace or paddle.CUDAPlace): This parameter represents
+                the executor run on which device.
+        config(dict, optional): configs for convert. if set None, will use
+                default config. It must be same with config that used in
+                'quant_aware'. Default is None.
+        scope(paddle.static.Scope, optional):  Scope records the mapping between
+                variable names and variables, similar to brackets in
+                programming languages. Usually users can use
+                `paddle.static.global_scope <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_.
+                When ``None`` will use
+                `paddle.static.global_scope() <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_
+                . Default: ``None``.
+        save_int8: Whether to return ``program`` which model parameters'
+                dtype is ``int8``. This parameter can only be used to
+                get model size. Default: ``False``.
+    Returns:
+        Tuple : freezed program which can be used for inference.
+                when ``save_int8`` is False, return ``freezed_program(paddle.static.Program)``.
+                when ``save_int8`` is True, return ``freezed_program(paddle.static.Program)``
+                and ``freezed_program_int8(paddle.static.Program)``
+    """
+    scope = paddle.static.global_scope() if not scope else scope
+
+    if config is None:
+        config = _quant_config_default
+    else:
+        assert isinstance(config, dict), "config must be dict"
+        config = _parse_configs(config)
+    _logger.info(f"convert config {config}")
+    test_graph = IrGraph(core.Graph(program.desc), for_test=True)
+
+    if config['onnx_format']:
+        quant_weight_pass = QuantWeightPass(scope, place)
+        for sub_graph in test_graph.all_sub_graphs():
+            quant_weight_pass.apply(sub_graph)
+        out_scale_infer_pass = AddQuantDequantForInferencePass(
+            scope=scope, place=place, quant_bits=config['activation_bits']
+        )
+        for sub_graph in test_graph.all_sub_graphs():
+            out_scale_infer_pass.apply(sub_graph)
+    else:
+        out_scale_infer_pass = OutScaleForInferencePass(scope=scope)
+        for sub_graph in test_graph.all_sub_graphs():
+            out_scale_infer_pass.apply(sub_graph)
+        # Freeze the graph after training by adjusting the quantize
+        # operators' order for the inference.
+        freeze_pass = QuantizationFreezePass(
+            scope=scope,
+            place=place,
+            weight_bits=config['weight_bits'],
+            activation_bits=config['activation_bits'],
+            weight_quantize_type=config['weight_quantize_type'],
+        )
+        if os.path.exists(VARS_MAPPING_TABLE):
+            test_graph.out_node_mapping_table = load_dict()
+        for sub_graph in test_graph.all_sub_graphs():
+            freeze_pass.apply(sub_graph)
+
+    freezed_program = test_graph.to_program()
+
+    # Move sub blocks persistable var to global block
+    global_block = freezed_program.global_block()
+    for _op in global_block.ops:
+        if _op.type == "while":
+            _block_id = _op.attr("sub_block").id
+            _block = freezed_program.block(_block_id)
+            persistables = []
+            for _name, _var in _block.vars.items():
+                if _var.persistable:
+                    global_block._clone_variable(_var)
+                    persistables.append(_name)
+            for _name in persistables:
+                _block._remove_var(_name)
+            persistables.extend(_op.input('X'))
+            _op.desc.set_input("X", persistables)
+
+    assert not (save_int8 and config['onnx_format']), (
+        "When onnx_format=True, already saved int8 weight,so you can't set save_int8=True."
+    )
+    if save_int8:
+        convert_int8_pass = ConvertToInt8Pass(scope=scope, place=place)
+        for sub_graph in test_graph.all_sub_graphs():
+            convert_int8_pass.apply(sub_graph)
+        freezed_program_int8 = test_graph.to_program()
+        return freezed_program, freezed_program_int8
+    else:
+        return freezed_program

python/paddle/static/quantization/utils.py

@@ -0,0 +1,293 @@
+#   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 sys
+
+import numpy as np
+
+from ...base.framework import IrNode, Operator
+from .quant_config import SUPPORT_QUANTIZATION_OP_DICT
+
+_channelwise_quant_axis1_ops = [
+    'conv2d_transpose',
+    'mul',
+    'matmul',
+    'matmul_v2',
+]
+
+
+def _get_op_input_var_names(op):
+    """
+    Get the input var names of the op.
+    Args:
+        op(IrNode, Operator): the input op.
+    Returns:
+        input_var_names or None.
+    """
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    var_names = []
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return []
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][0]
+    for name in name_list:
+        var_name = op.input(name)
+        if isinstance(var_name, list):
+            var_names.extend(var_name)
+        else:
+            var_names.append(var_name)
+    return var_names
+
+
+def _get_op_output_var_names(op):
+    """ """
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    var_names = []
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return []
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][1]
+    for name in name_list:
+        var_name = op.output(name)
+        if isinstance(var_name, list):
+            var_names.extend(var_name)
+        else:
+            var_names.append(var_name)
+    return var_names
+
+
+def _get_input_name_index(op, input_var_name):
+    """Get the input name and index of the var_name in the op"""
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return None
+
+    res = None
+    for argname in SUPPORT_QUANTIZATION_OP_DICT[op_name][0]:
+        var_names = op.input(argname)
+        for index, name in enumerate(var_names):
+            if name == input_var_name:
+                res = (argname, index)
+    return res
+
+
+def _get_output_name_index(op, output_var_name):
+    """Get the output name and index of the var_name in the op"""
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return None
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][1]
+    res = None
+    for name in name_list:
+        var_name = op.output(name)
+        for index, val in enumerate(var_name):
+            if val == output_var_name:
+                res = (name, index)
+    return res
+
+
+def load_variable_data(scope, var_name):
+    '''
+    Load variable value from scope
+    '''
+    var_node = scope.find_var(var_name)
+    assert var_node is not None, "Cannot find " + var_name + " in scope."
+    tensor = np.array(var_node.get_tensor())
+    if tensor.shape == ():
+        return tensor.reshape(1)
+    else:
+        return tensor
+
+
+def set_variable_data(scope, place, var_name, np_value):
+    '''
+    Set the value of var node by name, if the node exits,
+    '''
+    assert isinstance(np_value, np.ndarray), (
+        'The type of value should be numpy array.'
+    )
+    var_node = scope.find_var(var_name)
+    if var_node is not None:
+        tensor = var_node.get_tensor()
+        tensor.set(np_value, place)
+
+
+def quant_tensor(x, scale, quant_axis=0, weight_bits=8, onnx_format=False):
+    # symmetry quant
+    def _clip(x, scale):
+        x[x > scale] = scale
+        x[x < -scale] = -scale
+        return x
+
+    bnt = (1 << (weight_bits - 1)) - 1
+    if isinstance(scale, list) and len(scale) == 1:
+        scale = scale[0]
+    if isinstance(scale, list):
+        assert quant_axis in [-1, 0, 1], 'quant_axis should be 0 or 1 for now.'
+        for i, s in enumerate(scale):
+            if s == 0.0:
+                s = 1e-8
+            if quant_axis == 0:
+                if onnx_format:
+                    x[i] = np.round(x[i] / s * bnt)
+                    x[i] = np.clip(x[i], -bnt - 1, bnt)
+                else:
+                    x[i] = _clip(x[i], s)
+                    x[i] = x[i] / s * bnt
+            else:
+                if onnx_format:
+                    x[:, i] = np.round(x[:, i] / s * bnt)
+                    x[:, i] = np.clip(x[:, i], -bnt - 1, bnt)
+                else:
+                    x[:, i] = _clip(x[:, i], s)
+                    x[:, i] = x[:, i] / s * bnt
+    else:
+        scale = 1e-8 if scale == 0.0 else scale
+        if onnx_format:
+            x = np.round(x / scale * bnt)
+            x = np.clip(x, -bnt - 1, bnt)
+        else:
+            x = _clip(x, scale)
+            x = x / scale * bnt
+    return x
+
+
+def dequant_tensor(x, scale, quant_axis=0, weight_bits=8):
+    assert quant_axis in [0, 1], 'quant_axis should be 0 or 1 for now.'
+    bnt = (1 << (weight_bits - 1)) - 1
+    if isinstance(scale, list):
+        for i, s in enumerate(scale):
+            if s == 0.0:
+                s = 1e-8
+            if quant_axis == 0:
+                x[i] = x[i] * s / bnt
+            else:
+                x[:, i] = x[:, i] * s / bnt
+    else:
+        scale = 1e-8 if scale == 0.0 else scale
+        x = x * scale / bnt
+    return x
+
+
+def bias_correction_w(x, x_quant, scale_v, quant_axis, weight_bits=8):
+    '''
+    Bias correction for weight
+    '''
+    eps = 1e-8
+    bnt = (1 << (weight_bits - 1)) - 1
+    x_dequant = x_quant.copy()
+    if isinstance(scale_v, list):
+        if quant_axis == 0:
+            for i, s in enumerate(scale_v):
+                x_dequant[i] = x_dequant[i] * s / bnt
+            quant_bias = x - x_dequant
+            mean_bias = quant_bias.reshape(quant_bias.shape[0], -1).mean(-1)
+            std_orig = x.reshape(x.shape[0], -1).std(-1)
+            std_quant = x_dequant.reshape(x_dequant.shape[0], -1).std(-1)
+            std_bias = std_orig / (std_quant + eps)
+        else:
+            for i, s in enumerate(scale_v):
+                x_dequant[:, i] = x_quant[:, i] * s / bnt
+            quant_bias = x - x_dequant
+            mean_bias = np.array(
+                [quant_bias[:, i].mean() for i in range(quant_bias.shape[1])]
+            )
+            std_orig = np.array([x[:, i].std() for i in range(x.shape[1])])
+            std_quant = np.array(
+                [x_dequant[:, i].std() for i in range(x_dequant.shape[1])]
+            )
+            std_bias = std_orig / (std_quant + eps)
+    else:
+        x_dequant = x_quant * scale_v / bnt
+        mean_bias = (x - x_dequant).mean()
+        std_bias = x.std() / (x_dequant.std() + eps)
+    if mean_bias.ndim == 1:
+        std_bias = np.resize(std_bias, x.shape)
+        mean_bias = np.resize(mean_bias, x.shape)
+
+    x_dequant = (mean_bias + x_dequant) * std_bias
+    quantized_param_v = quant_tensor(
+        x_dequant, scale_v, quant_axis, weight_bits
+    )
+    return quantized_param_v
+
+
+def stable_sigmoid(x):
+    sig = np.where(x < 0, np.exp(x) / (1 + np.exp(x)), 1 / (1 + np.exp(-x)))
+    return sig
+
+
+def calculate_quant_cos_error(orig_tensor, qdq_tensor):
+    cos_sim = np.inner(orig_tensor.flatten(), qdq_tensor.flatten()) / (
+        np.linalg.norm(orig_tensor.flatten())
+        * np.linalg.norm(qdq_tensor.flatten())
+    )
+    return cos_sim
+
+
+def move_persistable_var_to_global_block(program):
+    # Move sub blocks persistable var to global block
+    global_block = program.global_block()
+    for _op in global_block.ops:
+        if _op.type == "while":
+            _block_id = _op.attr("sub_block").id
+            _block = program.block(_block_id)
+            persistables = []
+            for _name, _var in _block.vars.items():
+                if _var.persistable:
+                    global_block._clone_variable(_var)
+                    persistables.append(_name)
+            for _name in persistables:
+                _block._remove_var(_name)
+            persistables.extend(_op.input('X'))
+            _op.desc.set_input("X", persistables)
+
+
+def l2_loss(gt, pred):
+    return ((gt - pred) ** 2).mean()
+
+
+class tqdm:
+    def __init__(self, total, bar_format='Loading|{bar}', ncols=80):
+        self.total = total
+        self.bar_format = bar_format
+        self.ncols = ncols
+        self.n = 0
+
+    def update(self, n=1):
+        self.n += n
+        a = "=" * round((self.n / self.total) * self.ncols)
+        b = " " * (self.ncols - len(a))
+        prefix = self.bar_format.split('|')[0]
+        sys.stderr.write(f"\r{prefix}|{a}=>{b}| {self.n}/{self.total}")
+        sys.stderr.flush()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stderr.write('\n')