# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collections
import copy
import logging
import os
import queue
import re

import paddle
from paddle.base.core import (  # noqa: F401
    contains_spmd_rule,
    get_phi_spmd_rule,
)
from paddle.base.framework import Operator
from paddle.base.log_helper import get_logger
from paddle.distributed.fleet.meta_optimizers.common import OpRole
from paddle.framework import core

from ..process_mesh import ProcessMesh, compute_compatible_process_mesh
from .dist_attribute import OperatorDistAttr, TensorDistAttr
from .dist_context import _node_id
from .operators.common import (
    _gradient_sync_by_partial_ops,
    find_compatible_distributed_operator_impls,
    find_distributed_operator_impl_container,
)
from .process_group import get_world_process_group
from .utils import (
    __no_shape_var_type__,
    _g_gradient_clip_ops,
    get_pp_degree,
    is_gradient_clip_op,
    is_loss_grad_op,
    is_loss_op,
    is_naive_data_parallel,
    naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
    set_var_dist_attr,
)

_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
)
__skip_dims_mapping_op__ = [
    "create_py_reader",
    "create_double_buffer_reader",
    "while",
    "read",
]

_skip_propagation_prefix = "Auto_Parallel_Completion_Skipped"
_max_propagation_step = 500


def mark_as_sharding_propagation_skip_op(op):
    prefix = op.attr("op_namescope") if op.has_attr("op_namescope") else '/'
    op._set_attr('op_namescope', prefix + _skip_propagation_prefix)


def is_sharding_propagation_skip_op(op):
    if isinstance(op, paddle.base.libpaddle.OpDesc):
        op_desc = op
    elif isinstance(op, Operator):
        op_desc = op.desc
    else:
        raise RuntimeError(f"static mode operator is expected but got [{op}]")
    return op_desc.has_attr(
        "op_namescope"
    ) and _skip_propagation_prefix in op_desc.attr("op_namescope")


def compute_compatible_dim_mapping(dim_mapping_list):
    """Compute the compatible dim mapping given a list of dim mapping."""
    if not dim_mapping_list:
        return None

    def _compute_compatible_dim_mapping_of_two(dm1, dm2):
        if dm1 == -1:
            return True, dm2
        if dm2 == -1:
            return True, dm1
        if dm1 == dm2:
            return True, dm1
        return False, None

    compatible_result = -1
    for mapping in dim_mapping_list:
        compatible, compatible_result = _compute_compatible_dim_mapping_of_two(
            compatible_result, mapping
        )
        if not compatible:
            return None
    return compatible_result


def compute_compatible_dims_mapping(dims_mapping_list):
    """Compute the compatible dims mapping given a list of dims mapping.
    Each of dims mapping is also a list.
    """
    if not dims_mapping_list:
        return None
    length = len(dims_mapping_list[0])
    for dims_mapping in dims_mapping_list:
        if dims_mapping is None:
            return None
        if len(dims_mapping) != length:
            return None
    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 merge_process_mesh_two(pm1, pm2):
    process_set1 = set()
    process_set2 = set()
    if pm1 is None and pm2 is None:
        return None
    if pm1 is not None:
        process_set1 = set(pm1.process_ids)
    if pm2 is not None:
        process_set2 = set(pm2.process_ids)
    merged_process_set = process_set1.union(process_set2)
    merged_process_mesh = ProcessMesh(list(merged_process_set))
    return merged_process_mesh


def _validate_dims_mapping(dims_mapping, process_mesh):
    if dims_mapping is None:
        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 _can_apply_infer_spmd_rule(dist_op):
    enable = os.getenv("FLAGS_infer_spmd_enable", True)
    if isinstance(enable, str):
        enable = enable.lower()
        enable = True if enable == 'true' else False
    enable = bool(enable)

    # TODO remove me. ops to be adapted: squeeze2
    __adapted_ops__ = [
        "fused_rotary_position_embedding",
        "matmul_v2",
        "elementwise_div",
        "fused_softmax_mask_upper_triangle",
        "elementwise_add",
        "elementwise_mul",
        "assign",
        "scale",
        "dropout",
        "reduce_sum",
        "layer_norm",
        "lookup_table_v2",
        "reshape2",
        "transpose2",
        "split",
        "unsqueeze2",
        "silu",
        "concat",
        "expand_as_v2",
        "swiglu",
        "tile",
        "fused_rms_norm",
        "strided_slice",
        "stack",
        "gather_nd",
    ]
    parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
    if parallel_ce == "true":
        __adapted_ops__.append("softmax_with_cross_entropy")
    op_type = dist_op.serial_op.type
    return enable and contains_spmd_rule(op_type) and op_type in __adapted_ops__


def _update_op_dims_mapping_and_distoperatorimpl(
    dist_op, original_op_dist_attr, changed
):
    dist_op_container = find_distributed_operator_impl_container(dist_op)
    _logger.debug(
        f"Update Op [{dist_op.serial_op.type}] using DistOpContainer [{dist_op_container.type}]."
    )

    updated = dist_op_container.update_dims_mapping(dist_op)
    changed = updated or changed
    # TODO(ljz) remove the below code once we introduce general reshard to replace specific distopimpls
    reverted = dist_op_container.mapping_to_dist_operator_impl(
        dist_op, original_op_dist_attr
    )
    _logger.debug(
        f"Op [{dist_op.serial_op.type}] use dist op impl [{dist_op.dist_attr.impl_type}] idx [{dist_op.dist_attr.impl_idx}]."
    )
    return changed and not (reverted)


class Completer:
    def __init__(self, dist_context):
        assert dist_context is not None
        self._dist_context = dist_context
        self._has_prepared = False

    def _update_tensor_node_dims_mapping(self, tensor_node, fwd=True):
        changed = False
        if (not tensor_node.is_var()) or (tensor_node.var() is None):
            return False
        tensor_desc = tensor_node.var()
        # Skip reader tensor
        if tensor_desc.type() in __no_shape_var_type__:
            return False
        tensor_dist_attr = self._dist_context.get_tensor_dist_attr_for_graph(
            tensor_node
        )
        assert tensor_dist_attr is not None
        if tensor_dist_attr.is_annotated("dims_mapping"):
            return False

        tensor_dims_mapping = tensor_dist_attr.dims_mapping
        if fwd:
            dims_mapping_list = []
            for pred_op_node in tensor_node.inputs:
                if pred_op_node.op() is not None:
                    if (
                        pred_op_node.op().type() == "create_py_reader"
                        or pred_op_node.op().type()
                        == "create_double_buffer_reader"
                        or pred_op_node.op().type() == "read"
                        # or is_sharding_propagation_skip_op(pred_op_node.op()) # reshard should only fwd tensor propagation
                    ):
                        continue
                    op_dist_attr = (
                        self._dist_context.get_op_dist_attr_for_graph(
                            pred_op_node
                        )
                    )
                    if (
                        op_dist_attr.process_mesh
                        == tensor_dist_attr.process_mesh
                    ):
                        op_dims_mapping = op_dist_attr.get_output_dims_mapping(
                            tensor_desc.name()
                        )
                        dims_mapping_list.append(op_dims_mapping)
            dims_mapping_list.append(tensor_dims_mapping)
            compatible_dims_mapping = compute_compatible_dims_mapping(
                dims_mapping_list
            )
            if not _validate_dims_mapping(
                compatible_dims_mapping, tensor_dist_attr.process_mesh
            ):
                return False
            if (compatible_dims_mapping is not None) and (
                compatible_dims_mapping != tensor_dims_mapping
            ):
                tensor_dist_attr.dims_mapping = compatible_dims_mapping
                changed = True
        else:
            dims_mapping_list = []
            for succ_op_node in tensor_node.outputs:
                if succ_op_node.op() is not None:
                    if (
                        succ_op_node.op().type() == "create_py_reader"
                        or succ_op_node.op().type()
                        == "create_double_buffer_reader"
                        or succ_op_node.op().type() == "read"
                        or is_sharding_propagation_skip_op(succ_op_node.op())
                    ):
                        continue
                    op_dist_attr = (
                        self._dist_context.get_op_dist_attr_for_graph(
                            succ_op_node
                        )
                    )
                    if (
                        op_dist_attr.process_mesh
                        == tensor_dist_attr.process_mesh
                    ):
                        op_dims_mapping = op_dist_attr.get_input_dims_mapping(
                            tensor_desc.name()
                        )
                        dims_mapping_list.append(op_dims_mapping)
            dims_mapping_list.append(tensor_dims_mapping)
            compatible_dims_mapping = compute_compatible_dims_mapping(
                dims_mapping_list
            )
            if not _validate_dims_mapping(
                compatible_dims_mapping, tensor_dist_attr.process_mesh
            ):
                return False
            if (compatible_dims_mapping is not None) and (
                compatible_dims_mapping != tensor_dims_mapping
            ):
                tensor_dist_attr.dims_mapping = compatible_dims_mapping
                changed = True
        return changed

    def _update_op_node_dims_mapping(self, op_node, fwd=True):
        changed = False
        op_desc = op_node.op()

        # step0: skip corner cases
        if (not op_node.is_op()) or (op_node.op() is None):
            return False
        # Skip reader op
        if (
            op_desc.type() in __skip_dims_mapping_op__
            or is_sharding_propagation_skip_op(op_node.op())
        ):
            return False

        dist_op = self._dist_context.get_dist_op_for_graph(op_node)
        op_dist_attr = dist_op.dist_attr
        original_op_dist_attr = copy.deepcopy(op_dist_attr)

        # step 1: merge the dims mappings from tensor nodes to op nodes
        if fwd:
            node_list = op_node.inputs
        else:
            node_list = op_node.outputs
        for tensor_node in node_list:
            if not tensor_node.is_var() or tensor_node.var() is None:
                continue
            if tensor_node.var().type() == core.VarDesc.VarType.READER:
                continue

            tensor_desc = tensor_node.var()
            if fwd:
                annotated = op_dist_attr.is_annotated_input_dims_mapping(
                    tensor_desc.name()
                )
            else:
                annotated = op_dist_attr.is_annotated_output_dims_mapping(
                    tensor_desc.name()
                )
            if annotated:
                continue

            tensor_dist_attr = (
                self._dist_context.get_tensor_dist_attr_for_graph(tensor_node)
            )
            if op_dist_attr.process_mesh == tensor_dist_attr.process_mesh:
                tensor_dims_mapping = tensor_dist_attr.dims_mapping
                if fwd:
                    op_dims_mapping = op_dist_attr.get_input_dims_mapping(
                        tensor_desc.name()
                    )
                else:
                    op_dims_mapping = op_dist_attr.get_output_dims_mapping(
                        tensor_desc.name()
                    )

                compatible_dims_mapping = compute_compatible_dims_mapping(
                    [op_dims_mapping, tensor_dims_mapping]
                )
                if not _validate_dims_mapping(
                    compatible_dims_mapping, op_dist_attr.process_mesh
                ):
                    continue
                if (compatible_dims_mapping is not None) and (
                    compatible_dims_mapping != op_dims_mapping
                ):
                    if fwd:
                        op_dist_attr.set_input_dims_mapping(
                            tensor_desc.name(), compatible_dims_mapping
                        )
                    else:
                        op_dist_attr.set_output_dims_mapping(
                            tensor_desc.name(), compatible_dims_mapping
                        )
                    changed = True

        # step 2: Infer & Update dims mapping of op node using SPMD Rule.
        if _can_apply_infer_spmd_rule(dist_op):
            _logger.debug(
                f"Op [{dist_op.serial_op.type}] update dims mapping using New InferSPMD Rule."
            )
            return _update_op_dims_mapping_and_distoperatorimpl(
                dist_op, original_op_dist_attr, changed
            )
        else:
            _logger.debug(
                f"Op [{dist_op.serial_op.type}] update dims mapping using Original DistOp Rule."
            )
            # update_op_dims_mapping_v1()
            op_dist_impls = find_compatible_distributed_operator_impls(
                dist_op, fwd=fwd
            )
            if op_dist_impls is not None:
                not_compatible = True
                backup_op_dist_attr = copy.deepcopy(op_dist_attr)
                backup_changed = changed
                for op_dist_impl in op_dist_impls:
                    dim_changed = op_dist_impl.update_dims_mapping(dist_op)
                    if dim_changed:
                        changed = True
                    if (
                        op_dist_impl.is_auto_compatible(dist_op)
                        and dist_op.validate_dist_attr()
                    ):
                        op_dist_attr.impl_type = op_dist_impl.type
                        op_dist_attr.impl_idx = op_dist_impl.idx
                        not_compatible = False
                        break
                    else:
                        dist_op.dist_attr = backup_op_dist_attr
                        changed = backup_changed
                if not_compatible:
                    dist_op.dist_attr = original_op_dist_attr
                    changed = False
            else:
                dist_op.dist_attr = original_op_dist_attr
                changed = False

            return changed

    def _update_dims_mapping_between_graphs(self):
        changed = False
        for parent_node, child_node in self._node_pairs_between_graphs:
            parent_node_dist_attr = self._dist_context.get_dist_attr_for_graph(
                parent_node
            )
            child_node_dist_attr = self._dist_context.get_dist_attr_for_graph(
                child_node
            )
            if (
                parent_node_dist_attr.process_mesh
                != child_node_dist_attr.process_mesh
            ):
                continue
            parent_node_dims_mapping = parent_node_dist_attr.dims_mapping
            child_node_dims_mapping = child_node_dist_attr.dims_mapping
            compatible_dims_mapping = compute_compatible_dims_mapping(
                [parent_node_dims_mapping, child_node_dims_mapping]
            )
            if not _validate_dims_mapping(
                compatible_dims_mapping, parent_node_dist_attr.process_mesh
            ):
                return False
            if (compatible_dims_mapping is not None) and (
                compatible_dims_mapping != parent_node_dims_mapping
            ):
                parent_node_dist_attr.dims_mapping = compatible_dims_mapping
                changed = True
            if (compatible_dims_mapping is not None) and (
                compatible_dims_mapping != child_node_dims_mapping
            ):
                child_node_dist_attr.dims_mapping = compatible_dims_mapping
                changed = True
        return changed

    def _update_dims_mapping_for_special(self):
        # Set the dims_mapping of a tensor to the dims_mapping inside the op which produces it
        op_nodes = self._dist_context._serial_ordered_op_nodes
        # NOTE: this list may be changed if Paddle changes the existing rules.
        related_reader_ops = [
            "create_py_reader",
            "create_double_buffer_reader",
            "read",
        ]
        for op_node in op_nodes:
            if (
                op_node.op() is not None
                and op_node.op().type() in related_reader_ops
            ):
                continue
            op_dist_attr = self._dist_context.get_dist_attr_for_graph(op_node)
            for tensor_node in op_node.outputs:
                if tensor_node.is_var() and tensor_node.var() is not None:
                    if tensor_node.var().type() == core.VarDesc.VarType.READER:
                        continue
                    tensor_desc = tensor_node.var()
                    tensor_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_graph(
                            tensor_node
                        )
                    )
                    if (
                        op_dist_attr.process_mesh
                        == tensor_dist_attr.process_mesh
                    ):
                        op_dims_mapping = op_dist_attr.get_output_dims_mapping(
                            tensor_desc.name()
                        )
                        tensor_dist_attr.dims_mapping = op_dims_mapping

    def _update_dims_mapping(self):
        # Complete dims_mapping for each node
        step = 0
        reach_fix_point = False

        while (not reach_fix_point) and (step < _max_propagation_step):
            changed = False
            for is_fwd in [True, False]:
                all_nodes = (
                    self._dist_context.serial_ordered_nodes
                    if is_fwd
                    else reversed(self._dist_context.serial_ordered_nodes)
                )
                for node in all_nodes:
                    if node.is_var() and node.var() is not None:
                        tensor_changed = self._update_tensor_node_dims_mapping(
                            node, fwd=is_fwd
                        )
                        if tensor_changed:
                            changed = True
                    if node.is_op() and node.op() is not None:
                        op_changed = self._update_op_node_dims_mapping(
                            node, fwd=is_fwd
                        )
                        if op_changed:
                            changed = True
                graph_changed = self._update_dims_mapping_between_graphs()
                if graph_changed:
                    changed = True

            if changed:
                reach_fix_point = False
            else:
                reach_fix_point = True
            step += 1
        # NOTE: this will be removed after changing the reshard rule

        if step >= _max_propagation_step:
            _logger.debug(
                "Sharding Propagation reach the Max Step and is NOT Converge! The Sharding Propagation Iteration is Terminated."
            )

        self._update_dims_mapping_for_special()

    def _update_process_mesh_by_nearest(self, op_node, nearest_op_node):
        op_dist_attr = self._dist_context.get_dist_attr_for_graph(op_node)
        # Set the process mesh of the op node by its nearest op node
        if not op_dist_attr.is_annotated("process_mesh"):
            process_mesh = op_dist_attr.process_mesh
            nearest_op_dis_attr = self._dist_context.get_dist_attr_for_graph(
                nearest_op_node
            )
            nearest_process_mesh = nearest_op_dis_attr.process_mesh
            compatible_process_mesh = compute_compatible_process_mesh(
                [process_mesh, nearest_process_mesh]
            )
            if (
                compatible_process_mesh is not None
                and process_mesh != compatible_process_mesh
            ):
                op_dist_attr.process_mesh = compatible_process_mesh
        # Skip the process_mesh setting of inputs and outputs of while_op
        if op_dist_attr.op_type == "while":
            return
        # Set the process mesh of the op node's leaf-inputs
        for tensor_node in op_node.inputs:
            if tensor_node.is_var() and tensor_node.var() is not None:
                tensor_dist_attr = (
                    self._dist_context.get_tensor_dist_attr_for_graph(
                        tensor_node
                    )
                )
                if tensor_dist_attr.is_annotated("process_mesh"):
                    continue
                # Skip the non-leaf var node
                if len(tensor_node.inputs) != 0:
                    continue
                compatible_process_mesh = compute_compatible_process_mesh(
                    [tensor_dist_attr.process_mesh, op_dist_attr.process_mesh]
                )
                if (
                    compatible_process_mesh is not None
                    and tensor_dist_attr.process_mesh != compatible_process_mesh
                ):
                    tensor_dist_attr.process_mesh = compatible_process_mesh
                # Set the process mesh of the op node's outputs
        for tensor_node in op_node.outputs:
            if tensor_node.is_var() and tensor_node.var() is not None:
                tensor_dist_attr = (
                    self._dist_context.get_tensor_dist_attr_for_graph(
                        tensor_node
                    )
                )
                if tensor_dist_attr.is_annotated("process_mesh"):
                    continue
                compatible_process_mesh = compute_compatible_process_mesh(
                    [tensor_dist_attr.process_mesh, op_dist_attr.process_mesh]
                )
                if (
                    compatible_process_mesh is not None
                    and tensor_dist_attr.process_mesh != compatible_process_mesh
                ):
                    tensor_dist_attr.process_mesh = compatible_process_mesh

    def _update_process_mesh_for_specials(self):
        def _find_nearest_tensor_node_before(nodes, idx, var_name):
            for node in reversed(nodes[:idx]):
                if (
                    node.is_var()
                    and node.var() is not None
                    and node.var().name() == var_name
                ):
                    return node

        def _find_nearest_tensor_node_after(nodes, idx, var_name):
            for node in nodes[idx + 1 :]:
                if (
                    node.is_var()
                    and node.var() is not None
                    and node.var().name() == var_name
                ):
                    return node

        def _find_nodes_related_to_cond(source_node):
            related_nodes = []
            visited = set()
            frontier = []
            frontier.append(source_node)
            # BFS
            while len(frontier) != 0:
                cur = frontier[0]
                frontier = frontier[1:]
                if _node_id(cur) in visited:
                    continue
                # TODO: need more restrictions
                neighbors = cur.inputs + cur.outputs
                for node in neighbors:
                    if node.is_var() and node.var() is not None:
                        if (
                            node.var().type() != core.VarDesc.VarType.READER
                            and len(node.var().shape()) == 1
                        ):
                            frontier.append(node)
                            related_nodes.append(node)
                    if node.is_op() and node.op() is not None:
                        flag = True
                        if (
                            node.op().type() == "create_py_reader"
                            or node.op().type() == "create_double_buffer_reader"
                            or node.op().type() == "read"
                        ):
                            flag = False
                        for tensor_node in node.inputs:
                            if (
                                tensor_node.is_var()
                                and tensor_node.var() is not None
                            ):
                                if (
                                    tensor_node.var().type()
                                    in __no_shape_var_type__
                                    or len(tensor_node.var().shape()) != 1
                                ):
                                    flag = False
                                    break
                        for tensor_node in node.outputs:
                            if (
                                tensor_node.is_var()
                                and tensor_node.var() is not None
                            ):
                                if (
                                    tensor_node.var().type()
                                    in __no_shape_var_type__
                                    or len(tensor_node.var().shape()) != 1
                                ):
                                    flag = False
                                    break
                        if flag:
                            frontier.append(node)
                            related_nodes.append(node)
                visited.add(_node_id(cur))
            return related_nodes

        def _make_dims_mapping_replicate(dist_attr):
            if isinstance(dist_attr, TensorDistAttr):
                for i, _ in enumerate(dist_attr.dims_mapping):
                    dist_attr.dims_mapping[i] = -1
            if isinstance(dist_attr, OperatorDistAttr):
                for arg_name in dist_attr.inputs_dist_attrs.keys():
                    new_dims_mapping = []
                    dims_mapping = dist_attr.get_input_dims_mapping(arg_name)
                    for _ in dims_mapping:
                        new_dims_mapping.append(-1)
                    dist_attr.set_input_dims_mapping(arg_name, new_dims_mapping)
                for arg_name in dist_attr.outputs_dist_attrs.keys():
                    new_dims_mapping = []
                    dims_mapping = dist_attr.get_output_dims_mapping(arg_name)
                    for _ in dims_mapping:
                        new_dims_mapping.append(-1)
                    dist_attr.set_output_dims_mapping(
                        arg_name, new_dims_mapping
                    )

        # Amend the process meshes related to while_op
        for while_op_node, while_op_node_idx in self._while_op_nodes.values():
            sub_graph_id = while_op_node.op()._block_attr_id("sub_block")
            sub_graph = self._dist_context.serial_graph.get_sub_graph(
                sub_graph_id
            )
            sub_graph_nodes = list(sub_graph.all_nodes())
            while_dist_op = self._dist_context.get_dist_op_for_graph(
                while_op_node
            )
            while_op_dist_attr = while_dist_op.dist_attr

            # Step 1: set the process mesh of while_op to the merged process mesh of its subblock
            merged_process_mesh = while_op_dist_attr.process_mesh
            for node in sub_graph_nodes:
                if (node.is_var() and node.var() is not None) or (
                    node.is_op() and node.op() is not None
                ):
                    dist_attr = self._dist_context.get_dist_attr_for_graph(node)
                    merged_process_mesh = merge_process_mesh_two(
                        merged_process_mesh, dist_attr.process_mesh
                    )
            while_op_dist_attr.process_mesh = merged_process_mesh
            _make_dims_mapping_replicate(while_op_dist_attr)

            # Step 2: set the related nodes of while_op to the process mesh of while_op
            # Step 2.1: Find related nodes of cond var the graph of while_op
            cond_tensor_related_nodes = []
            cond_tensor_name = while_op_node.op().input("Condition")[0]
            cond_tensor_node = None
            for node in while_op_node.inputs:
                if (
                    node.is_var()
                    and node.var() is not None
                    and node.var().name() == cond_tensor_name
                ):
                    cond_tensor_node = node
                    cond_tensor_related_nodes.append(cond_tensor_node)
                    break

            cond_tensor_related_nodes.extend(
                _find_nodes_related_to_cond(cond_tensor_node)
            )

            # Step 2.2: Find related nodes of cond var in the subgraph of while_op
            cond_tensor_node = None
            for node in reversed(sub_graph_nodes):
                if (
                    node.is_var()
                    and node.var() is not None
                    and node.var().name() == cond_tensor_name
                    and len(node.outputs) == 0
                ):
                    cond_tensor_node = node
                    break

            cond_tensor_related_nodes.extend(
                _find_nodes_related_to_cond(cond_tensor_node)
            )
            # Step 2.3: Add the StepScopes output of while_op
            stepscopes_tensor_name = while_op_node.op().output("StepScopes")[0]
            stepscopes_tensor_node = None
            for output_node in while_op_node.outputs:
                if (
                    output_node.is_var()
                    and output_node.var() is not None
                    and output_node.var().name() == stepscopes_tensor_name
                ):
                    stepscopes_tensor_node = output_node
            cond_tensor_related_nodes.append(stepscopes_tensor_node)
            # Step 2.4: Set the process meshes of all nodes related to cond var to the process mesh of while op
            for node in cond_tensor_related_nodes:
                tensor_dist_attr = self._dist_context.get_dist_attr_for_graph(
                    node
                )
                tensor_dist_attr.process_mesh = merged_process_mesh
                _make_dims_mapping_replicate(tensor_dist_attr)

            # Step 3: set the process meshes of the inputs in while_op to the process meshes of the outside input nodes
            while_op_inputs_dist_attrs = while_op_dist_attr.inputs_dist_attrs
            for (
                tensor_name,
                tensor_dist_attr,
            ) in while_op_inputs_dist_attrs.items():
                nearest_tensor_node = _find_nearest_tensor_node_before(
                    self._dist_context.serial_ordered_nodes,
                    while_op_node_idx,
                    tensor_name,
                )
                nearest_tensor_dist_attr = (
                    self._dist_context.get_dist_attr_for_graph(
                        nearest_tensor_node
                    )
                )
                tensor_dist_attr.process_mesh = (
                    nearest_tensor_dist_attr.process_mesh
                )
                for node in while_op_node.inputs:
                    if node.var().name() == tensor_name:
                        node_dist_attr = (
                            self._dist_context.get_dist_attr_for_graph(node)
                        )
                        node_dist_attr.process_mesh = (
                            nearest_tensor_dist_attr.process_mesh
                        )

            # Step 4: set the process meshes of the outputs in while_op to the process meshes of the outside output nodes
            while_op_outputs_dist_attrs = while_op_dist_attr.outputs_dist_attrs
            for (
                tensor_name,
                tensor_dist_attr,
            ) in while_op_outputs_dist_attrs.items():
                nearest_tensor_node = _find_nearest_tensor_node_before(
                    self._dist_context.serial_ordered_nodes,
                    while_op_node_idx,
                    tensor_name,
                )
                if nearest_tensor_node is None:
                    nearest_tensor_node = _find_nearest_tensor_node_after(
                        self._dist_context.serial_ordered_nodes,
                        while_op_node_idx,
                        tensor_name,
                    )
                nearest_tensor_dist_attr = (
                    self._dist_context.get_dist_attr_for_graph(
                        nearest_tensor_node
                    )
                )
                tensor_dist_attr.process_mesh = (
                    nearest_tensor_dist_attr.process_mesh
                )
                for node in while_op_node.outputs:
                    if node.var().name() == tensor_name:
                        node_dist_attr = (
                            self._dist_context.get_dist_attr_for_graph(node)
                        )
                        node_dist_attr.process_mesh = (
                            nearest_tensor_dist_attr.process_mesh
                        )

        # Amend the process meshes related to array
        for array_node_list in self._array_nodes.values():
            merged_process_mesh = None
            for array_node in array_node_list:
                dist_attr = self._dist_context.get_dist_attr_for_graph(
                    array_node
                )
                merged_process_mesh = merge_process_mesh_two(
                    merged_process_mesh, dist_attr.process_mesh
                )
            for array_node in array_node_list:
                dist_attr = self._dist_context.get_dist_attr_for_graph(
                    array_node
                )
                dist_attr.process_mesh = merged_process_mesh
                _make_dims_mapping_replicate(dist_attr)

    def _update_process_mesh_between_graphs(self):
        for parent_node, child_node in self._node_pairs_between_graphs:
            parent_node_dist_attr = self._dist_context.get_dist_attr_for_graph(
                parent_node
            )
            child_node_dist_attr = self._dist_context.get_dist_attr_for_graph(
                child_node
            )
            parent_node_dist_attr.process_mesh = (
                child_node_dist_attr.process_mesh
            )
            compatible_process_mesh = compute_compatible_process_mesh(
                [
                    parent_node_dist_attr.process_mesh,
                    child_node_dist_attr.process_mesh,
                ]
            )
            if (
                compatible_process_mesh is not None
                and parent_node_dist_attr.process_mesh
                != compatible_process_mesh
            ):
                parent_node_dist_attr.process_mesh = compatible_process_mesh
            if (
                compatible_process_mesh is not None
                and child_node_dist_attr.process_mesh != compatible_process_mesh
            ):
                child_node_dist_attr.process_mesh = compatible_process_mesh

    def _update_process_mesh(self):
        ordered_op_nodes = self._dist_context._serial_ordered_op_nodes

        # Step 1: Set the annotated process meshes from tensors to the first ops using them
        ordered_tensor_nodes = self._dist_context._serial_ordered_tensor_nodes
        for tensor_node in ordered_tensor_nodes:
            tensor_dist_attr = (
                self._dist_context.get_tensor_dist_attr_for_graph(tensor_node)
            )
            if not tensor_dist_attr.is_annotated("process_mesh"):
                continue
            first_op_node = None
            for op_node in ordered_op_nodes:
                # TODO: Need a better rule for the control flow ops.
                # For now, do not set the process mesh of while_op from its inputs
                if op_node.op().type() == "while":
                    continue
                for input_tensor_node in op_node.inputs:
                    if _node_id(tensor_node) == _node_id(input_tensor_node):
                        first_op_node = op_node
                        break
                if first_op_node is not None:
                    break
            if first_op_node is None:
                continue
            op_dist_attr = self._dist_context.get_dist_attr_for_graph(
                first_op_node
            )
            if op_dist_attr is not None and not op_dist_attr.is_annotated(
                "process_mesh"
            ):
                compatible_process_mesh = compute_compatible_process_mesh(
                    [tensor_dist_attr.process_mesh, op_dist_attr.process_mesh]
                )
                if (
                    compatible_process_mesh is not None
                    and op_dist_attr.process_mesh != compatible_process_mesh
                ):
                    op_dist_attr.process_mesh = compatible_process_mesh

        # Step 2: set the process meshes of ops with the nearest op before them
        # Step 2.1: find the first op node which has the process mesh
        idx_of_first_op_node_has_process_mesh = -1
        for idx, op_node in enumerate(ordered_op_nodes):
            op_dist_attr = self._dist_context.get_dist_attr_for_graph(op_node)
            if (
                op_dist_attr.process_mesh is not None
                and idx_of_first_op_node_has_process_mesh == -1
            ):
                idx_of_first_op_node_has_process_mesh = idx
                # Reuse the following method to set the related tensors for same op node
                self._update_process_mesh_by_nearest(op_node, op_node)
        # Step 2.2: set the process meshes of ops by the nearest op node after the first op node
        if idx_of_first_op_node_has_process_mesh + 1 > len(ordered_op_nodes):
            return None
        for idx, op_node in enumerate(
            ordered_op_nodes[idx_of_first_op_node_has_process_mesh + 1 :]
        ):
            original_idx = idx_of_first_op_node_has_process_mesh + idx + 1
            nearest_op_node = ordered_op_nodes[original_idx - 1]
            nearest_op_dist_attr = self._dist_context.get_dist_attr_for_graph(
                nearest_op_node
            )
            op_dist_attr = self._dist_context.get_dist_attr_for_graph(op_node)
            assert nearest_op_dist_attr.process_mesh is not None
            self._update_process_mesh_by_nearest(op_node, nearest_op_node)
        # Step 2.3: set the process meshes of ops by the nearest op node before the first op node
        nearest_op_node = ordered_op_nodes[
            idx_of_first_op_node_has_process_mesh
        ]
        for op_node in ordered_op_nodes[:idx_of_first_op_node_has_process_mesh]:
            self._update_process_mesh_by_nearest(op_node, nearest_op_node)

        # Step 3: adjust the process meshes for special ops
        self._update_process_mesh_for_specials()

        # Step 4: adjust the process meshes between graphs
        self._update_process_mesh_between_graphs()

    def _prepare(self):
        if self._has_prepared:
            return
        self._while_op_nodes = {}
        self._array_nodes = {}
        self._node_pairs_between_graphs = []
        all_nodes = self._dist_context.serial_ordered_nodes
        for idx, node in enumerate(all_nodes):
            if node.is_op():
                if node.op().type() == "while":
                    self._while_op_nodes[_node_id(node)] = (node, idx)
                if node.op().type() == "read_from_array":
                    array_var_name = node.op().input("X")[0]
                    if self._array_nodes.get(array_var_name, None) is None:
                        self._array_nodes[array_var_name] = []
                    self._array_nodes[array_var_name].append(node)
                    # Add the array input node
                    self._array_nodes[array_var_name].append(node.inputs[0])
                if node.op().type() == "write_to_array":
                    array_var_name = node.op().output("Out")[0]
                    if self._array_nodes.get(array_var_name, None) is None:
                        self._array_nodes[array_var_name] = []
                    self._array_nodes[array_var_name].append(node)
                    self._array_nodes[array_var_name].append(node.outputs[0])
            if node.is_var() and node.var() is not None:
                if node.node.graph_id() != 0:
                    parent_nodes = (
                        self._dist_context._tensor_nodes_with_same_name[
                            node.node.graph_id() - 1
                        ].get(node.var().name(), None)
                    )
                    if parent_nodes is not None:
                        sorted_parent_nodes = sorted(
                            parent_nodes, key=lambda x: x[0]
                        )
                        for _, parent_node in sorted_parent_nodes:
                            self._node_pairs_between_graphs.append(
                                (parent_node, node)
                            )

        self._has_prepared = True

    def complete_forward_annotation(self, serial_main_program=None):
        """Complete annotation for the partial annotated serial_main_program.
        Arguments:
            serial_main_program: partial annotated serial_main_program.
        Returns:
            serial_main_program: completed annotated serial_main_program.
        """

        if serial_main_program is None:
            serial_main_program = self._dist_context.serial_main_program
        else:
            self._dist_context._serial_main_program = serial_main_program

        tensor_names, ops = self._get_tensor_names_and_ops_with_global_mesh(
            serial_main_program
        )
        if not is_naive_data_parallel(self._dist_context):
            self._dist_context.initialize(with_graph=True)
            self._prepare()
            self._update_process_mesh()
            self._update_dims_mapping()
            # Copy the corresponding distributed attribute from graph to serial_main_program
            self._dist_context.copy_dist_attr_from_graph_to_program()
        else:
            _logger.info("Default distributed attributed will be set.")
            self._dist_context.initialize(with_graph=False)
            # A fast and special completion for data parallel
            self._update_dist_attr_for_dp()

        self._complete_with_global_mesh(serial_main_program, tensor_names, ops)
        # NOTE:[HighOrderGrad] update vars and ops distributed attribute in high order gradient
        self._complete_high_order_grad_annotation(serial_main_program)
        self._complete_chunk_id(serial_main_program)
        # Do the validation check and amend some completion
        self._dist_context.amend_dist_attr_for_program()
        self._dist_context.validate_dist_attr_for_program()
        return serial_main_program

    def _get_tensor_names_and_ops_with_global_mesh(self, serial_main_program):
        if (
            not self._dist_context.strategy
            or not self._dist_context.strategy.pipeline.enable
        ):
            return [], []

        # step1: get tensor annotated with global mesh
        global_mesh = paddle.distributed.auto_parallel.get_mesh()
        if global_mesh is None:
            _logger.warning(
                "global_mesh is not set, tensor annotation with global mesh may be not work, please use paddle.distributed.auto_parallel.set_mesh(mesh) firstly."
            )
            return [], []
        global_mesh_process_ids = global_mesh._process_ids
        tensor_names_with_global_mesh = []
        block = serial_main_program.global_block()
        for var in block.vars.values():
            dist_var = self._dist_context.get_dist_tensor_for_program(var)
            mesh = dist_var.dist_attr.process_mesh
            if mesh is not None and sorted(mesh.process_ids) == sorted(
                global_mesh_process_ids
            ):
                tensor_names_with_global_mesh.append(var.name)

        # if no one tensor has global mesh, do nothing
        if len(tensor_names_with_global_mesh) == 0:
            return [], []

        # step2: get all tensors and ops should annotated with global mesh
        tensor_name_to_op = {}
        ops = block.ops
        for op in ops:
            output_tensor_names = op.output_arg_names
            for tensor_name in output_tensor_names:
                tensor_name_to_op[tensor_name] = op

        ops_with_global_mesh = []
        has_visited = set()
        tensor_name_queue = queue.Queue()
        for tensor_name in tensor_names_with_global_mesh:
            tensor_name_queue.put(tensor_name)
        tensor_names_with_global_mesh.clear()
        # BFS to find all tensors and ops should annotated with global mesh
        while not tensor_name_queue.empty():
            tensor_name = tensor_name_queue.get()
            if tensor_name in has_visited:
                continue

            has_visited.add(tensor_name)
            tensor_names_with_global_mesh.append(tensor_name)
            op = tensor_name_to_op[tensor_name]
            ops_with_global_mesh.append(op)
            input_arg_names = op.input_arg_names
            for input_name in input_arg_names:
                tensor_name_queue.put(input_name)
        return tensor_names_with_global_mesh, ops_with_global_mesh

    def _complete_with_global_mesh(
        self, serial_main_program, tensor_names, ops
    ):
        if len(tensor_names) == 0:
            return
        # step1: get global mesh
        block = serial_main_program.global_block()
        # tensor_names[0] is a tensor annotated with global mesh
        tensor = block._var_recursive(tensor_names[0])
        dist_tensor = self._dist_context.get_dist_tensor_for_program(tensor)
        global_mesh = dist_tensor.dist_attr.process_mesh

        # step2: set the global mesh to ops and tensors
        for op in ops:
            dist_op = self._dist_context.get_dist_op_for_program(op)
            dist_op.dist_attr.process_mesh = global_mesh
        for tensor_name in tensor_names:
            tensor = block._var_recursive(tensor_name)
            dist_tensor = self._dist_context.get_dist_tensor_for_program(tensor)
            dist_tensor.dist_attr.process_mesh = global_mesh

    def _complete_chunk_id(self, serial_main_program):
        def set_chunk_id(block, op, chunk_id, var_to_chunk_id):
            dist_op = self._dist_context.get_dist_op_for_program(op)
            dist_op.dist_attr.chunk_id = chunk_id
            for name in op.input_arg_names + op.output_arg_names:
                if "lod_tensor_blocking_queue" in name:
                    continue
                if name not in var_to_chunk_id:
                    var = block._find_var_recursive(name)
                    dist_tensor = (
                        self._dist_context.get_dist_tensor_for_program(var)
                    )
                    if (
                        dist_op.dist_attr.process_mesh
                        == dist_tensor.dist_attr.process_mesh
                    ):
                        dist_tensor.dist_attr.chunk_id = chunk_id
                        var_to_chunk_id[var.name] = chunk_id

        def set_process_mesh(block, op, process_mesh, var_to_process_mesh):
            dist_op = self._dist_context.get_dist_op_for_program(op)
            for name in op.input_arg_names:
                if name not in var_to_process_mesh:
                    var = block._find_var_recursive(name)
                    dist_tensor = (
                        self._dist_context.get_dist_tensor_for_program(var)
                    )
                    if (
                        dist_op.dist_attr.process_mesh
                        == dist_tensor.dist_attr.process_mesh
                    ):
                        dist_tensor.dist_attr.process_mesh = process_mesh
                        var_to_process_mesh[var.name] = process_mesh
            for name in op.output_arg_names:
                if name not in var_to_process_mesh:
                    var = block._find_var_recursive(name)
                    dist_tensor = (
                        self._dist_context.get_dist_tensor_for_program(var)
                    )
                    dist_tensor.dist_attr.process_mesh = process_mesh
                    var_to_process_mesh[var.name] = process_mesh
            dist_op.dist_attr.process_mesh = process_mesh

        if (
            not self._dist_context.strategy
            or not self._dist_context.strategy.pipeline.enable
        ):
            return

        pp_degree, sub_process_meshes = get_pp_degree(self._dist_context)
        vpp_degree = self._dist_context.strategy.pipeline.vpp_degree
        seg_method = self._dist_context.strategy.pipeline.vpp_seg_method
        schedule_mode = self._dist_context.strategy.pipeline.schedule_mode

        if pp_degree < 2 and vpp_degree > 1:
            raise ValueError(
                "VPP schedule mode only can be set in pipeline mode."
            )
        if vpp_degree > 1 and (
            not seg_method or schedule_mode not in ["VPP", "ZBVPP"]
        ):
            raise ValueError(
                "Please set right schedule_mode and vpp_seg_method for VPP and ZBVPP."
            )
        if vpp_degree < 2:
            return

        block = serial_main_program.global_block()
        ops = block.ops

        # Step1: search seg_method in op's struct_name
        # 1. get op_idx of each segment
        # 2. get process_mesh or each segment
        seg_op_deps = collections.OrderedDict()  # struct_name -> [idx]
        seg_op_mesh = collections.OrderedDict()  # struct_name -> process_mesh
        regex = re.compile(seg_method, re.IGNORECASE)

        start_op_index = 0
        for i, op in enumerate(ops):
            m = regex.search(op.struct_name)
            if m:
                start_op_index = i
                break

        total_op_num = len(ops)
        end_op_index = total_op_num - 1
        for i in reversed(range(total_op_num)):
            m = regex.search(ops[i].struct_name)
            if m:
                end_op_index = i
                break

        # all ops between start_op_index and end_op_index should not be ignored
        for i in range(start_op_index, end_op_index + 1):
            struct_name = ops[i].struct_name
            m = regex.search(struct_name)
            if not m:
                # only assign op created by reshard is allowed
                if (
                    ops[i].type == "assign"
                    and "reshard_api" in ops[i].output_arg_names[0]
                ):
                    # this assign op belongs to next segment
                    for j in range(i + 1, total_op_num):
                        m = regex.search(ops[j].struct_name)
                        if m:
                            break
                    assert m
                    struct_name = ops[j].struct_name
                else:
                    raise ValueError(
                        f"The op {ops[i]} should only be created by reshard"
                    )

            struct_name = struct_name[m.start(0) :].split("/")[0]
            dist_op = self._dist_context.get_dist_op_for_program(ops[i])
            if struct_name not in seg_op_deps:
                seg_op_deps[struct_name] = [i]
                seg_op_mesh[struct_name] = dist_op.dist_attr.process_mesh
            else:
                assert seg_op_deps[struct_name][-1] + 1 == i, (
                    "The segment's ops should be continuous."
                )
                pre_mesh = seg_op_mesh[struct_name]
                assert pre_mesh == dist_op.dist_attr.process_mesh, (
                    "The segment's ops should have same process_mesh."
                )
                seg_op_deps[struct_name].extend([i])

        num_chunks = pp_degree * vpp_degree
        assert len(seg_op_deps) % num_chunks == 0, (
            f"The number of layers[{seg_method}] ({len(seg_op_deps)}) should be divided by part number ({num_chunks})."
        )

        # Step2: analysis whether the pp_stage is non-decreasing among segments
        # 1. if non_decreasing is True, the ops' process_mesh will be changed by vpp strategy
        # 2. if non_decreasing is False, the ops's process_mesh will not be changed.
        non_decreasing = True
        seg_pp_stages = [-1]
        for seg_pm in seg_op_mesh.values():
            assert seg_pm in sub_process_meshes
            pp_stage = sub_process_meshes.index(seg_pm)
            if seg_pp_stages[-1] > pp_stage:
                non_decreasing = False
                break
            seg_pp_stages.append(pp_stage)

        if not non_decreasing:
            _logger.info("Cannot Use Auto VPP")
        else:
            _logger.info("Using Auto VPP")

        # Step3: Get op index boundary, pp_stage, chunk_id, struct_names of each segment
        seg_pp_stages = []
        seg_pp_stage = list(range(pp_degree))
        for _ in range(vpp_degree):
            seg_pp_stages.extend(seg_pp_stage)
            if schedule_mode == "ZBVPP":
                seg_pp_stage.reverse()

        seg_chunk_ids = [i // pp_degree for i in range(num_chunks)]
        part_size = len(seg_op_deps) // num_chunks
        segment_struct_names = []
        segment_parts = [0] * (num_chunks + 1)
        memory_counter, seg_idx = 0, 1
        struct_name = []
        for name, idxs in seg_op_deps.items():
            struct_name.append(name)
            memory_counter += 1
            if memory_counter == part_size:
                segment_parts[seg_idx] = idxs[-1] + 1
                memory_counter, seg_idx = 0, seg_idx + 1
                segment_struct_names.append(struct_name)
                struct_name = []
            segment_parts[num_chunks] = len(ops)

        # Step4: set right chunk_id and process_mesh for each op and var in each segment
        var_to_chunk_id = {}
        var_to_process_mesh = {}
        for seg_id in range(len(segment_parts) - 1):
            start_idx = segment_parts[seg_id]
            end_idx = segment_parts[seg_id + 1]
            pp_stage = seg_pp_stages[seg_id]
            chunk_id = seg_chunk_ids[seg_id]
            process_mesh = sub_process_meshes[pp_stage]
            struct_names = segment_struct_names[seg_id]
            seg_op_idx = []
            for name in struct_names:
                seg_op_idx.extend(seg_op_deps[name])

            _logger.info(
                f"stage=[{pp_stage}], chunk_id=[{chunk_id}], layer_name=[{struct_names}]"
            )
            _logger.info(
                f"start op: [{ops[start_idx].type}]: [{ops[start_idx].input_arg_names}] [{ops[start_idx].output_arg_names}]"
            )
            _logger.info(
                f"end op: [{ops[end_idx - 1].type}]: [{ops[end_idx - 1].input_arg_names}] [{ops[end_idx - 1].output_arg_names}]"
            )

            for idx in range(start_idx, end_idx):
                op = ops[idx]
                if op.has_attr("sub_block"):
                    block_id = op.attr('sub_block').id
                    sub_block = serial_main_program.blocks[block_id]
                    if non_decreasing and idx in seg_op_idx:
                        set_process_mesh(
                            block, op, process_mesh, var_to_process_mesh
                        )
                    set_chunk_id(block, op, chunk_id, var_to_chunk_id)

                    for sub_op in sub_block.ops:
                        if non_decreasing and idx in seg_op_idx:
                            set_process_mesh(
                                sub_block,
                                sub_op,
                                process_mesh,
                                var_to_process_mesh,
                            )
                        set_chunk_id(
                            sub_block, sub_op, chunk_id, var_to_chunk_id
                        )
                else:
                    if non_decreasing and idx in seg_op_idx:
                        set_process_mesh(
                            block, op, process_mesh, var_to_process_mesh
                        )
                    set_chunk_id(block, op, chunk_id, var_to_chunk_id)

        # Step5: set right chunk_id and process_mesh for loss op
        # Note(sonder): for zbvpp schedule mode, the loss will be calculated in the first stage when vpp_degree is even
        if schedule_mode == "ZBVPP" and vpp_degree % 2 == 0:
            for i in range(end_op_index, total_op_num):
                set_chunk_id(block, ops[i], vpp_degree - 1, var_to_chunk_id)
                set_process_mesh(
                    block, ops[i], sub_process_meshes[0], var_to_process_mesh
                )

    def _update_dist_attr_for_dp(self):
        # TODO: we must ensure the world process group contains all ranks
        ranks = get_world_process_group().ranks
        process_mesh = ProcessMesh(ranks)

        dist_tensors = self._dist_context._dist_tensors_for_program
        for dist_tensor in dist_tensors.values():
            dist_tensor.dist_attr.process_mesh = process_mesh

        dist_ops = self._dist_context._dist_ops_for_program
        for dist_op in dist_ops.values():
            serial_op = dist_op.serial_op
            op_dist_attr = dist_op.dist_attr
            op_dist_attr.process_mesh = process_mesh
            original_op_dist_attr = copy.deepcopy(op_dist_attr)

            if serial_op.type == "create_py_reader":
                continue

            for arg_name in serial_op.input_arg_names:
                serial_tensor = dist_op.get_serial_input(arg_name)
                if not serial_tensor.is_parameter:
                    dist_tensor = (
                        self._dist_context.get_dist_tensor_for_program(
                            serial_tensor
                        )
                    )
                    op_dist_attr = dist_op.dist_attr
                    op_dist_attr.process_mesh = (
                        dist_tensor.dist_attr.process_mesh
                    )
                    op_dist_attr.set_input_dims_mapping(
                        arg_name, dist_tensor.dist_attr.dims_mapping
                    )

            op_dist_impls = find_compatible_distributed_operator_impls(
                dist_op, fwd=True
            )
            if op_dist_impls is not None:
                not_compatible = True
                backup_op_dist_attr = copy.deepcopy(op_dist_attr)
                for op_dist_impl in op_dist_impls:
                    op_dist_impl.update_dims_mapping(dist_op)
                    if (
                        op_dist_impl.is_auto_compatible(dist_op)
                        and dist_op.validate_dist_attr()
                    ):
                        op_dist_attr.impl_type = op_dist_impl.type
                        op_dist_attr.impl_idx = op_dist_impl.idx
                        not_compatible = False
                        break
                    else:
                        dist_op.dist_attr = backup_op_dist_attr
                if not_compatible:
                    dist_op.dist_attr = original_op_dist_attr
            else:
                dist_op.dist_attr = original_op_dist_attr

            for arg_name in serial_op.output_arg_names:
                op_dist_attr = dist_op.dist_attr
                serial_tensor = dist_op.get_serial_output(arg_name)
                if serial_op.type in ["fill_constant"]:
                    old_dims_mapping = op_dist_attr.get_output_dims_mapping(
                        arg_name
                    )
                    if len(old_dims_mapping) > 0:
                        new_dims_mapping = [0] + [
                            -1 for _ in range(len(old_dims_mapping) - 1)
                        ]
                        op_dist_attr.set_output_dims_mapping(
                            arg_name, new_dims_mapping
                        )
                dist_tensor = self._dist_context.get_dist_tensor_for_program(
                    serial_tensor
                )
                dist_tensor.dist_attr.dims_mapping = (
                    op_dist_attr.get_output_dims_mapping(arg_name)
                )

    def _complete_tensor_dist_attr_by_op(self, serial_main_program=None):
        if serial_main_program is None:
            serial_main_program = self._dist_context.serial_main_program
        else:
            self._dist_context._serial_main_program = serial_main_program

        self._dist_context.initialize()

        self._prepare()

        has_set_dist_attr = set()

        all_nodes = self._dist_context.serial_ordered_nodes
        for node in all_nodes:
            if node.is_op():
                if node.op().type() in ["while"]:
                    continue
                dist_op = self._dist_context.get_dist_op_for_graph(node)
                op_dist_attr = dist_op.dist_attr
                for tensor_node in node.inputs:
                    if tensor_node.is_var() and tensor_node.var() is not None:
                        # Skip the non-leaf var node
                        if len(tensor_node.inputs) != 0:
                            continue
                        tensor_desc = tensor_node.var()
                        tensor_name = tensor_desc.name()
                        tensor = dist_op.get_serial_input(tensor_name)
                        # Use the first op to set the tensor dist attr
                        if tensor_name in has_set_dist_attr:
                            continue
                        tensor_dist_attr = (
                            self._dist_context.get_tensor_dist_attr_for_graph(
                                tensor_node
                            )
                        )
                        tensor_dist_attr.process_mesh = (
                            op_dist_attr.process_mesh
                        )
                        tensor_dist_attr.dims_mapping = (
                            op_dist_attr.get_input_dims_mapping(tensor_name)
                            if tensor.is_parameter
                            else [-1 for i in tensor_desc.shape()]
                        )
                        has_set_dist_attr.add(tensor_name)
                for tensor_node in node.outputs:
                    if tensor_node.is_var() and tensor_node.var() is not None:
                        tensor_name = tensor_node.var().name()
                        if tensor_name in has_set_dist_attr:
                            continue
                        tensor_dist_attr = (
                            self._dist_context.get_tensor_dist_attr_for_graph(
                                tensor_node
                            )
                        )
                        tensor_dist_attr.process_mesh = (
                            op_dist_attr.process_mesh
                        )
                        tensor_dist_attr.dims_mapping = (
                            op_dist_attr.get_output_dims_mapping(tensor_name)
                        )
                        has_set_dist_attr.add(tensor_name)

        self._update_process_mesh_for_specials()

        self._update_process_mesh_between_graphs()

        self._update_dims_mapping_for_special()

        self._update_dims_mapping_between_graphs()

        # Copy the corresponding distributed attribute from graph to serial_main_program
        self._dist_context.copy_dist_attr_from_graph_to_program()

        # Do the validation check and amend some completion
        self._dist_context.amend_dist_attr_for_program()

        self._dist_context.validate_dist_attr_for_program()

    def _complete_high_order_grad_annotation(self, serial_main_program=None):
        """
        NOTE:
            [HighOrderGrad] Complete the annotation of vars and ops only for high order gradient.
            This function is temporary to support high order gradient, and will be removed in the future.
        """

        if serial_main_program is None:
            serial_main_program = self._dist_context.serial_main_program
        else:
            self._dist_context._serial_main_program = serial_main_program

        def _is_grad_var_name(name):
            if "@GRAD" in name:
                return True
            return False

        def _get_op_by_id(ops, id):
            for op in ops:
                if op.desc.original_id() == id:
                    return op
            return None

        ops = list(serial_main_program.global_block().ops)
        vars = serial_main_program.global_block().vars
        dist_op_context = self._dist_context.dist_op_context
        grad_var_to_var = dist_op_context.grad_var_to_var

        if len(grad_var_to_var) < 2:
            return

        appended_grad_times = 0
        for idx in range(0, len(ops)):
            op = ops[idx]
            if int(op.attr('op_role')) == int(
                core.op_proto_and_checker_maker.OpRole.Forward
            ):
                continue

            if int(op.attr('op_role')) == int(
                core.op_proto_and_checker_maker.OpRole.Backward
            ) and int(ops[idx - 1].attr('op_role')) == int(
                core.op_proto_and_checker_maker.OpRole.Forward
            ):
                appended_grad_times += 1

            if int(op.attr('op_role')) == int(
                int(core.op_proto_and_checker_maker.OpRole.Backward)
                | int(core.op_proto_and_checker_maker.OpRole.Loss)
            ):
                assert op.type == "fill_constant"
                break

            # complete the annotation of grad op (xxx_grad op or sum op)
            # xxx_grad op will have a corresponding forward op in grad_op_id_to_op_id
            grad_op = ops[idx]
            if (
                grad_op.desc.original_id()
                in dist_op_context.grad_op_id_to_op_id
            ):
                # TODO support the case where one forward op corresponding to multiple xxx_grad op
                forward_op = _get_op_by_id(
                    ops,
                    dist_op_context.grad_op_id_to_op_id[
                        grad_op.desc.original_id()
                    ],
                )
                assert forward_op is not None

                fwd_op_dist_attr = (
                    self._dist_context.get_op_dist_attr_for_program(forward_op)
                )
                fwd_op_process_mesh = fwd_op_dist_attr.process_mesh
                grad_op_dist_attr = OperatorDistAttr()
                grad_op_dist_attr.process_mesh = fwd_op_process_mesh

                for input_name in grad_op.input_arg_names:
                    if (
                        input_name not in forward_op.input_arg_names
                        and input_name not in forward_op.output_arg_names
                    ):
                        if input_name in grad_var_to_var[appended_grad_times]:
                            fwd_name = grad_var_to_var[appended_grad_times][
                                input_name
                            ]
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_output_dims_mapping(
                                    fwd_name
                                )
                            )
                        else:
                            input_var = vars[input_name]
                            ref_dims_mapping = self._dist_context.get_tensor_dist_attr_for_program(
                                input_var
                            ).dims_mapping
                    else:
                        if input_name in forward_op.input_arg_names:
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_input_dims_mapping(
                                    input_name
                                )
                            )
                        else:
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_output_dims_mapping(
                                    input_name
                                )
                            )
                    assert ref_dims_mapping is not None, (
                        f"[{input_name}] 's dims mapping is NONE"
                    )
                    grad_op_dist_attr.set_input_dims_mapping(
                        input_name, ref_dims_mapping
                    )

                for output_name in grad_op.output_arg_names:
                    assert output_name in grad_var_to_var[appended_grad_times]
                    fwd_name = grad_var_to_var[appended_grad_times][output_name]
                    ref_dims_mapping = fwd_op_dist_attr.get_input_dims_mapping(
                        fwd_name
                    )
                    # var
                    output_var = vars[output_name]
                    tensor_dist_attr = TensorDistAttr()
                    tensor_dist_attr.dims_mapping = ref_dims_mapping
                    tensor_dist_attr.process_mesh = fwd_op_process_mesh
                    self._dist_context.set_tensor_dist_attr_for_program(
                        output_var, tensor_dist_attr
                    )
                    # op
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_name, ref_dims_mapping
                    )

                self._dist_context.set_op_dist_attr_for_program(
                    grad_op, grad_op_dist_attr
                )

            # grad ops that have not a corresponding mapping in grad_op_id_to_op_id
            else:
                if grad_op.type == 'sum':
                    assert all(map(_is_grad_var_name, grad_op.input_arg_names))
                    output_name = grad_op.output_arg_names[0]
                    assert (
                        output_name in grad_var_to_var[appended_grad_times]
                    ), (
                        f"sum op's output '{output_name}' has no corresponding var"
                    )
                    ref_fwd_var_name = grad_var_to_var[appended_grad_times][
                        output_name
                    ]
                    ref_fwd_var = vars[ref_fwd_var_name]
                    ref_fwd_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            ref_fwd_var
                        )
                    )
                    ref_fwd_dims_mapping = ref_fwd_dist_attr.dims_mapping
                    ref_fwd_process_mesh = ref_fwd_dist_attr.process_mesh
                    # output
                    tensor_dist_attr = TensorDistAttr()
                    tensor_dist_attr.dims_mapping = ref_fwd_dims_mapping
                    tensor_dist_attr.process_mesh = ref_fwd_process_mesh
                    output_var = vars[output_name]
                    self._dist_context.set_tensor_dist_attr_for_program(
                        output_var, tensor_dist_attr
                    )
                    # op
                    grad_op_dist_attr = OperatorDistAttr()
                    grad_op_dist_attr.process_mesh = ref_fwd_process_mesh
                    for var_name in grad_op.input_arg_names:
                        grad_op_dist_attr.set_input_dims_mapping(
                            var_name, ref_fwd_dims_mapping
                        )
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_name, ref_fwd_dims_mapping
                    )

                elif grad_op.type == 'fill_any_like':
                    ref_var_name = grad_op.input_arg_names[0]
                    ref_var = vars[ref_var_name]
                    ref_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            ref_var
                        )
                    )
                    ref_dims_mapping = ref_dist_attr.dims_mapping
                    ref_process_mesh = ref_dist_attr.process_mesh
                    # output
                    tensor_dist_attr = TensorDistAttr()
                    tensor_dist_attr.dims_mapping = ref_dims_mapping
                    tensor_dist_attr.process_mesh = ref_process_mesh
                    output_var_name = grad_op.output_arg_names[0]
                    output_var = vars[output_var_name]
                    self._dist_context.set_tensor_dist_attr_for_program(
                        output_var, tensor_dist_attr
                    )
                    # op
                    grad_op_dist_attr = OperatorDistAttr()
                    grad_op_dist_attr.process_mesh = ref_process_mesh
                    grad_op_dist_attr.set_input_dims_mapping(
                        ref_var_name, ref_dims_mapping
                    )
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_var_name, ref_dims_mapping
                    )

                elif grad_op.type in ['shape', 'fill_constant']:
                    continue

                else:
                    raise ValueError(f"got unexpected op [{grad_op.type}]")

                self._dist_context.set_op_dist_attr_for_program(
                    grad_op, grad_op_dist_attr
                )

    def complete_backward_annotation(self, serial_main_program=None):
        """Complete the annotation of vars and ops in the backward phase for parallel program."""

        if serial_main_program is None:
            serial_main_program = self._dist_context.serial_main_program
        else:
            self._dist_context._serial_main_program = serial_main_program

        def _is_grad_var_name(name):
            if "@GRAD" in name:
                return True
            return False

        def _get_forward_varname_from_grad_varname(grad_var_name):
            assert _is_grad_var_name(grad_var_name), (
                f"[{grad_var_name}] is not a grad var name."
            )
            return grad_var_name[: grad_var_name.find("@GRAD")]

        def _get_op_by_id(ops, id):
            for op in ops:
                if op.desc.original_id() == id:
                    return op
            return None

        def _complete_grad_op_with_forward_op(forward_op, grad_op, vars):
            fwd_op_dist_attr = self._dist_context.get_op_dist_attr_for_program(
                forward_op
            )
            grad_op_dist_attr = OperatorDistAttr()
            ref_process_mesh = fwd_op_dist_attr.process_mesh
            ref_chunk_id = fwd_op_dist_attr.chunk_id

            if grad_op.type == "concat" and forward_op.type == "split":
                split_input_var_name = forward_op.input("X")[0]
                ref_dims_mapping = fwd_op_dist_attr.get_input_dims_mapping(
                    split_input_var_name
                )
                # var
                output_var = vars[grad_op.desc.output('Out')[0]]
                set_var_dist_attr(
                    self._dist_context,
                    output_var,
                    ref_dims_mapping,
                    ref_process_mesh,
                    chunk_id=ref_chunk_id,
                )
                # op
                for input_name in grad_op.input_arg_names:
                    grad_op_dist_attr.set_input_dims_mapping(
                        input_name, ref_dims_mapping
                    )
                grad_op_dist_attr.set_output_dims_mapping(
                    output_var.name, ref_dims_mapping
                )
            else:
                # complete grad_op's input_dist_attrs, no need to complete input_var's tensor_dist_attr
                for input_name in grad_op.input_arg_names:
                    if (
                        input_name not in forward_op.input_arg_names
                        and input_name not in forward_op.output_arg_names
                    ):
                        if input_name in grad_var_to_var:
                            fwd_name = grad_var_to_var[input_name]
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_output_dims_mapping(
                                    fwd_name
                                )
                            )
                        else:
                            input_var = vars[input_name]
                            ref_dims_mapping = self._dist_context.get_tensor_dist_attr_for_program(
                                input_var
                            ).dims_mapping
                    else:
                        if input_name in forward_op.input_arg_names:
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_input_dims_mapping(
                                    input_name
                                )
                            )
                        else:
                            ref_dims_mapping = (
                                fwd_op_dist_attr.get_output_dims_mapping(
                                    input_name
                                )
                            )
                    assert ref_dims_mapping is not None, (
                        f"[{input_name}] 's dims mapping is NONE"
                    )
                    grad_op_dist_attr.set_input_dims_mapping(
                        input_name, ref_dims_mapping
                    )

                # complete grad_op's output_dist_attrs, and output_var's tensor_dist_attr
                for output_name in grad_op.output_arg_names:
                    if output_name == "@EMPTY@":
                        output_var = vars[output_name]
                        ref_dims_mapping = [
                            -1 for _ in range(len(output_var.shape))
                        ]
                        set_var_dist_attr(
                            self._dist_context,
                            output_var,
                            ref_dims_mapping,
                            ref_process_mesh,
                            chunk_id=ref_chunk_id,
                        )
                        grad_op_dist_attr.set_output_dims_mapping(
                            output_name, ref_dims_mapping
                        )
                        continue

                    assert output_name in grad_var_to_var
                    fwd_name = grad_var_to_var[output_name]
                    ref_dims_mapping = fwd_op_dist_attr.get_input_dims_mapping(
                        fwd_name
                    )
                    # var
                    output_var = vars[output_name]
                    set_var_dist_attr(
                        self._dist_context,
                        output_var,
                        ref_dims_mapping,
                        ref_process_mesh,
                        chunk_id=ref_chunk_id,
                    )
                    # op
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_name, ref_dims_mapping
                    )

            grad_op_dist_attr.process_mesh = ref_process_mesh
            grad_op_dist_attr.chunk_id = ref_chunk_id
            grad_op_dist_attr.impl_type = fwd_op_dist_attr.impl_type
            grad_op_dist_attr.impl_idx = fwd_op_dist_attr.impl_idx
            grad_op_dist_attr.chunk_id = fwd_op_dist_attr.chunk_id

            # inference partial backward
            def infer_backward_op_partial_status(
                vars, grad_op, grad_op_dist_attr
            ):
                # NOTE Since we use composite op in static mode which might have implicit Reduction of broadcast axes for calculating parameter's gradient.
                # Those implicit Reduction hinder the Partial inference in a normal way, and we need a special method to handle it.
                param_grads = []
                activation_grad = None
                broadcast_axis_indies = []
                if (
                    grad_op.type == "matmul_v2_grad"
                    and len(grad_op.output("Y@GRAD")) > 0
                ):
                    activation_grad = grad_op.input("Out@GRAD")[0]
                    param_grads.extend(grad_op.output("Y@GRAD"))
                    act_ndim = len(vars[activation_grad].shape)
                    param_ndim = len(vars[grad_op.output("Y@GRAD")[0]].shape)
                    # TODO handle case where trans_x or trans_y is true
                    # NOTE we regard axis m as broadcast axis since it is the contracting axis when calculate param grad.
                    if param_ndim <= 2:
                        if act_ndim > 1:
                            broadcast_axis_indies = list(range(act_ndim - 1))
                    elif act_ndim > param_ndim:
                        broadcast_axis_indies = list(
                            range(act_ndim - param_ndim)
                        )
                elif grad_op.type == "elementwise_add_grad":
                    activation_grad = grad_op.input("Out@GRAD")[0]
                    param_grads.extend(grad_op.output("Y@GRAD"))
                    param_var = grad_op.input("Y")[0]
                    broadcast_axis_indies = list(
                        range(
                            len(vars[activation_grad].shape)
                            - len(vars[param_var].shape)
                        )
                    )
                elif grad_op.type == "layer_norm_grad":
                    activation_grad = grad_op.input("Y@GRAD")[0]
                    param_grads.extend(grad_op.output("Bias@GRAD"))
                    param_grads.extend(grad_op.output("Scale@GRAD"))
                    begin_norm_axis = int(grad_op.attr("begin_norm_axis"))
                    broadcast_axis_indies = list(range(begin_norm_axis))
                elif grad_op.type == "lookup_table_v2_grad":
                    activation_grad = grad_op.input("Out@GRAD")[0]
                    param_grads.extend(grad_op.output("W@GRAD"))
                    broadcast_axis_indies = list(
                        range(len(vars[activation_grad].shape) - 1)
                    )
                else:
                    raise NotImplementedError(
                        f"Backward Partial is not adapted for {grad_op}"
                    )

                # resolute partial
                # NOTE We set the Partial status in op_dist_attr instead tensor_dist_attr
                # since the Partial will be reshard as Replicated immediately after op output in static mode.
                if len(param_grads) > 0:
                    activation_grad_dims_mapping = (
                        grad_op_dist_attr.get_input_dims_mapping(
                            activation_grad
                        )
                    )
                    for axis in broadcast_axis_indies:
                        if activation_grad_dims_mapping[axis] != -1:
                            partial_dim = activation_grad_dims_mapping[axis]
                            for p_grad_name in param_grads:
                                p_grad_dist_attr = (
                                    grad_op_dist_attr.get_output_dist_attr(
                                        p_grad_name
                                    )
                                )
                                p_grad_dist_attr._set_partial_dims(
                                    [partial_dim]
                                )

            if grad_op.type in _gradient_sync_by_partial_ops:
                infer_backward_op_partial_status(
                    vars, grad_op, grad_op_dist_attr
                )

            self._dist_context.set_op_dist_attr_for_program(
                grad_op, grad_op_dist_attr
            )

        loss_op = None
        first_backward_op_idx = -1
        for idx, op in enumerate(serial_main_program.global_block().ops):
            if is_loss_op(op):
                loss_op = op
            if is_loss_grad_op(op):
                assert op.type == "fill_constant"
                first_backward_op_idx = idx
                break

        assert first_backward_op_idx >= 0 and loss_op is not None, (
            "No backward procedure found in this program."
        )

        ops = list(serial_main_program.global_block().ops)
        vars = serial_main_program.global_block().vars
        dist_op_context = self._dist_context.dist_op_context
        grad_var_to_var = dist_op_context.grad_var_to_var[
            len(dist_op_context.grad_var_to_var)
        ]

        for idx in range(first_backward_op_idx, len(ops)):
            grad_op = ops[idx]
            # complete the initial grad loss op
            if idx == first_backward_op_idx:
                assert grad_op.type == "fill_constant"
                assert len(grad_op.input_arg_names) == 0, (
                    f"first backward op should has only ONE output, but got [{len(grad_op.input_arg_names)}]"
                )
                assert len(grad_op.output_arg_names) == 1, (
                    f"first backward op should has only ONE output, but got [{len(grad_op.output_arg_names)}]"
                )

                loss_var = vars[loss_op.output_arg_names[0]]
                loss_grad_var = vars[grad_op.output_arg_names[0]]
                assert loss_var.name + "@GRAD" == loss_grad_var.name
                dist_loss_var = self._dist_context.get_dist_tensor_for_program(
                    loss_var
                )
                dist_loss_op = self._dist_context.get_dist_op_for_program(
                    loss_op
                )

                set_var_dist_attr(
                    self._dist_context,
                    loss_grad_var,
                    dist_loss_var.dist_attr.dims_mapping,
                    dist_loss_var.dist_attr.process_mesh,
                    chunk_id=dist_loss_var.dist_attr.chunk_id,
                )
                naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
                    grad_op,
                    dist_loss_op.dist_attr.process_mesh,
                    dist_loss_op.dist_attr.get_output_dims_mapping(
                        loss_var.name
                    ),
                    self._dist_context,
                    chunk_id=dist_loss_op.dist_attr.chunk_id,
                )
                continue

            # complete the annotation of grad op (xxx_grad op or sum op)
            # xxx_grad op will have a corresponding forward op in grad_op_id_to_op_id
            if (
                grad_op.desc.original_id()
                in dist_op_context.grad_op_id_to_op_id
            ):
                # TODO support the case where one forward op corresponding to multiple xxx_grad op
                forward_op = _get_op_by_id(
                    ops[:first_backward_op_idx],
                    dist_op_context.grad_op_id_to_op_id[
                        grad_op.desc.original_id()
                    ],
                )
                assert forward_op is not None

                if grad_op.has_attr('sub_block') and forward_op.has_attr(
                    'sub_block'
                ):
                    _complete_grad_op_with_forward_op(forward_op, grad_op, vars)
                    grad_sub_block_id = grad_op.attr('sub_block').id
                    forward_sub_block_id = forward_op.attr('sub_block').id
                    grad_sub_block = serial_main_program.blocks[
                        grad_sub_block_id
                    ]
                    forward_sub_block = serial_main_program.blocks[
                        forward_sub_block_id
                    ]
                    for sub_grad_op in grad_sub_block.ops:
                        sub_forward_op = _get_op_by_id(
                            forward_sub_block.ops,
                            dist_op_context.grad_op_id_to_op_id[
                                sub_grad_op.desc.original_id()
                            ],
                        )
                        _complete_grad_op_with_forward_op(
                            sub_forward_op, sub_grad_op, grad_sub_block.vars
                        )
                else:
                    _complete_grad_op_with_forward_op(forward_op, grad_op, vars)

            # grad ops that have not a corresponding mapping in grad_op_id_to_op_id
            else:
                if grad_op.type in ['sum', 'grad_add']:
                    assert all(map(_is_grad_var_name, grad_op.input_arg_names))
                    output_name = grad_op.output_arg_names[0]
                    assert output_name in grad_var_to_var, (
                        f"sum op's output '{output_name}' has no corresponding var"
                    )
                    ref_fwd_var_name = grad_var_to_var[output_name]
                    ref_fwd_var = vars[ref_fwd_var_name]
                    ref_fwd_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            ref_fwd_var
                        )
                    )
                    ref_fwd_dims_mapping = ref_fwd_dist_attr.dims_mapping
                    ref_fwd_process_mesh = ref_fwd_dist_attr.process_mesh
                    ref_fwd_chunk_id = ref_fwd_dist_attr.chunk_id

                    # output
                    output_var = vars[output_name]
                    set_var_dist_attr(
                        self._dist_context,
                        output_var,
                        ref_fwd_dims_mapping,
                        ref_fwd_process_mesh,
                        chunk_id=ref_fwd_chunk_id,
                    )

                    # op
                    grad_op_dist_attr = OperatorDistAttr()
                    for var_name in grad_op.input_arg_names:
                        grad_op_dist_attr.set_input_dims_mapping(
                            var_name, ref_fwd_dims_mapping
                        )
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_name, ref_fwd_dims_mapping
                    )
                    # NOTE(zhaoyingli):
                    # The sum op is used to accumulate the grads' value of the same forward var,
                    # sum op's chunk_id is same with the last op which generate the grad.
                    ref_chunk_id = None
                    ref_process_mesh = None
                    for pre_idx in range(
                        idx - 1, first_backward_op_idx + 1, -1
                    ):
                        pre_grad_op = ops[pre_idx]
                        inter_arg_name = list(
                            set(pre_grad_op.output_arg_names)
                            & set(grad_op.input_arg_names)
                        )
                        if len(inter_arg_name) > 0:
                            pre_op_dist_attr = (
                                self._dist_context.get_op_dist_attr_for_program(
                                    pre_grad_op
                                )
                            )
                            ref_chunk_id = pre_op_dist_attr.chunk_id
                            ref_process_mesh = pre_op_dist_attr.process_mesh
                            break
                    assert (
                        ref_chunk_id is not None
                        and ref_process_mesh is not None
                    )
                    grad_op_dist_attr.process_mesh = ref_process_mesh
                    grad_op_dist_attr.chunk_id = ref_chunk_id
                    self._dist_context.set_op_dist_attr_for_program(
                        grad_op, grad_op_dist_attr
                    )

                elif grad_op.type == 'fill_any_like':
                    ref_var_name = grad_op.input_arg_names[0]
                    ref_var = vars[ref_var_name]
                    ref_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            ref_var
                        )
                    )
                    ref_dims_mapping = ref_dist_attr.dims_mapping
                    ref_process_mesh = ref_dist_attr.process_mesh
                    ref_chunk_id = ref_dist_attr.chunk_id
                    # var
                    output_var_name = grad_op.output_arg_names[0]
                    output_var = vars[output_var_name]
                    set_var_dist_attr(
                        self._dist_context,
                        output_var,
                        ref_dims_mapping,
                        ref_process_mesh,
                        chunk_id=ref_chunk_id,
                    )
                    # op
                    grad_op_dist_attr = OperatorDistAttr()
                    grad_op_dist_attr.process_mesh = ref_process_mesh
                    grad_op_dist_attr.chunk_id = ref_chunk_id
                    grad_op_dist_attr.set_input_dims_mapping(
                        ref_var_name, ref_dims_mapping
                    )
                    grad_op_dist_attr.set_output_dims_mapping(
                        output_var_name, ref_dims_mapping
                    )
                    self._dist_context.set_op_dist_attr_for_program(
                        grad_op, grad_op_dist_attr
                    )
                else:
                    raise ValueError(f"got unexpected op [{grad_op.type}]")

    def complete_update_annotation(self, serial_main_program):
        """Complete the annotation of vars and ops in the update phase for parallel program."""
        # Copy the dist tensors and dist ops annotated by users from the default context
        # global mesh
        from paddle.distributed.auto_parallel.static.process_group import (
            get_world_process_group,
        )

        world_ranks = get_world_process_group().ranks

        # Notice: serial_main_program is actually a dist_main_program of current rank,
        # and must be passed into this function.
        # TODO: We should fix this behavior.

        ops = list(serial_main_program.global_block().ops)
        vars = serial_main_program.global_block().vars
        learning_rate_completed = False

        for idx in range(len(ops)):
            # complete the annotation of the optimizer op.
            # TODO to add attribute for moment var
            op = ops[idx]
            if int(op.attr('op_role')) == int(OpRole.Optimize):
                if is_gradient_clip_op(op):
                    if op.type in _g_gradient_clip_ops:
                        # complete op dist_attr with global world ranks
                        op_dist_attr = OperatorDistAttr()
                        op_dist_attr.process_mesh = ProcessMesh(world_ranks)

                        for in_name in op.input_arg_names:
                            in_var = vars[in_name]
                            in_dist_attr = self._dist_context.get_tensor_dist_attr_for_program(
                                in_var
                            )
                            op_dist_attr.set_input_dims_mapping(
                                in_name, in_dist_attr.dims_mapping
                            )
                        for out_name in op.output_arg_names:
                            out_var = vars[out_name]
                            out_dist_attr = TensorDistAttr()
                            out_dist_attr.process_mesh = ProcessMesh(
                                world_ranks
                            )
                            out_dist_attr.dims_mapping = [
                                -1 for _ in out_var.shape
                            ]
                            self._dist_context.set_tensor_dist_attr_for_program(
                                out_var, out_dist_attr
                            )
                            op_dist_attr.set_output_dims_mapping(
                                out_name, out_dist_attr.dims_mapping
                            )
                    else:
                        # get ref_process_mesh and ref_dims_mapping from input_var
                        in_var = vars[op.input("X")[0]]
                        in_dist_attr = (
                            self._dist_context.get_tensor_dist_attr_for_program(
                                in_var
                            )
                        )
                        assert in_dist_attr is not None
                        ref_process_mesh = in_dist_attr.process_mesh
                        ref_dims_mapping = in_dist_attr.dims_mapping
                        ref_chunk_id = in_dist_attr.chunk_id

                        if (
                            op.type == "cast"
                            and ops[idx + 1].type == "elementwise_mul"
                        ):
                            ref_var = vars[ops[idx + 1].input("X")[0]]
                            ref_dist_attr = self._dist_context.get_tensor_dist_attr_for_program(
                                ref_var
                            )
                            assert ref_dist_attr is not None
                            ref_process_mesh = ref_dist_attr.process_mesh

                        # complete out_var's tensor_dist_attr
                        out_var = vars[op.output("Out")[0]]
                        out_dist_attr = (
                            self._dist_context.get_tensor_dist_attr_for_program(
                                out_var
                            )
                        )
                        if not out_dist_attr:
                            out_dist_attr = TensorDistAttr()
                            out_dist_attr.process_mesh = ref_process_mesh
                            out_dist_attr.chunk_id = ref_chunk_id
                            if out_var.shape == in_var.shape:
                                out_dist_attr.dims_mapping = ref_dims_mapping
                            else:
                                assert (
                                    len(out_var.shape) == 1
                                    and out_var.shape[0] == 1
                                )
                                out_dist_attr.dims_mapping = [
                                    -1 for _ in out_var.shape
                                ]
                            self._dist_context.set_tensor_dist_attr_for_program(
                                out_var, out_dist_attr
                            )

                        # complete op's dist_attr
                        op_dist_attr = OperatorDistAttr()
                        op_dist_attr.process_mesh = ref_process_mesh
                        for in_name in op.input_arg_names:
                            in_var = vars[in_name]
                            in_dist_attr = self._dist_context.get_tensor_dist_attr_for_program(
                                in_var
                            )
                            op_dist_attr.set_input_dims_mapping(
                                in_name, in_dist_attr.dims_mapping
                            )
                        for out_name in op.output_arg_names:
                            out_var = vars[out_name]
                            out_dist_attr = self._dist_context.get_tensor_dist_attr_for_program(
                                out_var
                            )
                            op_dist_attr.set_output_dims_mapping(
                                out_name, out_dist_attr.dims_mapping
                            )

                    self._dist_context.set_op_dist_attr_for_program(
                        op, op_dist_attr
                    )

                if "Grad" in op.input_names and "Param" in ops[idx].input_names:
                    assert len(op.input("Param")) == 1, (
                        "Only support one-to-one now."
                    )
                    assert len(op.input("Grad")) == 1, (
                        "Only support one-to-one now."
                    )
                    param = vars[op.input("Param")[0]]
                    grad_var = vars[op.input("Grad")[0]]

                    param_dist_attr = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            param
                        )
                    )
                    assert param_dist_attr is not None
                    ref_process_mesh = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            param
                        ).process_mesh
                    )
                    assert ref_process_mesh is not None
                    ref_dims_mapping = (
                        self._dist_context.get_tensor_dist_attr_for_program(
                            param
                        ).dims_mapping
                    )
                    assert ref_dims_mapping is not None
                    op_dist_attr = OperatorDistAttr()
                    op_dist_attr.process_mesh = ref_process_mesh
                    op_dist_attr.set_input_dims_mapping(
                        grad_var.name, ref_dims_mapping
                    )
                    op_dist_attr.set_input_dims_mapping(
                        param.name, ref_dims_mapping
                    )
                    op_dist_attr.set_output_dims_mapping(
                        param.name, ref_dims_mapping
                    )
                    learning_var = vars[op.input("LearningRate")[0]]
                    op_dist_attr.set_input_dims_mapping(
                        learning_var.name, [-1 for _ in learning_var.shape]
                    )
                    op_dist_attr.set_output_dims_mapping(
                        learning_var.name, [-1 for _ in learning_var.shape]
                    )

                    if not learning_rate_completed:
                        learning_rate_completed = True
                        var_dist_attr = TensorDistAttr()
                        var_dist_attr.process_mesh = ProcessMesh(world_ranks)
                        var_dist_attr.dims_mapping = [
                            -1 for _ in learning_var.shape
                        ]
                        self._dist_context.set_tensor_dist_attr_for_program(
                            learning_var, var_dist_attr
                        )

                    for input_name in op.desc.input_names():
                        if input_name in [
                            'Param',
                            'Grad',
                            'LearningRate',
                            "Beta1Tensor",
                            "Beta2Tensor",
                            "EpsilonTensor",
                        ]:
                            continue
                        if len(op.desc.input(input_name)) == 0:
                            continue

                        assert len(op.desc.input(input_name)) == 1
                        input_var = vars[op.desc.input(input_name)[0]]
                        input_var_attr = TensorDistAttr()

                        if (
                            "Beta1Pow" in input_name
                            or "Beta2Pow" in input_name
                            or "SkipUpdate" in input_name
                        ):
                            input_var_attr.dims_mapping = [-1]
                            op_dist_attr.set_input_dims_mapping(
                                input_var.name, [-1 for _ in input_var.shape]
                            )
                            op_dist_attr.set_output_dims_mapping(
                                input_var.name, [-1 for _ in input_var.shape]
                            )
                        else:
                            input_var_attr.dims_mapping = ref_dims_mapping
                            op_dist_attr.set_input_dims_mapping(
                                input_var.name, ref_dims_mapping
                            )
                            op_dist_attr.set_output_dims_mapping(
                                input_var.name, ref_dims_mapping
                            )
                        if "SkipUpdate" not in input_name:
                            input_var_attr.process_mesh = ref_process_mesh
                            self._dist_context.set_tensor_dist_attr_for_program(
                                input_var, input_var_attr
                            )

                    self._dist_context.set_op_dist_attr_for_program(
                        op, op_dist_attr
                    )
                    continue

    def complete_prim_annotation(self, serial_main_program=None):
        """
        fill default data parallel annotation for program with primitive operators.

        Arguments:
            serial_main_program: partial annotated serial_main_program.
        Returns:
            serial_main_program: completed annotated serial_main_program.
        """
        if serial_main_program is None:
            serial_main_program = self._dist_context.serial_main_program
        else:
            self._dist_context._serial_main_program = serial_main_program

        self._dist_context._is_initialized = True
        self._dist_context._init_dist_attr_for_program()
        self._init_global_mesh_for_program()
        # Do the validation check and amend some completion
        self._dist_context.amend_dist_attr_for_program()
        self._dist_context.validate_dist_attr_for_program()

    def _init_global_mesh_for_program(self):
        # Copy the dist tensors and dist ops annotated by users from the default context
        # global mesh
        from paddle.distributed.auto_parallel.static.process_group import (
            get_world_process_group,
        )

        world_ranks = get_world_process_group().ranks

        for block in self._dist_context._serial_main_program.blocks:
            for tensor in block.vars.values():
                # Copy the distributed tensors in the default context
                dist_tensor = self._dist_context.get_dist_tensor_for_program(
                    tensor
                )
                assert dist_tensor is not None
                dist_tensor.dist_attr.process_mesh = ProcessMesh(world_ranks)
            for op in block.ops:
                # Copy the distributed operators in the default context
                dist_op = self._dist_context.get_dist_op_for_program(op)
                assert dist_op is not None
                dist_op.dist_attr.process_mesh = ProcessMesh(world_ranks)

                # Find the most compatible implementations from the distributed operator
                op_dist_impls = find_compatible_distributed_operator_impls(
                    dist_op, fwd=True
                )
                if op_dist_impls is not None:
                    backup_op_dist_attr = copy.deepcopy(dist_op.dist_attr)
                    for op_dist_impl in op_dist_impls:
                        dim_changed = op_dist_impl.update_dims_mapping(dist_op)
                        if op_dist_impl.is_auto_compatible(dist_op):
                            # if op_dist_impl.type == "elementwise":
                            #     dist_op.dist_attr.impl_type = "default"
                            # else:
                            dist_op.dist_attr.impl_type = op_dist_impl.type
                            # op_dist_attr.impl_type = op_dist_impl.type
                            dist_op.dist_attr.impl_idx = op_dist_impl.idx
                            break
                        else:
                            dist_op.dist_attr = backup_op_dist_attr