paddlepaddle--paddle/python/paddle/distributed/auto_parallel/static/utils.py

# 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
import logging
import os
import threading
import warnings
from functools import reduce

import numpy as np

import paddle
from paddle.base.framework import use_pir_api
from paddle.base.libpaddle import pir
from paddle.base.wrapped_decorator import (
    wrap_decorator,
)
from paddle.framework import core
from paddle.framework.io_utils import is_belong_to_optimizer, is_parameter
from paddle.static import Variable

from ..process_mesh import ProcessMesh, merge_process_meshes
from .dist_attribute import DistTensorSpec, OperatorDistAttr, TensorDistAttr

OpRole = core.op_proto_and_checker_maker.OpRole
OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()

__no_shape_var_type__ = [
    core.VarDesc.VarType.READER,
    core.VarDesc.VarType.STEP_SCOPES,
    core.VarDesc.VarType.DENSE_TENSOR_ARRAY,
    core.VarDesc.VarType.FEED_MINIBATCH,
    core.VarDesc.VarType.FETCH_LIST,
]

__not_naive_data_parallel_op__ = ["expand_v2"]
_g_gradient_clip_ops = [
    "sum",
    "sqrt",
    "fill_constant",
    "elementwise_max",
    "elementwise_div",
    "stack",
    "reduce_sum",
]

partition_skip_op_list = [
    "builtin.combine",
    "builtin.split",
    "pd_op.pylayer",
    "cf.yield",
    "cf.tuple_push",
    "cf.tuple_pop",
    "cf.stack_create",
    "cf.has_elements",
]


def get_logger(log_level, name="auto_parallel"):
    logger = logging.getLogger(name)
    logger.propagate = False
    if not logger.handlers:
        logger.setLevel(log_level)
        log_handler = logging.StreamHandler()
        log_format = logging.Formatter(
            '%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
        )
        log_handler.setFormatter(log_format)
        logger.addHandler(log_handler)
    else:
        logger.setLevel(log_level)
    return logger


def is_valid_list_index(list, index):
    if index >= -len(list) and index < len(list):
        return True
    else:
        return False


def is_dim_shard(mapping):
    if mapping != -1:
        return True
    else:
        return False


def is_dim_replicate(mapping):
    if mapping == -1:
        return True
    else:
        return False


def verify_dims_mapping(dims_mapping, process_mesh):
    if dims_mapping is None:
        return False
    if not all(isinstance(d, int) for d in dims_mapping):
        return False
    for i in range(len(dims_mapping)):
        if dims_mapping[i] < -1 or dims_mapping[i] >= len(process_mesh.shape):
            return False
    for i in range(len(process_mesh.shape)):
        if dims_mapping.count(i) > 1:
            return False
    return True


def convert_to_dims_mapping(shard_spec, process_mesh):
    dims_mapping = []
    for shard in shard_spec:
        if shard is None:
            dims_mapping.append(-1)
        elif process_mesh.shape[process_mesh.dim_names.index(shard)] == 1:
            dims_mapping.append(-1)
        else:
            dims_mapping.append(process_mesh.dim_names.index(shard))
    return dims_mapping


def convert_to_shard_spec(dims_mapping, process_mesh):
    shard_spec = []
    for dim_mapping in dims_mapping:
        if dim_mapping == -1:
            shard_spec.append(None)
        else:
            shard_spec.append(process_mesh.dim_names[dim_mapping])
    return shard_spec


def verify_shard_spec(shard_spec, tensor_shape, process_mesh):
    if len(shard_spec) != len(tensor_shape):
        return False
    for shard in shard_spec:
        if shard is not None and not isinstance(shard, str):
            return False
        if shard is not None and shard not in process_mesh.dim_names:
            return False
    dims_mapping = convert_to_dims_mapping(shard_spec, process_mesh)
    if not verify_dims_mapping(dims_mapping, process_mesh):
        return False
    for i in range(len(tensor_shape)):
        if (
            dims_mapping[i] != -1
            and tensor_shape[i] > 0
            and tensor_shape[i] % process_mesh.shape[dims_mapping[i]] != 0
        ):
            return False
    return True


def compute_compatible_dim_mapping(dim_mappings):
    if not dim_mappings:
        return None
    compatible_mapping = dim_mappings[0]
    for mapping in dim_mappings:
        if compatible_mapping == -1:
            compatible_mapping = mapping
        elif mapping == -1:
            continue
        elif compatible_mapping == mapping:
            continue
        else:
            return None
    return compatible_mapping


def compute_compatible_dims_mapping(dims_mapping_list):
    if not dims_mapping_list:
        return None
    length = len(dims_mapping_list[0])
    for dims_mapping in dims_mapping_list:
        assert dims_mapping is not None, (
            "Dims mapping must not be None for compatible computation"
        )
        assert len(dims_mapping) == length, (
            "The length of dims_mapping in list must be same for compatible computation."
        )
    compatible_result = []
    for dim_mappings in zip(*dims_mapping_list):
        compatible_dim_mapping = compute_compatible_dim_mapping(
            list(dim_mappings)
        )
        if compatible_dim_mapping is None:
            return None
        compatible_result.append(compatible_dim_mapping)
    return compatible_result


def compute_compatible_process_mesh(process_mesh_list):
    compatible_process_mesh = None
    if not process_mesh_list:
        return compatible_process_mesh
    for process_mesh in process_mesh_list:
        if process_mesh is not None:
            if (
                compatible_process_mesh is None
                or compatible_process_mesh == process_mesh
            ):
                compatible_process_mesh = process_mesh
            else:
                return None
    return compatible_process_mesh


def compute_compatible_and_update_dim_mapping(dims_mapping_list, index_list):
    assert len(dims_mapping_list) == len(index_list)
    changed = False
    dim_mappings = []
    for i in range(len(dims_mapping_list)):
        assert is_valid_list_index(dims_mapping_list[i], index_list[i])
        dim_mappings.append(dims_mapping_list[i][index_list[i]])
    compatible_dim_mapping = compute_compatible_dim_mapping(dim_mappings)
    if compatible_dim_mapping is None:
        return False
    for i in range(len(dims_mapping_list)):
        if compatible_dim_mapping != dims_mapping_list[i][index_list[i]]:
            dims_mapping_list[i][index_list[i]] = compatible_dim_mapping
            changed = True
    return changed


def append_distributed_attr_suffix(name):
    """
    Append auto parallel suffix for distributed attribute name.
    """
    return name + core.kAutoParallelSuffix()


def remove_distributed_attr_suffix(name):
    """
    Remove auto parallel suffix from distributed attribute name.
    """
    return name.strip(core.kAutoParallelSuffix())


def check_distributed_attr_for_program(program, dist_context=None):
    from .dist_context import get_default_distributed_context

    if dist_context is None:
        dist_context = get_default_distributed_context()
    assert dist_context.is_initialized_for_program(), (
        "Distributed attributes must be initialized before check."
    )
    for block in program.blocks:
        for tensor in block.vars.values():
            dist_tensor = dist_context.get_dist_tensor_for_graph(tensor)
            tensor_dist_attr = dist_context.get_tensor_dist_attr_for_program(
                tensor
            )
            if (tensor_dist_attr is not None) and (not dist_tensor.is_valid()):
                return False
        for op in block.ops:
            dist_op = dist_context.get_dist_op_for_graph(tensor)
            op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
            if (op_dist_attr is not None) and (not dist_op.is_valid()):
                return False
    return True


def print_program_with_dist_attr(program, dist_context=None):
    """
    This function reuses the original program output ability with a distributed context.
    Using lock can avoid multiple threads change the default distributed context simultaneously.
    """
    lock = threading.Lock()
    lock.acquire()
    from .dist_context import (
        get_default_distributed_context,
        set_default_distributed_context,
    )

    if dist_context is None:
        dist_context = get_default_distributed_context()
        print(program, flush=True)
    else:
        original_default_context = get_default_distributed_context()
        set_default_distributed_context(dist_context)
        print(program, flush=True)
        set_default_distributed_context(original_default_context)
    lock.release()


def _get_comm_group(processes, shape, axis, rank):
    """
    Given a rank and the processes mesh the rank belongs to,
    compute the communication peers of the rank based on the give axis in the mesh.

    Example: 16 processes managed in a 4-Dimensional mesh with shape of [2, 2, 2, 2].
    the rank communication peers of rank 0 (included) are following:
    in axis 0: [0, 1]
    in axis 1: [0, 2]
    in axis 2: [0, 4]
    in axis 3: [0, 8]
    """

    # NOTE _linear_idx2coordinate assume processes mesh start with 0 and continuous
    # tricks to support processes mesh when it is not start with 0 or continuous
    assert rank in processes, (
        f"rank [{rank}] is NOT in processes group {processes}"
    )
    rank_relative = processes.index(rank)
    coordinate = _linear_idx2coordinate(shape, rank_relative)
    coordinates_in_group = [coordinate[:] for i in range(shape[axis])]

    # select comm group
    for i in range(shape[axis]):
        coordinates_in_group[i][axis] = i

    ranks_in_group_relative = [
        _coordinate2linear_idx(shape, coordinate)
        for coordinate in coordinates_in_group
    ]
    ranks_in_group = [processes[idx] for idx in ranks_in_group_relative]

    return sorted(ranks_in_group)


def _get_idx_in_axis(processes, shape, axis, rank):
    """
    Given a rank and the processes mesh the rank belongs to,
    compute the index of the rank in given axis.

    Example: 27 processes managed in a 3-Dimensional mesh with shape of [3, 3, 3].
    the index of rank 22 are:
    in axis 0: 1
    in axis 1: 1
    in axis 2: 2
    """

    # NOTE _linear_idx2coordinate assume processes mesh start with 0 and continuous
    #  tricks to support processes mesh when it is not start with 0 or continuous
    rank_relative = processes.index(rank)
    coordinate = _linear_idx2coordinate(shape, rank_relative)
    return coordinate[axis]


def _coordinate2linear_idx(mesh_shape, coordinate):
    """
    convert a coordinate in multidimensional mesh space into a scala idx in linear space.

    it use Row-major order for dimension conversion.
    so it has:  [most_significant_dim, ..., least_significant_dim]
    assume:

        the size of i-th dimension to be:  S[i]
        the index of j-th dimension is: I[j]

    linear_idx of a n dimensional coordinate is:

        I[n-1] * (S[n-2] * S[n-3] * S[n-4] *     ....    S[0]) +
        I[n-2] * (         S[n-3] * S[n-4] *     ....    S[0]) +
        I[n-3] * (                  S[n-4] *     ....    S[0]) +
        ...
        I[1]   * (                                       S[0]) +
        I[0]

    """
    # NOTE the following function work based on a strong an assumption
    # that the processes in mesh are
    #    1. starts from 0
    #    2. continuous
    # it will be wrong if the above condition does not meet,
    # e.g. process_mesh = { process_groups = [7, 8, 9,10, 12, 13, 14, 15], mesh = [2, 4]}
    # if you want a more general mapping, you should use cartesian product

    assert len(mesh_shape) == len(coordinate), (
        f"coordinate should have the same size as mesh shape, but got shape: {mesh_shape}, coordinate: {coordinate}"
    )
    for i in range(len(mesh_shape)):
        assert coordinate[i] >= 0, (
            f"index in dimension [{i}] is least than zero. coordinate: {coordinate}"
        )
        assert coordinate[i] < mesh_shape[i], (
            f"index beyond extent in dimension [{i}]. shape: {mesh_shape}, coordinate: {coordinate}"
        )

    base = mesh_shape[-1]
    linear_idx = coordinate[-1]

    # row major order
    for i in range(len(mesh_shape) - 2, -1, -1):
        linear_idx += base * coordinate[i]
        base *= mesh_shape[i]

    return linear_idx


def _linear_idx2coordinate(mesh_shape, linear_idx):
    """
    mapping a linear scala into multidimensional mesh space, return it coordinate in that space.

    it is the inverse function of _coordinate2linear_idx.
    assume:

        the size of i-th dimension to be:  S[i]
        the index of j-th dimension is: I[j]

    the coordinate given linear_idx is:

        I[0] = linear_idx                                  % S[0]
        I[0] = (linear_idx / S[0])                         % S[1]
        I[0] = (linear_idx / (S[0] * S[1]))                % S[2]
        ....

    """

    assert linear_idx >= 0, f"linear index [{linear_idx}] is least than zero"
    assert linear_idx < np.prod(mesh_shape), (
        f"linear index beyond the extent of mesh shape. shape: {mesh_shape}, linear index: {linear_idx}"
    )

    base = 1
    coordinate = [-1] * len(mesh_shape)

    for i in reversed(range(len(mesh_shape))):
        offset = linear_idx / base
        coordinate[i] = int(offset % mesh_shape[i])
        base *= mesh_shape[i]

    # row major order
    return coordinate


def _get_corresponding_rank(dist_context, target_mesh, rank):
    # TODO(JZ-LIANG) a hack method to support varying mesh in Pipeline parallelism case.
    # we assume that all mesh are evenly divide from a parent mesh and should have same size.
    # to revise this in future.

    coordinate = None
    for mesh in dist_context.process_meshes:
        if rank in mesh.process_ids and mesh.shape == target_mesh.shape:
            coordinate = _linear_idx2coordinate(
                mesh.shape, mesh.process_ids.index(rank)
            )
            break

    # assert coordinate is not None, "could NOT found rank [{}] in any registered mesh".format(
    #     rank)
    if coordinate is not None:
        return target_mesh.process_ids[
            _coordinate2linear_idx(mesh.shape, coordinate)
        ]
    else:
        return target_mesh.process_ids[0]


def _get_unshard_dist_shape(var, dist_attr):
    var_shape = var.shape
    mapping = dist_attr.dims_mapping
    mesh = dist_attr.process_mesh.shape
    assert len(var_shape) == len(mapping), (
        f"variable shape [{var_shape}] and dim_mapping [{mapping}] is NOT match !"
    )
    new_shape = []
    for idx in range(len(var_shape)):
        if var_shape[idx] == -1 or mapping[idx] == -1:
            new_shape.append(var_shape[idx])
        else:
            new_shape.append(var_shape[idx] * mesh[mapping[idx]])

    return new_shape


def make_data_unshard(dist_main_prog, dist_startup_prog, dist_context=None):
    from .dist_context import get_default_distributed_context

    if dist_context is None:
        dist_context = get_default_distributed_context()

    for var in dist_main_prog.list_vars():
        if var.is_data:
            tensor_dist_attr = dist_context.get_tensor_dist_attr_for_program(
                var
            )
            inverse_shape = _get_unshard_dist_shape(var, tensor_dist_attr)
            var.desc.set_shape(inverse_shape)
            dim_mapping = tensor_dist_attr.dims_mapping
            dim_mapping = [-1] * len(dim_mapping)
            tensor_dist_attr.dims_mapping = dim_mapping
            dist_context.set_tensor_dist_attr_for_program(var, tensor_dist_attr)


def _update_addition_info(addition_info):
    """Update default addition_info with inputs"""
    add_info = {"epoch": 0, "batch": 0, "batch_size": 0}
    if not addition_info:
        return add_info
    elif not isinstance(addition_info, dict):
        raise TypeError(
            "The type of 'addition_info' should be 'dict', "
            f"but got '{type(addition_info)}'."
        )
    else:
        for item, value in addition_info.items():
            if item not in ["epoch", "batch", "batch_size"]:
                raise ValueError(
                    "The key of 'addition_info' should be one of the "
                    f"['epoch', 'batch', 'batch_size'], but got '{item}'."
                )
            if not isinstance(value, int):
                raise ValueError(
                    "The value of 'addition_info' should be 'int', "
                    f"but got '{type(value)}'."
                )
            add_info[item] = value
        return add_info


def _check_valid_path(file_path):
    """Validity check of input file path"""
    if not file_path:
        return file_path
    elif isinstance(file_path, list):
        for file in file_path:
            if not isinstance(file, str):
                raise TypeError(
                    "The type of file path should be 'str', "
                    f"but got '{type(file)}'."
                )
            if not os.path.exists(file):
                raise ValueError(f"The file path '{file}' does not exist.")
        return file_path
    else:
        raise TypeError(
            "The type of file path should be 'list', "
            f"but got '{type(file_path)}'."
        )


def _check_param_dict(param_dict):
    if not param_dict:
        raise ValueError("'param_dict' cannot be None.")
    elif not isinstance(param_dict, dict):
        raise TypeError(
            "The type of 'param_dict' should be 'dict', "
            f"but got '{type(param_dict)}'."
        )
    else:
        for name, value in param_dict.items():
            if not isinstance(name, str):
                raise TypeError(
                    "The type of key of 'param_dict' should be 'str', "
                    f"but got '{type(name)}'."
                )
            if not isinstance(value, paddle.base.DenseTensor):
                raise TypeError(
                    "The type of value of 'param_dict' should be 'DenseTensor', "
                    f"but got '{type(value)}'."
                )
        return param_dict


def _check_dist_attr(dist_attr):
    if not dist_attr:
        return dist_attr
    elif not isinstance(dist_attr, dict):
        raise TypeError(
            "The type of 'dist_attr' should be 'dict', "
            f"but got '{type(dist_attr)}'."
        )
    else:
        for name, value in dist_attr.items():
            if not isinstance(name, str):
                raise TypeError(
                    "The type of param name of 'dist_attr' should be 'str', "
                    f"but got '{type(name)}'."
                )
            if not isinstance(value, dict):
                raise TypeError(
                    "The type of distributed attribute should be 'dict', "
                    f"but got '{type(value)}'"
                )
            attr = [
                'process_shape',
                'process_group',
                'dims_mapping',
                'dim_names',
            ]
            if list(value.keys()) != attr:
                raise ValueError(
                    "The key of distributed attribute should be "
                    "'['process_shape', 'process_group', 'dims_mapping']', "
                    f"but got {value.keys()}."
                )
        return dist_attr


def save_distributed_checkpoint(
    program,
    checkpoint_path,
    dist_attr_path,
    addition_info=None,
    is_integrated=False,
    dist_context=None,
):
    """
    Save model parameter state, optimizer state, distributed attribute and
    additional information of each rank.

    Args:
        program(Program): The program to be saved.
        checkpoint_path(str): The path of the checkpoint file to be saved.
        dist_attr_path(str): The path of distributed attribute file to be saved.
        addition_info(dict, optional): Additional information, key should be selected in ['epoch', 'batch', 'batch_size'].
            Default values are 0, when 'addition_info' is None. Default: None.
        is_integrated(bool, optional): Whether to integrate param before save. Default: False.
        dist_context(DistributedContext ,optional): collect related distributed information for program

    Returns:
        None

    Examples:
        .. code-block:: pycon

            >>> import os
            >>> from paddle.distributed.auto_parallel.static.utils import save_distributed_checkpoint

            >>> step = 16000
            >>> global_batch_size = 32
            >>> path = os.path.join("./output", "step_%d" % step)
            >>> os.makedirs(path, exist_ok=True)
            >>> program = paddle.static.Program()

            >>> add_info = {'batch': step, "batch_size": global_batch_size}
            >>> save_distributed_checkpoint(program, path, path, add_info)

    """
    from .dist_context import get_default_distributed_context

    assert isinstance(program, paddle.static.Program)
    assert isinstance(is_integrated, bool)
    if dist_context is None:
        dist_context = get_default_distributed_context()
    addition_info = _update_addition_info(addition_info)

    if not is_integrated:
        _save_distributed_state_dict(program, addition_info, checkpoint_path)
        _save_distributed_attribute(program, dist_attr_path, dist_context)
    else:
        # TODO: integrate param before save
        raise NotImplementedError(
            "Integrating parameter has not been implemented."
        )


def load_distributed_checkpoint(checkpoint_path, dist_attr_path):
    """
    Load parameter, optimizer, distributed attribute and addition_info.

    Args:
        checkpoint_path(list[str]): model parameter file path, must be in order of rank id.
        dist_attr_path(list[str]): distributed attribute file path, must be in order of rank id.

    Returns:
        param_dict(dict): parameters' value of all ranks.
        dist_attr(dict): parameters' distributed attribute.
        addition_info(dict): additional information user saved in last training.

    Notes:
        The return, 'addition_info', is belonging to the first file of checkpoint_path by default.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('Depends on external files.')
            >>> from paddle.distributed.auto_parallel.static.utils import load_distributed_checkpoint

            >>> ckpt_path = [
            ...     './model_state_rank0.pdmodel',
            ...     './model_state_rank1.pdmodel',
            ... ]
            >>> dist_attr_path = [
            ...     './dist_attr_rank0.pdattr',
            ...     './dist_attr_rank1.pdattr',
            ... ]
            >>> param_dict, dist_attr, add_info = load_distributed_checkpoint(
            ...     ckpt_path,
            ...     dist_attr_path,
            ... )
    """
    assert _check_valid_path(checkpoint_path), (
        "'checkpoint_path' cannot be None."
    )
    assert _check_valid_path(dist_attr_path), "'dist_attr_path' cannot be None."

    state_dict_info = _load_distributed_state_dict(checkpoint_path)
    dist_attr = _load_distributed_attribute(dist_attr_path)
    param_dict = state_dict_info["model"]
    addition_info = state_dict_info["addition_info"]
    return param_dict, dist_attr, addition_info


def load_checkpoint_into_program(
    checkpoint_path, dist_attr_path, program, dist_context=None
):
    """
    Load parameter, optimizer, distributed attribute and addition_info into model.

    Args:
        checkpoint_path(list[str]): model parameter file path, must be in order of rank id.
        dist_attr_path(list[str]): distributed attribute file path, must be in order of rank id.
        program(Program): the program to be updated with checkpoint_path.
        dist_context(DistributedContext ,optional): collect related distributed information for program

    Returns:
        addition_info(dict): user saved in last train.

    Notes:
        The return, 'addition_info', is belonging to the first file of checkpoint_path by default.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('Depends on external files.')
            >>> from paddle.distributed.auto_parallel.static.utils import load_checkpoint_into_program

            >>> exe.run(startup_program)
            >>> ckpt_path = [
            ...     './model_state_rank0.pdmodel',
            ...     './model_state_rank1.pdmodel',
            ... ]
            >>> dist_attr_path = [
            ...     './dist_attr_rank0.pdattr',
            ...     './dist_attr_rank1.pdattr',
            ... ]
            >>> load_checkpoint_into_program(ckpt_path, dist_attr_path, main_program)
    """
    from .dist_context import get_default_distributed_context

    assert isinstance(program, paddle.static.Program)
    assert _check_valid_path(checkpoint_path), (
        "'checkpoint_path' cannot be None."
    )
    assert _check_valid_path(dist_attr_path), "'dist_attr_path' cannot be None."
    if dist_context is None:
        dist_context = get_default_distributed_context()
    all_state_dict_info = _load_distributed_state_dict(checkpoint_path)
    all_pre_dist_attr = _load_distributed_attribute(dist_attr_path)
    all_cur_dist_attr = get_dist_attr(program, dist_context)
    all_param_dict = all_state_dict_info["model"]
    addition_info = all_state_dict_info["addition_info"]
    sliced_param_dict = merge_and_slice_parameter(
        all_param_dict, all_pre_dist_attr, all_cur_dist_attr
    )
    load_parameter_into_program(sliced_param_dict, program)

    return addition_info


def load_parameter_into_program(param_dict, program):
    """
    Load parameters into program.

    Args:
        param_dict(dict): parameters' name and value.
        program(Program): the program to be updated
    """
    assert isinstance(param_dict, dict)
    assert program and isinstance(program, paddle.static.Program)
    if not param_dict:
        return
    program.set_state_dict(param_dict)


def _save_distributed_attribute(program, dist_attr_path, dist_context):
    """Save distributed attribute of all parameters"""
    # TODO: just save a complete distributed attribute file
    rank_id = paddle.distributed.get_rank()
    dist_attr_name = os.path.join(
        dist_attr_path, f"dist_attr_rank{rank_id}.pdattr"
    )
    dist_attr_dict = {
        "model": get_dist_attr(program, dist_context),
        "world_size": paddle.distributed.get_world_size(),
    }
    paddle.save(dist_attr_dict, dist_attr_name)
    logging.info(f"Already saved distributed attribute to '{dist_attr_path}'.")


def _load_distributed_attribute(dist_attr_path):
    """Load parameters' distributed attribute from dist_attr_path"""
    total_dist_attr = {}
    for dist_attr_file in dist_attr_path:
        dist_attr = paddle.load(dist_attr_file)
        pre_world_size = dist_attr["world_size"]
        assert pre_world_size == len(dist_attr_path), (
            "The number of 'dist_attr_path' must be equal to the last training world size."
        )
        for name, attr in dist_attr["model"].items():
            if name not in total_dist_attr:
                total_dist_attr[name] = attr

    return total_dist_attr


def _save_distributed_state_dict(program, addition_info, checkpoint_path):
    """Save parameters' state_dict"""
    rank = paddle.distributed.get_rank()
    ckpt_file_name = os.path.join(
        checkpoint_path, f"model_state_rank{rank}.pdmodel"
    )
    state_dict = {
        "model": program.state_dict(),
        "world_size": paddle.distributed.get_world_size(),
        "addition_info": addition_info,
    }
    paddle.save(state_dict, ckpt_file_name)
    logging.info(f"Already saved model to '{checkpoint_path}'.")


def _load_distributed_state_dict(checkpoint_path):
    """Load parameters' state_dict from checkpoint_path"""
    all_state_dict = {}
    for idx, ckpt_file in enumerate(checkpoint_path):
        state_dict_info = paddle.load(ckpt_file, return_numpy=True)
        pre_world_size = state_dict_info["world_size"]
        assert pre_world_size == len(checkpoint_path), (
            "The number of 'checkpoint_path' must be equal to the last training world size."
        )
        if idx == 0:
            addition_info = state_dict_info["addition_info"]
        for name, value in state_dict_info["model"].items():
            if name in all_state_dict:
                all_state_dict[name].append(np.array(value))
            else:
                all_state_dict[name] = [np.array(value)]

    all_state_dict_info = {
        "model": all_state_dict,
        "addition_info": addition_info,
    }
    return all_state_dict_info


def get_dist_attr(program, dist_context=None):
    """
    Get distributed attribute of current rank.

    Args:
        program(Program): main program for training
    """
    dist_attr = {}
    if use_pir_api():
        ops = program.global_block().ops
        for op in ops:
            if op.name() == "builtin.parameter" or (
                op.name() == "pd_op.data"
                and op.has_attr("persistable")
                and op.attrs()["persistable"]
            ):
                op_dist_attr = op.dist_attr
                var_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()
                var_name = (
                    op.str_attr("parameter_name")
                    if op.name() == "builtin.parameter"
                    else op.str_attr("name")
                )
                process_mesh = var_dist_attr.process_mesh
                dist_attr[var_name] = {
                    "process_shape": process_mesh.shape,
                    "process_group": process_mesh.process_ids,
                    "dims_mapping": var_dist_attr.dims_mapping,
                    "dim_names": process_mesh.dim_names,
                }
    else:
        from .dist_context import get_default_distributed_context

        assert isinstance(program, paddle.static.Program)
        if dist_context is None:
            dist_context = get_default_distributed_context()
        for var in program.list_vars():
            if is_parameter(var) or is_belong_to_optimizer(var):
                tensor_dist_attr = (
                    dist_context.get_tensor_dist_attr_for_program(var)
                )
                process_mesh = tensor_dist_attr.process_mesh
                dims_mapping = tensor_dist_attr.dims_mapping
                dim_names = tensor_dist_attr.process_mesh.dim_names
                dist_attr[var.name] = {
                    "process_shape": process_mesh.shape,
                    "process_group": process_mesh.process_ids,
                    "dims_mapping": dims_mapping,
                    "dim_names": dim_names,
                }
    return dist_attr


def merge_and_slice_parameter(dist_param_dict, pre_dist_attr, cur_dist_attr):
    """
    Merge parameters with previous dist_attr and slice parameters with current dist_attr

    Args:
        dist_param_dict(dict): parameters' value of all ranks.
        pre_dist_attr(dict): parameters' dist_attr of last training process.
        cur_dist_attr(dict): parameters' dist_attr of current training process.

    Returns:
        dist_param_dict(dict): parameters' value of current rank.
    """
    assert _check_dist_attr(pre_dist_attr), "'pre_dist_attr' cannot be None."
    assert isinstance(dist_param_dict, dict), (
        f"The type of 'dist_param_dict' should be 'dict', but got {type(dist_param_dict)}."
    )
    for name, value in dist_param_dict.items():
        if not isinstance(name, str):
            raise TypeError(
                "The key of 'dist_param_dict' is parameter's name, "
                f"and its type should be 'str', but got {type(name)}."
            )
        if not isinstance(value, list) or not all(
            isinstance(v, np.ndarray) for v in value
        ):
            raise TypeError(
                "The value of 'dist_param_dict' is parameter's value of all ranks, "
                "and its type should be 'list(numpy.ndarray)'."
            )

    if cur_dist_attr is None:
        return {}

    param_not_in_pre = []
    param_not_in_cur = []
    logging.info("Start to merge and slice parameters.")
    for var_name in cur_dist_attr.keys():
        if var_name not in pre_dist_attr:
            param_not_in_pre.append(var_name)
            continue

        pre_attr = pre_dist_attr[var_name]
        cur_attr = cur_dist_attr[var_name]
        if pre_attr == cur_attr:
            # skip merge and slice
            rank_id = paddle.distributed.get_rank()
            index = cur_attr["process_group"].index(rank_id)
            param = dist_param_dict[var_name][index]
            dist_param_dict[var_name] = param
            continue

        pre_param = dist_param_dict[var_name]
        pre_dims_mapping = pre_attr["dims_mapping"]
        cur_dims_mapping = cur_attr["dims_mapping"]
        if len(set(pre_dims_mapping)) > 1 or -1 not in pre_dims_mapping:
            complete_param = _merge_parameter_with_dist_attr(
                pre_param, pre_attr
            )
            dist_param_dict[var_name] = complete_param
        else:
            complete_param = pre_param[0]
            dist_param_dict[var_name] = complete_param

        if len(set(cur_dims_mapping)) > 1 or -1 not in cur_dims_mapping:
            sliced_param = _slice_parameter_with_dist_attr(
                complete_param, cur_attr
            )
            dist_param_dict[var_name] = sliced_param

    for var_name in pre_dist_attr:
        if var_name not in cur_dist_attr:
            param_not_in_cur.append(var_name)
            dist_param_dict.pop(var_name)

    if param_not_in_pre:
        warnings.warn(
            f"Parameters '{param_not_in_pre}' are not found in last training process."
        )
    if param_not_in_cur:
        warnings.warn(
            f"Parameters '{param_not_in_cur}' are not found in current training process."
        )

    return dist_param_dict


def _merge_parameter_with_dist_attr(param_list, dist_attr):
    """Merge parameter with distributed attribute"""
    from .reshard import Resharder

    dims_mapping = dist_attr["dims_mapping"]
    process_shape = dist_attr["process_shape"]
    process_group = dist_attr["process_group"]
    # get the complete shape of the parameter
    complete_shape = Resharder.compute_complete_shape(
        param_list[0].shape, process_shape, dims_mapping
    )
    # merge the parameter with dist_attr
    partition_param_list = []
    merged_partition = []
    for process in process_group:
        partition_index = Resharder.compute_partition_index(
            process, complete_shape, dims_mapping, process_shape, process_group
        )
        index = process_group.index(process)
        if partition_index not in merged_partition:
            merged_partition.append(partition_index)
            _merge_parameter(
                partition_param_list,
                param_list[index],
                partition_index,
                complete_shape,
            )

    assert len(partition_param_list) == 1 or not partition_param_list, (
        "Fail to merge parameter"
    )
    complete_param = partition_param_list[0][0]
    return complete_param


def _slice_parameter_with_dist_attr(param, dist_attr):
    """Slice parameter with distributed attribute"""
    param = (
        np.array(param) if isinstance(param, paddle.base.DenseTensor) else param
    )
    dims_mapping = dist_attr["dims_mapping"]
    process_shape = dist_attr["process_shape"]
    process_group = dist_attr["process_group"]
    # slice the parameter with dist_attr
    partition_index_list = _get_split_indices(
        param.shape, dims_mapping, process_shape, process_group
    )
    sliced_param_list = _slice_parameter(
        param, partition_index_list, len(partition_index_list)
    )
    # get the current parameter's index in sliced_param_list
    rank_id = paddle.distributed.get_rank()
    sliced_param_index = _get_sliced_param_index(
        rank_id, param.shape, dims_mapping, process_shape, process_group
    )
    sliced_param = sliced_param_list[sliced_param_index]
    return sliced_param


def _merge_parameter(
    partition_param_list, param, partition_index, complete_shape
):
    """
    Merge partial parameters to a complete one.

    Returns:
        None

    Examples:
        .. code-block:: pycon

            >>> import numpy as np
            >>> from paddle.distributed.auto_parallel.static.utils import _merge_parameter

            >>> partition_param_list = [(np.array([[[1.11, 1.12]]]), [[0, 1], [0, 1], [0, 2]])]
            >>> param = np.array([[[1.13, 1.14]]])
            >>> partition_index = [[0, 1], [0, 1], [2, 4]]
            >>> complete_shape = [2, 2, 4]

            >>> _merge_parameter(partition_param_list, param, partition_index, complete_shape)
            >>> print(partition_param_list)
            [(array([[[1.11, 1.12, 1.13, 1.14]]]), [[0, 1],[0, 1],[0, 4]])]

    """
    from .reshard import Resharder

    if len(partition_param_list) == 1:
        is_complete_data = True
        for idx, item in enumerate(partition_param_list[0][1]):
            if item[0] != 0 or item[1] != complete_shape[idx]:
                is_complete_data = False
                break
        if is_complete_data:
            return

    if not partition_param_list:
        partition_param_list.append((param, partition_index))
    else:
        i = 0
        while i < len(partition_param_list):
            (
                concat_axis,
                first_order,
                new_partition,
            ) = Resharder.compute_concat_info(
                partition_param_list[i][1], partition_index
            )
            if concat_axis != -1:
                if first_order == 0:
                    new_param = np.concatenate(
                        (partition_param_list[i][0], param), axis=concat_axis
                    )
                else:
                    new_param = np.concatenate(
                        (param, partition_param_list[i][0]), axis=concat_axis
                    )

                partition_param_list.pop(i)
                _merge_parameter(
                    partition_param_list,
                    new_param,
                    new_partition,
                    complete_shape,
                )
                break
            i += 1


def _complete_op_dist_attr(program, block=None):
    if block is None:
        block = program.global_block()
    for op in block.ops:
        for sub_block in op.blocks():
            _complete_op_dist_attr(program, block=sub_block)
        if op.name() in partition_skip_op_list:
            continue

        if op.dist_attr is None:
            meshes = []
            operand_attrs = []
            result_attrs = []
            for operand in op.operands_source():
                tmp_attr = operand.dist_attr()
                if tmp_attr is None:
                    operand_attrs.append(pir.Attribute())
                    value_mesh = None
                    tmp_op_dist_attr = operand.get_defining_op().dist_attr
                    if tmp_op_dist_attr is not None:
                        value_mesh = tmp_op_dist_attr.process_mesh
                else:
                    operand_attrs.append(tmp_attr)
                    value_mesh = tmp_attr.process_mesh
                if value_mesh is not None and value_mesh not in meshes:
                    meshes.append(value_mesh)

            for result in op.results():
                tmp_attr = result.dist_attr()
                if tmp_attr is None:
                    result_attrs.append(pir.Attribute())
                else:
                    result_attrs.append(tmp_attr)
                    if tmp_attr.process_mesh not in meshes:
                        meshes.append(tmp_attr.process_mesh)
            if len(meshes) > 0:
                if len(meshes) == 1:
                    mesh = meshes[0]
                else:
                    mesh = merge_process_meshes(meshes)
                op.dist_attr = pir.create_op_dist_attribute(
                    mesh,
                    operand_attrs,
                    result_attrs,
                )


def _slice_parameter(complete_param, partition_index_list, length):
    """
    Slice a complete parameter.

    Returns:
        sliced_param_list(list): sliced parameters with 'partition_index_list'

    Examples:
        .. code-block:: pycon

            >>> import numpy as np
            >>> from paddle.distributed.auto_parallel.static.utils import _slice_parameter

            >>> complete_param = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
            >>> rank = 2
            >>> complete_shape = [1, 1, 6]
            >>> dims_mapping = [-1, -1, 0]
            >>> process_shape = [3]
            >>> process_group = [0, 1, 2]

            >>> sliced_param_list = _slice_parameter(complete_param, [[], [], [2, 4]], 3)
            >>> print(sliced_param_list)
            [array([[[1.11, 1.12]]]), array([[[1.13, 1.14]]]), array([[[1.15, 1.16]]])]

    """
    sliced_param_list = []
    axis = len(complete_param.shape) - length
    sliced_param = np.split(
        complete_param, partition_index_list[axis], axis=axis
    )
    if length == 1:
        return sliced_param
    for param in sliced_param:
        sliced_param_list.extend(
            _slice_parameter(param, partition_index_list, length - 1)
        )
    return sliced_param_list


def _get_sliced_param_index(
    rank, complete_shape, dims_mapping, process_shape, process_group
):
    """
    Get sliced_param's index of current rank in all sliced parameters list.

    Returns:
        sliced_param_index(int): the index of sliced param in sliced_param_list

    Examples:
        .. code-block:: pycon

            >>> import numpy as np
            >>> from paddle.distributed.auto_parallel.static.utils import _get_sliced_param_index

            >>> complete_param = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
            >>> rank = 2
            >>> complete_shape = [1, 1, 6]
            >>> dims_mapping = [-1, -1, 0]
            >>> process_shape = [3]
            >>> process_group = [0, 1, 2]

            >>> slice_param = _slice_parameter(complete_param, [[], [], [2, 4]], 3)
            >>> print(slice_param)
            [array([[[1.11, 1.12]]]), array([[[1.13, 1.14]]]), array([[[1.15, 1.16]]])]

            >>> index = _get_sliced_param_index(
            ...     rank,
            ...     complete_shape,
            ...     dims_mapping,
            ...     process_shape,
            ...     process_group,
            ... )
            >>> print(index)
            2
    """
    from .reshard import Resharder

    partition_index = Resharder.compute_partition_index(
        rank, complete_shape, dims_mapping, process_shape, process_group
    )
    sliced_param_index = 0
    for i, shape in enumerate(complete_shape):
        if dims_mapping[i] == -1:
            slice_shape = shape
        else:
            slice_shape = shape // process_shape[dims_mapping[i]]
        if slice_shape == 1:
            index = partition_index[i][0]
        else:
            index = (partition_index[i][0] + 1) // slice_shape
        sliced_param_index = sliced_param_index * (shape // slice_shape) + index
    return sliced_param_index


def _get_split_indices(
    complete_shape, dims_mapping, process_shape, process_group
):
    """
    Get split indices of every dimension.

    Returns:
        split_indices_list(list): the split indices of every dimension of the parameter

    Examples:
        .. code-block:: pycon

            >>> import numpy as np
            >>> from paddle.distributed.auto_parallel.static.utils import _get_split_indices

            >>> complete_param = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
            >>> complete_shape = [1, 1, 6]
            >>> dims_mapping = [-1, -1, 0]
            >>> process_shape = [3]
            >>> process_group = [0, 1, 2]

            >>> index = _get_split_indices(complete_shape, dims_mapping, process_shape, process_group)
            >>> print(index)
            [[], [], [2, 4]]
    """
    from .reshard import Resharder

    split_indices_list = []
    for process in process_group:
        partition_index = Resharder.compute_partition_index(
            process, complete_shape, dims_mapping, process_shape, process_group
        )
        if split_indices_list:
            for dim in range(len(partition_index)):
                split_indices_list[dim].extend(partition_index[dim])
        else:
            split_indices_list = partition_index
    split_indices_list = list(
        map(
            lambda x, y: list(set(x) - {y} - {0}),
            split_indices_list,
            complete_shape,
        )
    )
    split_indices_list = [sorted(x) for x in split_indices_list]
    return split_indices_list


def is_forward_op(op):
    op_role = int(op.attr('op_role'))
    return OP_ROLE_KEY in op.attr_names and (
        op_role == int(OpRole.Forward) or op_role == int(OpRole.Loss)
    )


def is_backward_op(op):
    return OP_ROLE_KEY in op.attr_names and int(
        op.all_attrs()[OP_ROLE_KEY]
    ) & int(OpRole.Backward)


def is_optimize_op(op):
    return OP_ROLE_KEY in op.attr_names and int(
        op.all_attrs()[OP_ROLE_KEY]
    ) & int(OpRole.Optimize)


def is_lr_sched_op(op):
    return OP_ROLE_KEY in op.attr_names and int(
        op.all_attrs()[OP_ROLE_KEY]
    ) & int(OpRole.Optimize.LRSched)


def is_loss_op(op):
    return OP_ROLE_KEY in op.attr_names and int(
        op.all_attrs()[OP_ROLE_KEY]
    ) == (int(OpRole.Forward) | int(OpRole.Loss))


def is_loss_grad_op(op):
    if OP_ROLE_KEY not in op.attr_names:
        return False
    op_role = int(op.all_attrs()[OP_ROLE_KEY])
    return op_role & int(OpRole.Backward) and op_role & int(OpRole.Loss)


def is_gradient_clip_op(op):
    return op.desc.has_attr("op_namescope") and op.desc.attr(
        "op_namescope"
    ).startswith("/gradient_clip")


def is_reshard_op(op):
    return op.desc.has_attr(
        "op_namescope"
    ) and "/auto_parallel/reshard" in op.desc.attr('op_namescope')


def is_prim_op(op):
    return op.type.endswith("_p")


def is_comm_op(op):
    return op.has_attr("ring_id")


def get_loss_op(block):
    loss_ops = []
    for op in block.ops:
        if is_loss_op(op):
            assert len(op.desc.output_arg_names()) == 1, (
                "loss op should only output loss var"
            )
            loss_ops.append(op)

    assert len(loss_ops) == 1, "num of loss op is not equal to one"
    return loss_ops[0]


def set_var_dist_attr(dist_context, var, dims_mapping, process_mesh, **kwargs):
    tensor_dist_attr = TensorDistAttr()
    tensor_dist_attr.dims_mapping = dims_mapping
    # TODO get global mesh group
    if isinstance(process_mesh, (list, np.ndarray)):
        tensor_dist_attr.process_mesh = ProcessMesh(process_mesh)
    elif isinstance(process_mesh, core.ProcessMesh):
        tensor_dist_attr.process_mesh = process_mesh
    else:
        raise ValueError(
            f"{process_mesh} must be a instance of ProcessMesh or list, but receive {type(process_mesh)}"
        )
    if kwargs.get("mark_annotated"):
        tensor_dist_attr.mark_annotated("dims_mapping")
        tensor_dist_attr.mark_annotated("process_mesh")
    if kwargs.get("chunk_id"):
        tensor_dist_attr.chunk_id = kwargs["chunk_id"]
    dist_context.set_tensor_dist_attr_for_program(var, tensor_dist_attr)
    return tensor_dist_attr


def naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
    new_op, process_mesh, ref_mapping, ctx, **kwargs
):
    assert process_mesh is not None
    assert ref_mapping is not None

    new_op_dist_attr = OperatorDistAttr()

    for input_varname in new_op.desc.input_arg_names():
        new_op_dist_attr.set_input_dims_mapping(input_varname, ref_mapping)
    for output_varname in new_op.desc.output_arg_names():
        new_op_dist_attr.set_output_dims_mapping(output_varname, ref_mapping)

    new_op_dist_attr.process_mesh = process_mesh
    if kwargs.get("chunk_id"):
        new_op_dist_attr.chunk_id = kwargs["chunk_id"]
    ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)


def naive_set_dist_op_attr_for_program_by_mesh(
    new_op, process_mesh, ctx, **kwargs
):
    assert process_mesh is not None

    new_op_dist_attr = OperatorDistAttr()

    for input_varname in new_op.desc.input_arg_names():
        var = new_op.block.var(input_varname)
        mapping = ctx.get_tensor_dist_attr_for_program(var).dims_mapping
        new_op_dist_attr.set_input_dims_mapping(input_varname, mapping)
    for output_varname in new_op.desc.output_arg_names():
        var = new_op.block.var(output_varname)
        mapping = ctx.get_tensor_dist_attr_for_program(var).dims_mapping
        new_op_dist_attr.set_output_dims_mapping(output_varname, mapping)

    new_op_dist_attr.process_mesh = process_mesh
    if "is_recompute" in kwargs:
        new_op_dist_attr.is_recompute = kwargs["is_recompute"]
    if "chunk_id" in kwargs:
        new_op_dist_attr.chunk_id = kwargs["chunk_id"]
    ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)


def update_op_dims_mapping_by_default_dist_impl(dist_op):
    changed = False
    op_dist_attr = dist_op.dist_attr
    op_desc = dist_op.serial_op.desc
    # The following statement will be replaced by a more elegant way
    if op_desc.type() == "shape" or op_desc.type() == "slice":
        return False
    output_names = op_desc.output_names()
    xshape_arg_names = []
    if "XShape" in output_names:
        xshape_arg_names = op_desc.output("XShape")
    batch_dim_mappings = []
    for arg_name in op_desc.input_arg_names():
        serial_tensor = dist_op.get_serial_input(arg_name)
        if serial_tensor.is_parameter:
            continue
        dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
        if len(dims_mapping) > 1:
            for idx, mapping in enumerate(dims_mapping[1:]):
                assert mapping == -1, (
                    f"{op_desc.type()} only the batch dimension (0-dim) can be sharded, but the dimension {idx} is sharded by {mapping} part."
                )
        if len(dims_mapping) >= 1:
            batch_dim_mappings.append(dims_mapping[0])
    for arg_name in op_desc.output_arg_names():
        serial_tensor = dist_op.get_serial_output(arg_name)
        if serial_tensor.is_parameter:
            continue
        dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
        if arg_name not in xshape_arg_names:
            if len(dims_mapping) > 1:
                for idx, mapping in enumerate(dims_mapping[1:]):
                    assert mapping == -1, (
                        f"{op_desc.type()} only the batch dimension (0-dim) can be sharded, but the dimension {idx} is sharded by {mapping} part."
                    )
            if len(dims_mapping) >= 1:
                batch_dim_mappings.append(dims_mapping[0])
        else:
            assert dims_mapping[0] == -1, (
                f"{op_desc.type()} only the batch dimension (1-dim) of XShape can be sharded, but the dimension 0 is sharded by {mapping} part."
            )
            if len(dims_mapping) > 2:
                for idx, mapping in enumerate(dims_mapping[2:]):
                    assert mapping == -1, (
                        f"{op_desc.type()} only the batch dimension (1-dim) of XShape can be sharded, but the dimension {idx} is sharded by {mapping} part."
                    )
            batch_dim_mappings.append(dims_mapping[1])

    compatible_dim_mapping = compute_compatible_dim_mapping(batch_dim_mappings)
    assert compatible_dim_mapping is not None, (
        "There is no compatible dim mapping."
    )
    for arg_name in op_desc.input_arg_names():
        serial_tensor = dist_op.get_serial_input(arg_name)
        if serial_tensor.is_parameter:
            continue
        dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
        if len(dims_mapping) >= 1 and compatible_dim_mapping != dims_mapping[0]:
            dims_mapping[0] = compatible_dim_mapping
            changed = True
    for arg_name in op_desc.output_arg_names():
        serial_tensor = dist_op.get_serial_output(arg_name)
        if serial_tensor.is_parameter:
            continue
        dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
        if arg_name not in xshape_arg_names:
            if (
                len(dims_mapping) >= 1
                and compatible_dim_mapping != dims_mapping[0]
            ):
                dims_mapping[0] = compatible_dim_mapping
                changed = True
        else:
            if compatible_dim_mapping != dims_mapping[1]:
                dims_mapping[1] = compatible_dim_mapping
                changed = True

    return changed


def update_op_dims_mapping_by_elementwise_like_dist_impl(dist_op):
    changed = False
    op_dist_attr = dist_op.dist_attr
    op_desc = dist_op.serial_op.desc
    input_arg_names = op_desc.input_arg_names()
    input_dims_mapping_dict = {}
    input_dims_mapping_lens = {}
    max_dims_mapping_len = -1
    for arg_name in input_arg_names:
        dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
        if max_dims_mapping_len < len(dims_mapping):
            max_dims_mapping_len = len(dims_mapping)
        input_dims_mapping_dict[arg_name] = dims_mapping
        input_dims_mapping_lens[arg_name] = len(dims_mapping)

    dims_mapping_list = []
    for arg_name in input_arg_names:
        if input_dims_mapping_lens[arg_name] < max_dims_mapping_len:
            new_dims_mapping = [-1 for _ in range(max_dims_mapping_len)]
            for i in range(input_dims_mapping_lens[arg_name]):
                new_idx = (
                    max_dims_mapping_len - input_dims_mapping_lens[arg_name]
                ) + i
                new_dims_mapping[new_idx] = input_dims_mapping_dict[arg_name][i]
            dims_mapping_list.append(new_dims_mapping)
        else:
            dims_mapping_list.append(input_dims_mapping_dict[arg_name])
    output_arg_names = op_desc.output_arg_names()
    for arg_name in output_arg_names:
        dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
        assert len(dims_mapping) == max_dims_mapping_len
        dims_mapping_list.append(dims_mapping)

    compatible_dims_mapping = compute_compatible_dims_mapping(dims_mapping_list)
    assert compatible_dims_mapping is not None, (
        "There is no compatible dim mapping."
    )

    for arg_name in input_arg_names:
        if input_dims_mapping_lens[arg_name] < max_dims_mapping_len:
            new_dims_mapping = [
                -1 for _ in range(input_dims_mapping_lens[arg_name])
            ]
            for i in range(input_dims_mapping_lens[arg_name]):
                new_idx = (
                    max_dims_mapping_len - input_dims_mapping_lens[arg_name]
                ) + i
                new_dims_mapping[i] = compatible_dims_mapping[new_idx]
            if new_dims_mapping != input_dims_mapping_dict[arg_name]:
                op_dist_attr.set_input_dims_mapping(arg_name, new_dims_mapping)
                changed = True
        else:
            if compatible_dims_mapping != input_dims_mapping_dict[arg_name]:
                op_dist_attr.set_input_dims_mapping(
                    arg_name, compatible_dims_mapping
                )
                changed = True

    for arg_name in output_arg_names:
        dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
        if compatible_dims_mapping != dims_mapping:
            op_dist_attr.set_output_dims_mapping(
                arg_name, compatible_dims_mapping
            )
            changed = True

    return changed


def get_all_distributed_main_program(
    serial_program_info, dist_context, parallelizer
):
    "Get all distributed main programs by dist_context."
    from .dist_context import DistributedOperatorContext

    cluster = serial_program_info.cluster
    copied_parallelizer = copy.deepcopy(parallelizer)
    all_dist_main_program = []
    ranks = (
        paddle.distributed.get_world_size()
        if cluster is None
        else len(cluster.get_all_devices("GPU"))
    )
    for rank_id in range(ranks):
        used_dist_context = copy.deepcopy(dist_context)
        used_dist_context._dist_op_context = DistributedOperatorContext()
        (
            _,
            _,
            dist_startup_program,
            dist_main_program,
            _,
        ) = copied_parallelizer._get_dist_program(rank_id, used_dist_context)
        all_dist_main_program.append(dist_main_program)

    return all_dist_main_program


class SerialProgramInfo:
    def __init__(
        self, train_program, startup_program, loss, optimizer, cluster=None
    ):
        self._train_program = train_program
        self._startup_program = startup_program
        self._loss = loss
        self._optimizer = optimizer
        self._cluster = cluster

    @property
    def train_program(self):
        return self._train_program

    @property
    def startup_program(self):
        return self._startup_program

    @property
    def loss(self):
        return self._loss

    @property
    def optimizer(self):
        return self._optimizer

    @property
    def cluster(self):
        return self._cluster


def get_standalone_cost_data(distributed_programs):
    def _compute_runtime(op_cost, op, vars):
        runtime = 0
        try:
            runtime = float(op_cost["op_time"])
        except:
            return runtime
        op_config = op_cost["config"]
        total_static_input_size = 0
        total_actual_input_size = 0
        parsed_info = op_config.split("\n")
        variable = "(Variable)"
        for info in parsed_info:
            variable = (
                "(Variable)" if "(Variable)" in info else "(list<Variable>"
            )
            if variable in info:
                arg_name_lower = info[: info.find(variable) - 1]
                shape_left_boundary = info.find("[")
                shape_right_boundary = info.find("]")
                assert (
                    shape_left_boundary > 0
                    and shape_right_boundary > 0
                    and shape_right_boundary > shape_left_boundary
                ), "Get shape failed."
                shape = info[
                    shape_left_boundary + 1 : shape_right_boundary
                ].split(",")
                shape = [int(x.strip()) for x in shape]
                dtype_factor = 1
                total_static_input_size += reduce(lambda x, y: x * y, shape, 1)
                if op.type == "c_embedding":
                    arg_name_lower = (
                        "w" if arg_name_lower == "weight" else "ids"
                    )
                for arg_name in op.input_names:
                    if arg_name.lower() == arg_name_lower:
                        for var_name in op.input(arg_name):
                            var = vars[var_name]
                            total_actual_input_size += reduce(
                                lambda x, y: x * y, var.shape
                            )
                        break
        assert total_static_input_size > 0 and total_actual_input_size > 0, (
            "Get input size failed."
        )

        actual_runtime = (
            total_actual_input_size / total_static_input_size * runtime
        )
        return actual_runtime

    import paddle.cost_model as cm

    cost_model = cm.CostModel()
    cost_model.static_cost_data()
    DEFAULT_MULTIPLE = 2
    OP_NAME_MAPPING = {
        "c_embedding": "embedding",
        "matmul_v2": "matmul",
        "transpose2": "transpose",
        "reshape2": "reshape",
        "unsqueeze2": "unsqueeze",
        "reduce_sum": "sum",
        "elementwise_div": "divide",
    }

    standalone_cost_data = []
    # skip ops
    not_enum_ops = [
        "create_py_reader",
        "create_double_buffer_reader",
        "read",
        "assign",
    ]
    for distributed_program in distributed_programs:
        cost_data = {}
        vars = distributed_program.global_block().vars
        for op in distributed_program.global_block().ops:
            runtime = 0
            if op.type in not_enum_ops:
                cost_data[op.desc.id()] = runtime
                continue
            dtype = (
                str(vars[op.input_arg_names[0]].dtype)
                if op.input_arg_names
                else "float32"
            )
            if int(op.attr('op_role')) == int(OpRole.Backward):
                if "_grad" in op.type:
                    forward_op_name = op.type[:-5]
                    if forward_op_name in OP_NAME_MAPPING.keys():
                        forward_op_name = OP_NAME_MAPPING[forward_op_name]
                    op_cost = cost_model.get_static_op_time(
                        forward_op_name, forward=False, dtype=dtype
                    )
                    if op_cost:
                        runtime = _compute_runtime(op_cost, op, vars)
                    else:
                        op_cost = cost_model.get_static_op_time(
                            forward_op_name, dtype=dtype
                        )
                        if op_cost:
                            runtime = 2 * _compute_runtime(op_cost, op, vars)
            elif int(op.attr('op_role')) == int(OpRole.Forward):
                op_name = (
                    OP_NAME_MAPPING[op.type]
                    if op.type in OP_NAME_MAPPING.keys()
                    else op.type
                )
                op_cost = cost_model.get_static_op_time(op_name)
                if op_cost:
                    runtime = _compute_runtime(op_cost, op, vars)

            cost_data[op.desc.id()] = runtime

        standalone_cost_data.append(cost_data)

    return standalone_cost_data


def set_dist_op_desc_original_id(dist_op_desc, op_desc, dist_context):
    op_id = op_desc.id()
    op_original_id = op_desc.original_id()
    # First, try to set the original id to the id of the op_desc
    if op_id in dist_context._dist_ops_for_program:
        dist_op_desc.set_original_id(op_id)
        return
    # Second, try to set the original id to the original_id of the op_desc
    elif op_original_id in dist_context._dist_ops_for_program:
        dist_op_desc.set_original_id(op_original_id)
        return
    # Third, print error information if we cannot find the original id
    else:
        raise AssertionError(
            "Cannot find the original id in the distributed context"
        )


def to_list(value):
    if value is None:
        return value
    if isinstance(value, (list, tuple)):
        return list(value)
    return [value]


def debug_program(program, path, name):
    filename = os.path.join(
        path, f"{name}_program.{paddle.distributed.get_rank()}"
    )
    with open(filename, 'w') as f:
        f.write(str(program))


def ring_id_to_process_group(ring_id):
    from .process_group import get_all_process_groups

    for g in get_all_process_groups():
        if g.id == ring_id:
            return g
    return None


def find_higher_order_backward_op(program):
    higher_order_op_suffix = ['_grad_grad', 'triple_grad']
    for block in program.blocks:
        for op in block.ops:
            for suffix in higher_order_op_suffix:
                if suffix in op.type:
                    return True

    return False


def get_var_numel(var):
    """
    input:
        - var: variable
    return:
        number of element in var
    """
    assert isinstance(var, Variable)
    assert -1 not in var.shape
    return reduce(lambda x, y: x * y, var.shape, 1)


def get_lr(optimizer):
    if isinstance(optimizer, paddle.optimizer.Optimizer):
        return optimizer.get_lr()
    elif isinstance(optimizer, paddle.static.Optimizer):
        if isinstance(optimizer._learning_rate, float):
            return optimizer._learning_rate
        else:
            return optimizer._learning_rate()
    else:
        raise TypeError(
            "'optimizer' must be object of class `paddle.optimizer.Optimizer`"
            f" or `paddle.static.Optimizer`, but got {type(optimizer)}."
        )


def initialize_pg_in_full_mode(all_process_groups, cur_rank):
    import socket

    from ...collective import _get_global_env

    has_recv_by_socket = []
    # This is a magic number
    magic_num = 500
    genv = _get_global_env()
    cur_rank_ip, cur_rank_port = genv.current_endpoint.split(":")
    cur_rank_recv_port = int(cur_rank_port) + magic_num
    server_socket = None
    # Large enough for recv rank
    buff_size = 1024
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((cur_rank_ip, cur_rank_recv_port))
    # The 10 is an empirical value
    server_socket.listen(10)
    client_sockets = {}
    for process_group in all_process_groups:
        if cur_rank not in process_group.ranks:
            continue
        if len(process_group.ranks) == 2:
            index = process_group.ranks.index(cur_rank)
            is_send = True if index == 0 else False
            if is_send:
                recv_rank = process_group.ranks[1]
                recv_rank_ip, recv_rank_port = genv.trainer_endpoints[
                    recv_rank
                ].split(":")
                connect_port = int(recv_rank_port) + magic_num
                client_socket = socket.socket(
                    socket.AF_INET, socket.SOCK_STREAM
                )
                client_socket.connect((recv_rank_ip, connect_port))
                client_socket.send(str(cur_rank).encode('utf-8'))
                rank = client_socket.recv(buff_size).decode('utf-8')
                rank = int(rank)
                if rank != recv_rank:
                    raise ValueError(
                        f"Please check comm pair, the recv rank should be {recv_rank} but got {rank}."
                    )
                else:
                    print(
                        f"It is able to instantiate {process_group.ranks} as sender now."
                    )
                client_socket.close()
            else:
                send_rank = process_group.ranks[0]
                while True:
                    if send_rank not in has_recv_by_socket:
                        client_socket, recv_addr = server_socket.accept()
                        rank = int(client_socket.recv(buff_size).decode())
                        client_sockets[rank] = client_socket
                        has_recv_by_socket.append(rank)
                    else:
                        client_sockets[send_rank].send(
                            str(cur_rank).encode("utf-8")
                        )
                        client_sockets[send_rank].close()
                        print(
                            f"It is able to instantiate {process_group.ranks} as receiver now."
                        )
                        break
        process_group.instantiate()
    server_socket.close()


def is_recompute_op(op):
    return (
        op.has_attr('op_namescope')
        and "/auto_parallel/rc" in op.attr('op_namescope')
        and 'exclude_rc' not in op.attr('op_namescope')
    )


def is_recompute_exclude_op(op):
    return op.has_attr('op_namescope') and 'exclude_rc' in op.attr(
        'op_namescope'
    )


def set_recompute_segments(model, losses, strategy, program):
    from ...passes.auto_parallel_recompute import RecomputeState

    if not losses:
        return

    recompute = strategy.recompute
    if not recompute.enable:
        return

    # NOTE: hack to enable recompute in engine api for GPT-3
    # TODO support more PaddleNLP/CV models here
    # extract ckpts by specific model
    ckpts = []
    if isinstance(model, paddle.nn.Layer):
        if (
            hasattr(model, "gpt")
            and model.__class__.__name__
            in [
                'GPTForPretraining',
                'GPTForPretrainingAuto',
            ]
            and hasattr(model.gpt, "checkpoints")
        ):
            ckpts = model.gpt.checkpoints
            # last recompute segment is not need to recompute
            if len(ckpts) > 2:
                ckpts.pop()
        else:
            ckpts = recompute.checkpoints
    else:
        ckpts = recompute.checkpoints

    if not ckpts:
        return

    block = program.global_block()
    rc_state = RecomputeState(block, block.ops)
    rc_state.build_stats()
    checkpoints = rc_state.sort_checkpoints(ckpts)

    segments = []
    start_idx = -1
    pre_segment_end_idx = -1
    while start_idx + 1 < len(checkpoints):
        if start_idx == -1:
            ckpt_name = checkpoints[start_idx + 1]
            if ckpt_name not in rc_state.var_op_deps:
                start_idx += 1
                continue
            op_idx_list = rc_state.var_op_deps[ckpt_name]["var_as_output_ops"]
            if op_idx_list and max(op_idx_list) > 0:
                segments.append([0, max(op_idx_list) + 1])
        else:
            flag, min_idx, max_idx = rc_state.is_subgraph(
                [checkpoints[start_idx]], [checkpoints[start_idx + 1]]
            )
            if flag:
                min_idx = rc_state._update_segment_start(
                    min_idx, pre_segment_end_idx
                )
                segments.append([min_idx, max_idx + 1])
            else:
                logging.debug(
                    f"Could not recompute op range [{min_idx}] - [{max_idx + 1}] "
                )
        start_idx += 1

    for i, segment in enumerate(segments):
        for j in range(segment[0], segment[1]):
            block.ops[j]._set_attr(
                'op_namescope', "/auto_parallel/rc_" + str(i)
            )


def get_input_split_info(cur_rank, var, dist_context):
    # deduce how the input data is split among the cluster
    tensor_dist_attr = dist_context.get_tensor_dist_attr_for_program(var)
    process_mesh = tensor_dist_attr.process_mesh
    dims_mapping = tensor_dist_attr.dims_mapping

    if cur_rank not in process_mesh.process_ids:
        rank_id = _get_corresponding_rank(dist_context, process_mesh, cur_rank)
    else:
        rank_id = cur_rank

    batch_size_axis = dims_mapping[0]
    if batch_size_axis > -1 and process_mesh.shape[batch_size_axis] > 1:
        group_ranks = _get_comm_group(
            process_mesh.process_ids,
            process_mesh.shape,
            batch_size_axis,
            rank_id,
        )
        return len(group_ranks), group_ranks.index(rank_id)

    return 1, 0


def validate_opt(optimizer):
    if optimizer is not None:
        optimizer._parameter_list = None
        optimizer._param_groups = None
        if optimizer._grad_clip and isinstance(
            optimizer._grad_clip, paddle.nn.ClipGradByGlobalNorm
        ):
            optimizer._grad_clip._async_add_n = True
    return optimizer


def set_data_parallel(x):
    from ..interface import ProcessMesh, shard_tensor
    from .process_group import get_world_process_group

    world_ranks = get_world_process_group().ranks
    process_mesh = ProcessMesh(world_ranks, ['dp'])
    shard_spec = ['dp' if len(world_ranks) > 1 else None] + [
        None for _ in range(len(x.shape) - 1)
    ]

    return shard_tensor(x, process_mesh, shard_spec)


def is_naive_data_parallel(dist_context):
    # Naive data parallel only completes dist_attr once from the front to back.
    if not dist_context.data_parallel:
        return False

    ops_type = [
        op.type
        for op in dist_context._original_serial_main_program.global_block().ops
    ]
    if (
        not set(ops_type) & set(__not_naive_data_parallel_op__)
    ) and dist_context.data_parallel:
        return True
    return False


def _copy_tensor_dist_attr_to_cpp(cpp_dist_attr, py_dist_attr):
    py_process_mesh = py_dist_attr.process_mesh
    if py_process_mesh is not None:
        cpp_dist_attr.process_mesh = core.ProcessMesh(
            py_process_mesh.shape,
            py_process_mesh.process_ids,
            ["d" + str(i) for i in range(len(py_process_mesh.shape))],
        )
    cpp_dist_attr.dims_mapping = py_dist_attr.dims_mapping
    cpp_dist_attr.annotated = py_dist_attr.annotated


def _copy_tensor_dist_attr_from_cpp(cpp_dist_attr, py_dist_attr):
    from ..process_mesh import ProcessMesh

    cpp_process_mesh = cpp_dist_attr.process_mesh
    if cpp_process_mesh is not None:
        py_dist_attr.process_mesh = ProcessMesh(
            shape=cpp_process_mesh.shape,
            process_ids=cpp_process_mesh.process_ids,
        )
    py_dist_attr.dims_mapping = cpp_dist_attr.dims_mapping
    py_dist_attr.annotated = cpp_dist_attr.annotated


def _copy_op_dist_attr_to_cpp(cpp_dist_attr, py_dist_attr):
    py_process_mesh = py_dist_attr.process_mesh
    if py_process_mesh is not None:
        cpp_dist_attr.process_mesh = core.ProcessMesh(
            py_process_mesh.shape,
            py_process_mesh.process_ids,
            ["d" + str(i) for i in range(len(py_process_mesh.shape))],
        )
    cpp_dist_attr.impl_type = py_dist_attr.impl_type
    cpp_dist_attr.impl_idx = py_dist_attr.impl_idx
    cpp_dist_attr.is_recompute = py_dist_attr.is_recompute
    cpp_dist_attr.annotated = py_dist_attr.annotated
    for name, py_tensor_dist_attr in py_dist_attr.inputs_dist_attrs.items():
        cpp_tensor_dist_attr = cpp_dist_attr.get_input_dist_attr(name)
        _copy_tensor_dist_attr_to_cpp(cpp_tensor_dist_attr, py_tensor_dist_attr)
    for name, py_tensor_dist_attr in py_dist_attr.outputs_dist_attrs.items():
        cpp_tensor_dist_attr = cpp_dist_attr.get_output_dist_attr(name)
        _copy_tensor_dist_attr_to_cpp(cpp_tensor_dist_attr, py_tensor_dist_attr)


def _copy_op_dist_attr_from_cpp(cpp_dist_attr, py_dist_attr):
    from ..process_mesh import ProcessMesh

    cpp_process_mesh = cpp_dist_attr.process_mesh
    if cpp_process_mesh is not None:
        py_dist_attr.process_mesh = ProcessMesh(
            shape=cpp_process_mesh.shape,
            process_ids=cpp_process_mesh.process_ids,
        )
    py_dist_attr.impl_type = cpp_dist_attr.impl_type
    py_dist_attr.impl_idx = cpp_dist_attr.impl_idx
    py_dist_attr.is_recompute = cpp_dist_attr.is_recompute
    py_dist_attr.annotated = cpp_dist_attr.annotated
    for name, cpp_tensor_dist_attr in cpp_dist_attr.inputs_dist_attrs.items():
        py_tensor_dist_attr = py_dist_attr.get_input_dist_attr(name)
        _copy_tensor_dist_attr_from_cpp(
            cpp_tensor_dist_attr, py_tensor_dist_attr
        )
    for name, cpp_tensor_dist_attr in cpp_dist_attr.outputs_dist_attrs.items():
        py_tensor_dist_attr = py_dist_attr.get_output_dist_attr(name)
        _copy_tensor_dist_attr_from_cpp(
            cpp_tensor_dist_attr, py_tensor_dist_attr
        )


def _copy_dist_attr_to_cpp(dist_context):
    for dist_tensor in dist_context._dist_tensors_for_program.values():
        _copy_tensor_dist_attr_to_cpp(
            dist_tensor.serial_tensor.dist_attr, dist_tensor.dist_attr
        )

    for dist_op in dist_context._dist_ops_for_program.values():
        _copy_op_dist_attr_to_cpp(
            dist_op.serial_op.dist_attr, dist_op.dist_attr
        )


def _copy_dist_attr_from_cpp(dist_context):
    for dist_tensor in dist_context._dist_tensors_for_program.values():
        _copy_tensor_dist_attr_from_cpp(
            dist_tensor.serial_tensor.dist_attr, dist_tensor.dist_attr
        )

    for dist_op in dist_context._dist_ops_for_program.values():
        _copy_op_dist_attr_from_cpp(
            dist_op.serial_op.dist_attr, dist_op.dist_attr
        )


def _copy_dist_attr_to_cpp_for_graph(dist_context):
    for node in dist_context.serial_ordered_nodes:
        if node.is_var() and node.var() is not None:
            py_dist_attr = dist_context.get_tensor_dist_attr_for_graph(node)
            cpp_dist_attr = node.var().dist_attr
            _copy_tensor_dist_attr_to_cpp(cpp_dist_attr, py_dist_attr)
        if node.is_op() and node.op() is not None:
            py_dist_attr = dist_context.get_op_dist_attr_for_graph(node)
            cpp_dist_attr = node.op().dist_attr
            _copy_op_dist_attr_to_cpp(cpp_dist_attr, py_dist_attr)


def _copy_dist_attr_from_cpp_for_graph(dist_context):
    for node in dist_context.serial_ordered_nodes:
        if node.is_var() and node.var() is not None:
            py_dist_attr = dist_context.get_tensor_dist_attr_for_graph(node)
            cpp_dist_attr = node.var().dist_attr
            _copy_tensor_dist_attr_from_cpp(cpp_dist_attr, py_dist_attr)
        if node.is_op() and node.op() is not None:
            py_dist_attr = dist_context.get_op_dist_attr_for_graph(node)
            cpp_dist_attr = node.op().dist_attr
            _copy_op_dist_attr_from_cpp(cpp_dist_attr, py_dist_attr)


def insert_dependencies_for_two_ops(
    block,
    idx,
    prior_op,
    posterior_op,
    dist_context,
    is_recompute=False,
    sync=False,
    op_namescope=None,
):
    """
    dependency: prior_op should be run before posterior_op
    """

    if is_sequential_run():
        return

    assert len(prior_op.output_arg_names) >= 1, (
        f"first op of dependency should at least have one output. [{prior_op}]"
    )
    assert len(posterior_op.input_arg_names) >= 1, (
        f"second op of dependency should at least have one input. [{posterior_op}]"
    )
    prior_op_mesh = dist_context.get_op_dist_attr_for_program(
        prior_op
    ).process_mesh
    posterior_mesh = dist_context.get_op_dist_attr_for_program(
        posterior_op
    ).process_mesh
    assert prior_op_mesh == posterior_mesh, (
        f"two ops of dependency should have same mesh but got [{prior_op_mesh}] and [{posterior_mesh}]"
    )

    def _select_best_depend_var(vars):
        # parameter should not be dep var since it maybe partition in sharding pass
        vars = [var for var in vars if not var.is_parameter]
        assert len(vars) > 0
        vars_with_numels = [(var, get_var_numel(var)) for var in vars]
        vars_with_numels.sort(key=lambda x: x[1])

        return vars_with_numels[-1][0]

    first_var = _select_best_depend_var(
        [block.var(name) for name in prior_op.output_arg_names]
    )
    second_var = _select_best_depend_var(
        [block.var(name) for name in posterior_op.input_arg_names]
    )

    return insert_dependencies_for_vars(
        block,
        idx,
        first_var,
        second_var,
        dist_context,
        OpRole.Backward,
        process_mesh=prior_op_mesh,
        is_recompute=is_recompute,
        sync=sync,
        op_namescope=op_namescope,
        use_nop=False,
    )


def insert_dependencies_for_vars(
    block,
    idx,
    prior_vars,
    post_vars,
    dist_context,
    oprole,
    process_mesh=None,
    is_recompute=False,
    sync=False,
    op_namescope=None,
    use_nop=False,
    skip_insert_when_sequential_run=True,
):
    """
    dependency: op that generates prior_vars should be run before op that generates post_vars
    """

    if skip_insert_when_sequential_run and is_sequential_run():
        return

    if isinstance(prior_vars, Variable):
        prior_vars = [prior_vars]
    if isinstance(post_vars, Variable):
        post_vars = [post_vars]
    for prior_var in prior_vars:
        assert block.has_var(prior_var.name)
    for post_var in post_vars:
        assert block.has_var(post_var.name)

    post_dist_attr = dist_context.get_tensor_dist_attr_for_program(post_vars[0])
    if process_mesh is None:
        process_mesh = post_dist_attr.process_mesh
    assert process_mesh is not None

    use_nop = True
    if use_nop:
        depend_op = block._insert_op_without_sync(
            idx,
            type='nop',
            inputs={
                "X": prior_vars,
            },
            outputs={"Out": post_vars},
        )
    else:
        depend_op = block._insert_op_without_sync(
            idx,
            type='depend',
            inputs={
                "X": post_vars,
                "Dep": prior_vars,
            },
            outputs={"Out": post_vars},
        )
    depend_op._set_attr(OP_ROLE_KEY, oprole)

    # TODO: condition can be removed when add correct dist_attr for coalesce vars and ops in sharding_pass
    if is_recompute or process_mesh != [-1]:
        depend_op_dist_attr = OperatorDistAttr()
        depend_op_dist_attr.impl_idx = 0
        depend_op_dist_attr.impl_type = "default"
        depend_op_dist_attr.process_mesh = process_mesh
        depend_op_dist_attr.is_recompute = is_recompute
        depend_op_dist_attr.chunk_id = post_dist_attr.chunk_id
        for input_varname in depend_op.desc.input_arg_names():
            var = block.var(input_varname)
            mapping = dist_context.get_tensor_dist_attr_for_program(
                var
            ).dims_mapping
            depend_op_dist_attr.set_input_dims_mapping(input_varname, mapping)
        for output_varname in depend_op.desc.output_arg_names():
            var = block.var(output_varname)
            mapping = dist_context.get_tensor_dist_attr_for_program(
                var
            ).dims_mapping
            depend_op_dist_attr.set_output_dims_mapping(output_varname, mapping)
        dist_context.set_op_dist_attr_for_program(
            depend_op, depend_op_dist_attr
        )

    if op_namescope is not None:
        depend_op._set_attr('op_namescope', f"/{op_namescope}")

    if sync:
        block._sync_with_cpp()

    return depend_op


def is_dep_skip_op(op):
    if "c_" in op.type:
        return True

    return False


def _dygraph_guard_(func):
    def __impl__(*args, **kwargs):
        if paddle.framework.in_dynamic_mode():
            return func(*args, **kwargs)
        else:
            with paddle.base.dygraph.guard():
                return func(*args, **kwargs)

    return __impl__


dygraph_guard = wrap_decorator(_dygraph_guard_)


def is_sequential_run():
    return bool(
        paddle.get_flags("FLAGS_new_executor_sequential_run")[
            "FLAGS_new_executor_sequential_run"
        ]
    )


def get_pp_degree(dist_context):
    if len(dist_context.process_meshes) < 2:
        return 0, []

    sub_process_meshes = get_sub_process_mesh(dist_context.process_meshes)
    return len(sub_process_meshes), sub_process_meshes


def get_sub_process_mesh_by_program(dist_program):
    all_ops = dist_program.global_block().ops
    process_meshes = []

    for idx, op in enumerate(all_ops):
        if "pd_op" in op.name() and op.dist_attr:
            process_mesh = op.dist_attr.process_mesh
            if process_mesh not in process_meshes:
                process_meshes.append(process_mesh)

    sub_process_meshes = get_sub_process_mesh(process_meshes)
    sub_process_meshes = sorted(
        sub_process_meshes, key=lambda x: x.process_ids[0]
    )

    return sub_process_meshes


def get_sub_process_mesh(process_meshes):
    process_ids = set()
    sub_process_meshes = copy.deepcopy(process_meshes)

    for pm in sub_process_meshes:
        process_ids |= set(pm.process_ids)

    global_pm_idx = []
    has_sub_pm = False
    for idx, pm in enumerate(sub_process_meshes):
        if len(set(pm.process_ids)) == len(process_ids):
            global_pm_idx.append(idx)
        elif set(pm.process_ids) < process_ids:
            has_sub_pm = True

    if has_sub_pm:
        for idx in reversed(global_pm_idx):
            sub_process_meshes.pop(idx)

    return sub_process_meshes


def get_pp_stage(dist_context, rank):
    pp_idx = None
    for idx, process_mesh in enumerate(dist_context.process_meshes):
        if rank in process_mesh.process_ids:
            pp_idx = idx
            break
    return pp_idx


def get_pp_stage_by_pp_degree(pp_degree):
    cur_rank = paddle.distributed.get_rank()
    return get_pp_stage_by_rank(cur_rank, pp_degree)


def get_pp_stage_by_process_mesh(process_mesh, pp_degree):
    pp_stage_for_process_mesh = None
    for rank in process_mesh.process_ids:
        pp_stage = get_pp_stage_by_rank(rank, pp_degree)
        if pp_stage_for_process_mesh is not None:
            if pp_stage != pp_stage_for_process_mesh:
                return None
            assert pp_stage == pp_stage_for_process_mesh, (
                f"Can't get pp_stage by process_mesh with different pp_stage {pp_stage} and {pp_stage_for_process_mesh}"
            )
        pp_stage_for_process_mesh = pp_stage

    return pp_stage_for_process_mesh


def get_pp_stage_by_rank(rank, pp_degree):
    word_size = paddle.distributed.get_world_size()
    pp_group_size = word_size // pp_degree
    pp_stage = rank // pp_group_size
    return pp_stage


def wrap_data_for_completion(
    dist_op, input_names: list, output_names: list, attr_names: list
):
    """
    Get data used in inferring distributed attributes, including:
      1. DistTensorSpec for each input and output tensor of this dist_op.
      2. Operator attributes of this dist_op, e.g. transpose_x in matmul op.

    Args:
      dist_op: the DistributedOperator
      input_names: list, name of the dist_op's input tensors
      output_names: list, name of the dist_op's output tensors
      attr_names: list, attribute name of the dist_op's corresponding serial op

    Returns:
      input_specs: list, DistTensorSpec for each input tensor of the dist_op
      output_specs: list, DistTensorSpec for each output tensor of the dist_op
      attrs: dict, attribute map of the dist op

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('Depends on other ops.')
            >>> from paddle.distributed.auto_parallel.static.utils import wrap_data_for_completion

            >>> op_desc = dist_op.serial_op.desc
            >>> input_name_list = []
            >>> output_name_list = []

            >>> input_name_list.append(op_desc.input('X')[0])  # 'X' is the arg name for op
            >>> input_name_list.append(op_desc.input('Y')[0])
            >>> output_name_list.append(op_desc.output('Out')[0])

            >>> attr_name_list = ['trans_x', 'trans_y']
            >>> input_specs, output_specs, attrs = wrap_data_for_completion(
            ...     dist_op,
            ...     input_name_list,
            ...     output_name_list,
            ...     attr_name_list,
            ... )

    """

    input_specs = []
    output_specs = []
    attrs = {}

    serial_op = dist_op.serial_op

    # Construct each input tensor's DistTensorSpec with shape and dist_attr
    for name in input_names:
        tensor_dist_attr = dist_op.dist_attr.get_input_dist_attr(name)
        var = serial_op.block._var_recursive(name)
        tensor_shape = var.shape
        dist_spec = DistTensorSpec(tensor_shape, tensor_dist_attr)
        input_specs.append(dist_spec)

    # Construct each output tensor's DistTensorSpec with shape and dist_attr
    for name in output_names:
        tensor_dist_attr = dist_op.dist_attr.get_output_dist_attr(name)
        var = serial_op.block._var_recursive(name)
        tensor_shape = var.shape
        dist_spec = DistTensorSpec(tensor_shape, tensor_dist_attr)
        output_specs.append(dist_spec)

    for attr_name in attr_names:
        attrs[attr_name] = serial_op.desc.attr(attr_name)

    return input_specs, output_specs, attrs


def get_dist_tensor_spec(dist_op, name, is_input=True):
    tensor_shape = dist_op.serial_op.block._var_recursive(name).shape
    if is_input:
        tensor_dist_attr = dist_op.dist_attr.get_input_dist_attr(name)
    else:
        tensor_dist_attr = dist_op.dist_attr.get_output_dist_attr(name)
    return DistTensorSpec(tensor_shape, tensor_dist_attr)


# get grad_var_to_var from distributed context, recording the mapping from backward grad variable to forward variable
# which is used for decomposing backward ops when enabling prim after distributed
def get_grad_var_to_var(dist_context):
    # get grad_var_to_var in distributed context
    grad_var_to_var_map = dist_context._dist_op_context.grad_var_to_var
    assert len(grad_var_to_var_map.keys()) == 1, "invalid grad_var_to_var"
    grad_var_to_var = grad_var_to_var_map[1]
    return grad_var_to_var


# update grad_var_to_var manually according to different distributed pass or strategy, thus recording complete and correct mapping between backward to forward
def update_grad_var_to_var(program, strategy, grad_var_to_var):
    from paddle.distributed.fleet.meta_optimizers.common import (
        OP_ROLE_KEY,
        OpRole,
    )

    # update grad_var_to_var according to different distributed pass
    first_backward_op_idx = -1
    for idx, op in enumerate(program.global_block().ops):
        # process @RESHARD variable in distributed training
        if (
            op.has_attr("op_namescope")
            and op.attr("op_namescope") == "/auto_parallel/reshard"
        ):
            reshard_op_types = [
                "split",
                "assign",
                "cast",
                "c_concat",
                "concat",
                "slice",
                "all_gather",
            ]
            if op.desc.type() in reshard_op_types:
                input_names = op.desc.input_names()
                if (
                    "X" in input_names
                    or "Input" in input_names
                    or "x" in input_names
                ):
                    inputs = (
                        op.desc.input("X")
                        if "X" in input_names
                        else (
                            op.desc.input("Input")
                            if "Input" in input_names
                            else op.desc.input("x")
                        )
                    )
                output_names = op.desc.output_names()
                if "Out" in output_names or "out" in output_names:
                    outputs = (
                        op.desc.output("Out")
                        if "Out" in output_names
                        else op.desc.output("out")
                    )
                if inputs[0] in grad_var_to_var.keys():
                    for output in outputs:
                        grad_var_to_var[output] = grad_var_to_var[inputs[0]]

        # process amp pass in distributed training
        if (
            strategy.amp.enable
            and op.has_attr(OP_ROLE_KEY)
            and (op.attr(OP_ROLE_KEY) & int(OpRole.Backward))
            and (op.attr(OP_ROLE_KEY) & int(OpRole.Loss))
        ):
            first_backward_op_idx = idx

    # process amp pass in distributed training
    if first_backward_op_idx != -1:
        scale_loss_op = program.global_block().ops[first_backward_op_idx - 1]
        scale_loss_var_name = scale_loss_op.desc.output("Out")[0]
        first_backward_op = program.global_block().ops[first_backward_op_idx]
        scale_loss_grad_var_name = first_backward_op.desc.output("Out")[0]
        if scale_loss_grad_var_name not in grad_var_to_var.keys():
            grad_var_to_var[scale_loss_grad_var_name] = scale_loss_var_name


def set_all_ops_op_role(block, op_role):
    all_ops = block.ops
    for op in all_ops:
        if op.op_role == -1:
            op.op_role = op_role
        for sub_block in op.blocks():
            set_all_ops_op_role(sub_block, op_role)


def fuse_param_func(
    fuse_params, is_qkv=False, num_heads=None, num_key_value_heads=None
):
    """fuse function for fusing weights

    (1) fuse_attention_qkv
        q => [q1,q2,q3,q4]
        k => [k1,k2,k3,k4] or [k1,k2] for GQA
        v => [v1,v2,v3,v4] or [v1,v2] for GQA
        fused weight => [q1,k1,v1,q2,k2,v2,q3,k3,v3,q4,k4,v4]
                or for GQA [q1,q2,k1,v1,q3,q4,k2,v2]
    (2) fuse_attention_ffn
        directly fuse weights to 1 parts
        [gate_weight], [up_weight] => [gate_weight, up_weight]

    Args:
        fuse_params (_type_): to be fused weights
        is_qkv (bool, optional): for attention qkv weights. Defaults to False.
        num_heads (_type_, optional): query heads. Defaults to None.
        num_key_value_heads (_type_, optional): key and value heads. Defaults to None.

    Returns:
        _type_: fused weights
    """
    concat_fn = paddle.concat
    split_fn = paddle.split

    if is_qkv:
        # fuse_attention_qkv
        assert num_heads, (
            f"num_heads should be number of heads for Q, but got {num_heads}"
        )
        assert num_key_value_heads, (
            f"num_key_value_heads should be number of key_value_heads for K and V, but got {num_key_value_heads}"
        )
        assert len(fuse_params) == 3, (
            f"fuse_params length is not equal 3, it should be Q K V list. but got length {len(fuse_params)}"
        )
        num_query_groups = num_heads // num_key_value_heads
        q_list = split_fn(fuse_params[0], num_heads, axis=-1)
        k_list = split_fn(fuse_params[1], num_key_value_heads, axis=-1)
        v_list = split_fn(fuse_params[2], num_key_value_heads, axis=-1)

        qkv_pairs = []
        for i in range(num_key_value_heads):
            qkv_pairs += q_list[
                i * num_query_groups : (i + 1) * num_query_groups
            ]
            qkv_pairs.append(k_list[i])
            qkv_pairs.append(v_list[i])
        return concat_fn(qkv_pairs, axis=-1)
    else:
        # fuse_attention_ffn
        return concat_fn(fuse_params, axis=-1)


def split_param_func(
    fused_param,
    split_nums=2,
    is_qkv=False,
    num_heads=None,
    num_key_value_heads=None,
):
    """split function for splitting weights

    (1) fuse_attention_qkv
        fused weight => [q1,k1,v1,q2,k2,v2,q3,k3,v3,q4,k4,v4]
                or for GQA [q1,q2,k1,v1,q3,q4,k2,v2]
        after split
        q => [q1,q2,q3,q4]
        k => [k1,k2,k3,k4] or [k1,k2] for GQA
        v => [v1,v2,v3,v4] or [v1,v2] for GQA
    (2) fuse_attention_ffn
        directly split weight to 2 parts
        [gate_weight, up_weight] => [gate_weight], [up_weight]

    Args:
        fused_param (_type_): len(fused_param)=1, only one weight to be split
        split_nums (int, optional): split_nums. Defaults to 2.
        is_qkv (bool, optional): for attention qkv weights. Defaults to False.
        num_heads (_type_, optional): query heads. Defaults to None.
        num_key_value_heads (_type_, optional): key and value heads. Defaults to None.

    Returns:
        _type_: split weights
    """
    concat_fn = paddle.concat
    split_fn = paddle.split

    if is_qkv:
        # fuse_attention_qkv
        assert num_heads, (
            f"num_heads should be number of heads for Q, but got {num_heads}"
        )
        assert num_key_value_heads, (
            f"num_key_value_heads should be number of key_value_heads for K and V, but got {num_key_value_heads}"
        )
        num_query_groups = num_heads // num_key_value_heads
        q_list, k_list, v_list = [], [], []
        split_heads = split_fn(
            fused_param, num_heads + 2 * num_key_value_heads, axis=-1
        )
        for i in range(num_key_value_heads):
            q_list += split_heads[
                i * (num_query_groups + 2) : (i + 1) * (num_query_groups + 2)
                - 2
            ]
            k_list.append(split_heads[(i + 1) * (num_query_groups + 2) - 2])
            v_list.append(split_heads[(i + 1) * (num_query_groups + 2) - 1])
        return (
            concat_fn(q_list, axis=-1),
            concat_fn(k_list, axis=-1),
            concat_fn(v_list, axis=-1),
        )
    else:
        # fuse_attention_ffn
        return split_fn(fused_param, split_nums, axis=-1)


def split_mesh(global_mesh: ProcessMesh, sub_mesh_dim: int):
    mesh_shape = global_mesh.shape
    mesh_ndim = len(mesh_shape)
    if sub_mesh_dim >= mesh_ndim or (
        sub_mesh_dim < 0 and -sub_mesh_dim > mesh_ndim
    ):
        raise ValueError(
            f"The sub_mesh_dim should between (-{mesh_ndim}, {mesh_ndim}]"
        )
    if sub_mesh_dim < 0:
        sub_mesh_dim += mesh_ndim

    process_ids = np.array(global_mesh.process_ids).reshape(mesh_shape)
    split_process_ids = np.split(
        process_ids, mesh_shape[sub_mesh_dim], axis=sub_mesh_dim
    )
    sub_mesh_list = []
    for sub_process_ids in split_process_ids:
        sub_mesh_list.append(
            ProcessMesh(sub_process_ids, global_mesh.dim_names)
        )

    return sub_mesh_list


# Note: This function is intended for internal use within the PaddlePaddle framework for optimizing computational graphs.
def update_pylayer_output(trivial_value):
    """
    Update the subblock within a pylayer operation by modifying its output argument.

    This function optimizes a pylayer operation by removing unnecessary outputs from the 'cf.yield' step.

    Args:
        trivale_value (pir::Value): The output argument of the pylayer operation to be modified.

    Example:
        (1) Original pylayer operation:
            (%1, %2) = "pd_op.pylayer" (%0) {
                () = "cf.tuple_pop" [id:1]
                (%3, %4) = "dist_op.xxx" [id:2]
                () = "cf.yield" [id:3] (%3, %4)
            }
        (2) After calling `update_pylayer_output(%4)`, the updated pylayer operation removes the unused output:
            (%1) = "pd_op.pylayer" (%0) {
                () = "cf.tuple_pop" [id:1]
                (%3) = "dist_op.xxx" [id:2]
                () = "cf.yield" [id:3] (%3)
            }

    Args:
        trivale_value(pir::Value): The output argument of the pylayer op to be updated.
    """
    define_op = trivial_value.get_defining_op()
    if define_op.get_parent_block().parent_op.name() != "pd_op.pylayer":
        return
    paddle.pir.set_insertion_point(define_op)
    fake_value = paddle.static.data(
        name="_fake_pylayer_out",
        shape=trivial_value.shape,
        dtype=trivial_value.dtype,
    )
    fake_value.set_type(trivial_value.type())
    trivial_value.replace_all_uses_with(fake_value)