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

from collections import OrderedDict

import numpy as np

import paddle
from paddle.utils.flops import flops

from ..cluster import DeviceType, LinkType, get_default_cluster
from ..dist_tensor import DistributedTensor
from ..process_group import get_process_group
from ..utils import _get_comm_group, _get_idx_in_axis

COMM_OP_TYPE = [
    "send_v2",
    "recv_v2",
    "broadcast",
    "all_gather",
    "all_reduce",
    "c_allreduce_sum",
    "c_identity",
]
NON_COMP_TYPE = ["while", *COMM_OP_TYPE]
_g_op_cost_factory = {}


def build_comp_desc_from_op(op):
    """Build the description of computation op."""
    # NOTE: The desc is for serial op.
    from ..reshard import get_var_with_recursion

    desc = {}
    # The desc of concat op is {"op": "concat", "inputs": {"X": [(paddle.float32, [20, 20]), (paddle.float32, [20, 20])]}, "outputs": {"Out": [(paddle.float32, [20, 40])], "attrs": {"axis": -1}}}
    vars = op.block.vars
    desc["op"] = op.type
    input_desc = OrderedDict()
    for input_name in op.input_names:
        var_name_list = op.input(input_name)
        var_desc = []
        for var_name in var_name_list:
            var = get_var_with_recursion(var_name, op.block, op.block.program)
            shape = var.shape
            var_desc.append((var.dtype, shape))
        input_desc[input_name] = var_desc
    desc["inputs"] = input_desc

    output_desc = OrderedDict()
    for out_name in op.output_names:
        var_name_list = op.output(out_name)
        var_desc = []
        for var_name in var_name_list:
            var = get_var_with_recursion(var_name, op.block, op.block.program)
            shape = var.shape
            var_desc.append((var.dtype, shape))
            desc["dtype"] = var.dtype
        output_desc[out_name] = var_desc
    desc["outputs"] = output_desc

    attr_desc = op.all_attrs
    desc["attrs"] = attr_desc

    return desc


def build_comp_desc_from_dist_op(dist_op, dist_context):
    """Build descriptions of computation op distributed on the processes."""
    from ..reshard import get_var_with_recursion

    op_descs = {}
    op = dist_op.serial_op
    dist_attr = dist_op.dist_attr
    process_mesh = dist_attr.process_mesh
    assert process_mesh, "Process mesh must not be None."
    processes = process_mesh.process_ids
    for process in processes:
        desc = {}
        desc["op"] = op.type
        attr_desc = op.all_attrs()
        # NOTE: The attrs of desc is replica of serial op, there may be a bug if shape need to be partitioned involved in attrs.
        desc["attrs"] = attr_desc
        input_desc = OrderedDict()
        output_desc = OrderedDict()

        # Get partitioned shape of input
        input_var_desc = {}
        for input_name in op.input_names:
            var_name_list = op.input(input_name)
            input_var_desc[input_name] = []
            for var_name in var_name_list:
                var = get_var_with_recursion(
                    var_name, op.block, op.block.program
                )
                # Use op input_dims_mapping
                dims_mapping = dist_attr.get_input_dims_mapping(var_name)
                global_sizes = var.shape
                # NOTE: When support uneven partition, the shard_sizes will be got from dist_attr.
                shard_sizes = None
                topology = process_mesh.shape
                shape = DistributedTensor.get_local_sizes(
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
                input_var_desc[input_name].append(shape)

                # For special op such as embedding and its grad op
                if (
                    op.type == "c_embedding"
                    or op.type == "lookup_table_v2"
                    or op.type == "c_embedding_grad"
                    or op.type == "lookup_table_v2_grad"
                ):
                    if input_name == "W":
                        embedding_row_dim_mapping = (
                            dist_attr.get_input_dims_mapping(
                                op.input(input_name)[0]
                            )[0]
                        )
                        relative_idx = _get_idx_in_axis(
                            processes,
                            dist_attr.process_mesh.shape,
                            embedding_row_dim_mapping,
                            process,
                        )
                        per_part_size = shape[0]
                        relative_idx = relative_idx * per_part_size
                        desc["attrs"]["start_index"] = relative_idx

        desc["inputs"] = input_var_desc

        for out_name in op.output_names:
            var_name_list = op.output(out_name)
            var_desc = []
            for var_name in var_name_list:
                # Use op output_dims_mapping
                var = get_var_with_recursion(
                    var_name, op.block, op.block.program
                )
                dist_attr = dist_op.dist_attr
                dims_mapping = dist_attr.get_output_dims_mapping(var_name)
                process_mesh = dist_attr.process_mesh
                global_sizes = var.shape
                shard_sizes = None
                processes = process_mesh.process_ids
                topology = process_mesh.shape
                shape = DistributedTensor.get_local_sizes(
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
                var_desc.append((var.dtype, shape))
                desc["dtype"] = var.dtype

                # For special op such as fill_constant_batch_size_like
                if op.type == "fill_constant_batch_size_like":
                    # Modify shape attr according to how output are partitioned
                    out_name = var_name_list[0]
                    dims_mapping = dist_attr.get_output_dims_mapping(out_name)
                    process_mesh_shape = dist_attr.process_mesh.shape
                    shape_list = op.attr("shape")
                    # Modify target shape
                    for idx, axis in enumerate(dims_mapping):
                        if axis >= 0:
                            shape_list[idx] = (
                                shape_list[idx] // process_mesh_shape[axis]
                            )
                    desc["attrs"]["shape"] = shape_list
            output_desc[out_name] = var_desc

        desc["outputs"] = output_desc

        op_descs[process] = desc

    return op_descs


def build_comp_desc_str_for_predict(desc):
    # NOTE: The description format may change in the future.
    def _parse_dtype(dtype):
        dtype_str = ""
        if dtype == paddle.float32:
            dtype_str = "float32"
        elif dtype == paddle.float16:
            dtype_str = "float16"
        elif dtype == paddle.int32:
            dtype_str = "int32"
        elif dtype == paddle.int64:
            dtype_str = "int64"
        elif dtype == paddle.unit8:
            dtype_str = "unit8"
        else:
            raise TypeError(f"Unsupported dtype {dtype}")
        return dtype_str

    assert isinstance(desc, dict)
    desc_str_list = []
    desc_str = None
    dtype_str_list = []
    dims_list = []
    shape_list = []

    desc_str_list.append(desc["op"])
    inputs = desc["inputs"]
    for key, item in inputs.items():
        for dtype, shape in item:
            dtype_str_list.append(_parse_dtype(dtype))
            shape_list += list(shape)
            dims = len(shape)
            dims_list.append(dims)

    dtype_str = "*".join(dtype_str_list)
    dims_list = [str(item) for item in dims_list]
    dims_str = "*".join(dims_list)

    shape_list = [str(item) for item in shape_list]
    shape_str = "[" + ",".join(shape_list) + "]"
    desc_str_list += [dtype_str, dims_str, shape_str]
    desc_str = "_".join(desc_str_list)
    attrs = desc["attrs"]
    parse_result = (desc_str, attrs)
    return parse_result


def build_comm_desc_from_dist_op(
    op_type,
    dist_op,
    ctx,
    var_names,
    attrs=None,
    parallel_axis=None,
    group_ranks=None,
):
    """Build descriptions of communication op distributed on the processes."""
    from ..reshard import get_var_with_recursion

    specific_op_type = []
    dist_attr = dist_op.dist_attr
    assert dist_attr, "Dist attr must not be None."
    process_mesh = dist_attr.process_mesh
    assert process_mesh, "Process mesh must not be None."

    processes = process_mesh.process_ids
    op_descs = {}
    for process in processes:
        rank_id = process
        desc = {}
        desc["op"] = op_type
        op_attrs = None
        comm_group_ranks = None

        if op_type not in specific_op_type:
            serial_op = dist_op.serial_op
            input_list = []
            # The var_names usually contain just one item.
            for var_name in var_names:
                dist_attr = dist_op.dist_attr
                has_found = False
                # Find var_name in serial op input or output
                for name in dist_op.serial_op.input_arg_names:
                    # If a tensor is the input of multi ops, sum the grad of all ops, so the name will be varname@RENAME@block@0 and so on.
                    if var_name in name:
                        var_name = name
                        has_found = True
                        break

                if not has_found:
                    for name in dist_op.serial_op.output_arg_names:
                        if var_name in name:
                            var_name = name
                            has_found = True
                            break
                assert has_found
                var = get_var_with_recursion(
                    var_name, serial_op.block, serial_op.block.program
                )

                dims_mapping = (
                    dist_attr.get_input_dims_mapping(var_name)
                    if var_name in dist_op.serial_op.input_arg_names
                    else dist_attr.get_output_dims_mapping(var_name)
                )
                global_sizes = var.shape
                shard_sizes = None
                topology = process_mesh.shape
                shape = DistributedTensor.get_local_sizes(
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
                input_list.append((var.dtype, shape))

            # NOTE: The input_name of comm ops used usually is X.
            if op_type == "all_reduce":
                desc["inputs"] = {"x": input_list}
            else:
                desc["inputs"] = {"X": input_list}

            # Get comm group by parallel_axis or the given group_ranks.
            if parallel_axis is not None:
                process_mesh_shape = process_mesh.shape
                process_mesh_group = process_mesh.process_ids
                comm_group_ranks = _get_comm_group(
                    process_mesh_group,
                    process_mesh_shape,
                    parallel_axis,
                    rank_id,
                )
            elif group_ranks is not None:
                comm_group_ranks = group_ranks
            else:
                raise ValueError(
                    "The parallel_axis and group_ranks can not be None in the same."
                )

            if attrs is not None:
                assert isinstance(attrs, dict)
                op_attrs = attrs
            else:
                op_attrs = {}

            desc["attrs"] = op_attrs
            desc["group_ranks"] = comm_group_ranks

            op_descs[rank_id] = desc

    return op_descs


def build_comm_desc(op_type, group_ranks, dtype, shape, attrs=None):
    """Build a comm desc directly."""
    desc = {}
    desc["op"] = op_type
    desc["group_ranks"] = group_ranks
    if op_type == "all_reduce":
        desc["inputs"] = {"x": [(dtype, shape)]}
    else:
        desc["inputs"] = {"X": [(dtype, shape)]}
    desc["attrs"] = attrs
    return desc


def build_comm_costs_from_descs(
    op_cost_class, ctx, processes, descs, cluster, is_dp=False
):
    """Build comm costs by descriptions"""
    comm_context = CommContext(cluster)
    group_ranks_list = []
    comm_op_cost_list = []
    for process in processes:
        desc = descs[process]
        group_ranks = desc["group_ranks"]
        if group_ranks not in group_ranks_list:
            group_ranks_list.append(group_ranks)
            comm_op_cost = op_cost_class(
                op_desc=desc, comm_context=comm_context
            )
            if is_dp:
                comm_op_cost.cost.time *= 0.9
            comm_op_cost_list.append(comm_op_cost)
    return comm_op_cost_list


def build_comp_costs_from_descs(op_cost_class, ctx, processes, descs, cluster):
    """Build comp costs by descriptions."""
    costs = {}
    for process in processes:
        costs[process] = op_cost_class(
            op_desc=descs[process], cluster=cluster, rank=process
        )
    return costs


def build_dp_costs(
    result, dist_op, ctx, var_names, attrs, parallel_axis, cluster
):
    """DP cost contains a allreduce_sum op cost and a scale op cost"""
    # The costs will be appended in the given result.
    from ..reshard import get_var_with_recursion

    dist_attr = dist_op.dist_attr
    process_mesh = dist_attr.process_mesh
    processes = process_mesh.process_ids
    assert len(var_names) == 1
    vars = dist_op.serial_op.block.vars
    var_name = var_names[0]
    has_found = False
    is_input = True
    for name in dist_op.serial_op.input_arg_names:
        if var_name in name:
            var_name = name
            has_found = True
            break

    if not has_found:
        for name in dist_op.serial_op.output_arg_names:
            if var_name in name:
                var_name = name
                has_found = True
                is_input = False
                break
    if not has_found:
        return

    all_reduce_sum_descs = build_comm_desc_from_dist_op(
        "all_reduce",
        dist_op,
        ctx,
        var_names,
        attrs=attrs,
        parallel_axis=parallel_axis,
    )
    comm_cost_list = build_comm_costs_from_descs(
        _g_op_cost_factory["all_reduce"],
        ctx,
        processes,
        all_reduce_sum_descs,
        cluster,
        is_dp=True,
    )
    result.append(comm_cost_list)

    # The scale op just on the group_ranks
    for comm_cost in comm_cost_list:
        group_ranks = comm_cost.group_ranks
        dp_degree = len(group_ranks)
        scale_costs = {}
        op_type = "scale"
        for rank in group_ranks:
            desc = {}
            desc["op"] = op_type
            desc["inputs"] = {}
            dims_mapping = (
                dist_attr.get_input_dims_mapping(var_name)
                if is_input
                else dist_attr.get_output_dims_mapping(var_name)
            )
            var = get_var_with_recursion(
                var_name,
                dist_op.serial_op.block,
                dist_op.serial_op.block.program,
            )
            global_sizes = var.shape
            shard_sizes = None
            topology = process_mesh.shape
            shape = DistributedTensor.get_local_sizes(
                global_sizes,
                dims_mapping,
                topology,
                processes,
                rank,
                shard_sizes,
            )
            desc["inputs"]["X"] = [(var.dtype, shape)]
            attrs = {"scale": 1.0 / dp_degree}
            desc["attrs"] = attrs
            desc["dtype"] = var.dtype
            scale_op_cost = _g_op_cost_factory["scale"](
                op_desc=desc, cluster=cluster, rank=rank
            )
            scale_costs[rank] = scale_op_cost
        result.append(scale_costs)


class CommContext:
    _instance = None
    _has_instance = False

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            _has_instance = True
        return cls._instance

    def __init__(self, cluster):
        if CommContext._has_instance:
            return
        self.beta = {}
        self.hops = {}
        assert cluster is not None
        self.cluster = cluster
        # if cluster has no info about those vars, it will be set by default
        self.base_ring = None
        self.base_tree = None
        self.intra_ring = None
        self.intra_tree = None
        self.inter_ring = None
        self.inter_tree = None
        self.switch = None
        self._post_init()

    def _post_init(self):
        alpha_latency = self.cluster.alpha_latency
        if alpha_latency is None:
            # set default
            self.base_ring = 8.4
            self.base_tree = 0.0
            # NVL in default
            self.intra_ring = 3.4
            self.intra_tree = 28
            # NET in default
            self.inter_ring = 9.6
            self.inter_tree = 28
            self.switch = 10.0
        else:
            base_ring = alpha_latency.base_ring
            self.base_ring = base_ring if base_ring is not None else 8.4

            base_tree = alpha_latency.base_tree
            self.base_tree = base_tree if base_tree is not None else 0.0

            intra_ring = alpha_latency.intra_ring
            if intra_ring == LinkType.NVL:
                self.intra_ring = 3.4
            elif intra_ring == LinkType.PHB:
                self.intra_ring = 5.7
            elif intra_ring is not None:
                self.intra_ring = intra_ring
            else:
                # NVL Default
                self.intra_ring = 3.4

            intra_tree = alpha_latency.intra_tree
            if intra_tree == LinkType.NVL:
                self.intra_tree = 28
            elif intra_tree == LinkType.PHB:
                self.intra_tree = 28
            elif intra_tree is not None:
                self.intra_tree = intra_tree
            else:
                # NVL Default
                self.intra_tree = 28

            inter_ring = alpha_latency.inter_ring
            if inter_ring == LinkType.NET:
                self.inter_ring = 9.6
            elif inter_ring is not None:
                self.inter_ring = inter_ring
            else:
                # NET Default
                self.inter_ring = 9.6

            inter_tree = alpha_latency.inter_tree
            if inter_tree == LinkType.NET:
                self.inter_tree = 28
            elif inter_tree is not None:
                self.inter_tree = inter_tree
            else:
                # NET Default
                self.inter_tree = 28

            switch = alpha_latency.switch
            self.switch = switch if switch is not None else 10

            assert self.base_ring is not None
            assert self.base_tree is not None
            assert self.intra_ring is not None
            assert self.intra_tree is not None
            assert self.inter_ring is not None
            assert self.inter_tree is not None
            assert self.switch is not None

    def get_max_beta(self, ranks):
        # NOTE: Get beta by ring, even in the case of tree such as tree broadcast
        ranks = self.cluster.convert_rank_to_device_id(ranks)
        key = ','.join(map(str, sorted(ranks)))
        if key in self.beta:
            return self.beta[key]
        max_beta = None
        for i in range(len(ranks)):
            for j in range(i + 1, len(ranks)):
                forward_order_beta = self.cluster.get_beta(ranks[i], ranks[j])
                backward_order_beta = self.cluster.get_beta(ranks[j], ranks[i])
                beta = max(backward_order_beta, forward_order_beta)
                if max_beta is None or beta > max_beta:
                    max_beta = beta
        if max_beta is None:
            max_beta = 0
        self.beta[key] = max_beta
        return max_beta

    def get_hops(self, ranks):
        key = ','.join(map(str, sorted(ranks)))
        hops = 0
        for i in range(len(ranks)):
            for j in range(i + 1, len(ranks)):
                hop = self.cluster.get_hop(ranks[i], ranks[j])
                hops += hop
        self.hops[key] = hops

        return hops


class Cost:
    def __init__(self, time=0, memory=0, flops=0):
        self.time = time
        self.memory = memory
        self.flops = flops

    def _check_time(self, val):
        assert val >= 0, "Time must be greater than or equal to 0."

    def _check_memory(self, val):
        assert isinstance(val, int) and val >= 0, (
            "Memory must be int and greater than equal to 0."
        )

    def _check_flops(self, val):
        assert isinstance(val, int) and val >= 0, (
            "FLOPs must be int and greater than equal to 0."
        )

    @property
    def time(self):
        return self._time

    @time.setter
    def time(self, val):
        self._check_time(val)
        self._time = val

    @property
    def memory(self):
        return self._memory

    @memory.setter
    def memory(self, val):
        self._check_memory(val)
        self._memory = val

    @property
    def flops(self):
        return self._flops

    @flops.setter
    def flops(self, val):
        self._check_flops(val)
        self._flops = val

    def __add__(self, rhs):
        assert isinstance(rhs, Cost)
        time = self.time + rhs.time
        memory = self.memory + rhs.memory
        flops = self.flops + rhs.flops
        assert time >= 0 and memory >= 0 and flops >= 0
        return Cost(time, memory, flops)

    def __sub__(self, rhs):
        assert isinstance(rhs, Cost)
        time = self.time - rhs.time
        memory = self.memory - rhs.memory
        flops = self.flops - rhs.flops
        assert time >= 0 and memory >= 0 and flops >= 0
        return Cost(time, memory, flops)


class OpCost:
    OP_TYPE = "op"

    def __init__(self, op=None, op_desc=None):
        self._op = op
        self._op_desc = op_desc
        self._cost = None

    @property
    def op(self):
        return self._op

    @property
    def op_desc(self):
        return self._op_desc

    @property
    def time(self):
        return self.cost.time

    @property
    def memory(self):
        return self.cost.memory

    @property
    def flops(self):
        return self.cost.flops

    @property
    def cost(self):
        return self._cost

    def calc_time(self):
        return 0

    def calc_memory(self):
        return 0

    def calc_flops(self):
        return 0

    def calc_cost(self):
        time = self.calc_time()
        memory = self.calc_memory()
        flops = self.calc_flops()
        cost = Cost(time, memory, flops)
        return cost

    def __add__(self, rhs):
        assert isinstance(rhs, (OpCost, Cost))
        time = 0
        memory = 0
        flops = 0
        if isinstance(rhs, OpCost):
            time = self.cost.time + rhs.cost.time
            memory = self.cost.memory + rhs.cost.memory
            flops = self.cost.flops + rhs.cost.flops
            assert time >= 0 and memory >= 0 and flops >= 0
        elif isinstance(rhs, Cost):
            time = self.time + rhs.time
            memory = self.memory + rhs.memory
            flops = self.flops + rhs.flops
            assert time >= 0 and memory >= 0 and flops >= 0
        return Cost(time, memory, flops)

    def __sub__(self, rhs):
        assert isinstance(rhs, (OpCost, Cost))
        time = 0
        memory = 0
        flops = 0
        if isinstance(rhs, OpCost):
            time = self.cost.time - rhs.cost.time
            memory = self.cost.memory - rhs.cost.memory
            flops = self.cost.flops - rhs.cost.flops
            assert time >= 0 and memory >= 0 and flops >= 0
        elif isinstance(rhs, Cost):
            time = self.time - rhs.time
            memory = self.memory - rhs.memory
            flops = self.flops - rhs.flops
            assert time >= 0 and memory >= 0 and flops >= 0
        return Cost(time, memory, flops)


class CommOpCost(OpCost):
    OP_TYPE = "COMM"

    def __init__(self, op=None, op_desc=None, comm_context=None):
        super().__init__(op=op, op_desc=op_desc)
        self._check_comm_op_type()
        self._comm_context = comm_context
        self._group_ranks = None
        self._comm_count = None
        self._hops = None
        self._rank_count = len(self.group_ranks)
        self._machine_count = None
        self._cost = self.calc_cost()

    @property
    def comm_context(self):
        return self._comm_context

    @property
    def comm_count(self):
        from ..reshard import get_var_with_recursion

        if self._comm_count is None:
            dtype = None
            shape = None
            if self.op is not None:
                vars = self.op.block.vars
                # NOTE: The tensor communicated input_name is "X" in default. Otherwise, this function should be overridden
                try:
                    var_name = self.op.input("X")[0]
                except:
                    var_name = self.op.output("Out")[0]
                var = get_var_with_recursion(
                    var_name, self.op.block, self.op.block.program
                )
                dtype = var.dtype
                shape = var.shape
            elif self.op_desc is not None:
                dtype = self.op_desc["inputs"]["X"][0][0]
                shape = self.op_desc["inputs"]["X"][0][1]

            factor = None
            if dtype == paddle.float32 or dtype == paddle.int32:
                factor = 4
            elif dtype == paddle.int64:
                factor = 8
            elif dtype == paddle.uint8:
                factor = 1
            elif dtype == paddle.float16:
                factor = 2
            elif dtype == paddle.bool:
                factor = 8
            else:
                raise ValueError(f"Unsupported comm dtype {dtype}")
            comm_count = int(np.prod(shape)) * factor
            self._comm_count = comm_count

        return self._comm_count

    @property
    def rank_count(self):
        return self._rank_count

    @property
    def machine_count(self):
        if self._machine_count is None:
            cluster = self._comm_context.cluster
            self._machine_count = cluster.get_involved_machine_count(
                self.group_ranks
            )
        return self._machine_count

    @property
    def hops(self):
        if self._hops is None:
            self._hops = self.comm_context.get_hops(self.group_ranks)
        return self._hops

    @property
    def group_ranks(self):
        if self._group_ranks is None:
            if self.op_desc is not None:
                self._group_ranks = self.op_desc["group_ranks"]
            elif self.op is not None:
                ring_id = self.op.attr("ring_id")
                process_group = get_process_group(ring_id)
                if process_group is None:
                    raise ValueError(
                        f"There not exists process group whose ring_id is {ring_id}."
                    )
                self._group_ranks = process_group.ranks
        return self._group_ranks

    @classmethod
    def _check_comm_op_type(cls):
        if cls.OP_TYPE != "COMM":
            if cls.OP_TYPE not in COMM_OP_TYPE:
                raise TypeError(
                    f"Please Check op type in {COMM_OP_TYPE}, but got {cls.OP_TYPE}."
                )


class CompOpCost(OpCost):
    OP_TYPE = "COMP"

    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
        super().__init__(op=op, op_desc=op_desc)
        self._check_comp_op_type()
        self.cluster = cluster
        self.rank = rank
        self._cost = self.calc_cost()

    @classmethod
    def _check_comp_op_type(cls):
        if cls.OP_TYPE != "COMP":
            if cls.OP_TYPE in NON_COMP_TYPE:
                raise TypeError(
                    f"Please Check op type not in {NON_COMP_TYPE}, but got {cls.OP_TYPE}."
                )

    def get_rank_gflops(self, rank, dtype):
        device = self.cluster.get_device(rank)
        gflops = 7800
        if dtype == paddle.float64:
            gflops = device.dp_gflops
        elif dtype == paddle.float32:
            gflops = device.sp_gflops
        elif dtype == paddle.float16 or dtype == paddle.bfloat16:
            gflops = device.hp_gflops
        return gflops

    def calc_flops(self):
        if not self.op_desc:
            return 0
        if "_grad" in self.__class__.OP_TYPE:
            op_type = self.__class__.OP_TYPE[: len(self.__class__.OP_TYPE) - 5]
            return 2 * flops(
                op_type, self.op_desc["inputs"], self.op_desc["attrs"]
            )
        return flops(
            self.__class__.OP_TYPE,
            self.op_desc["inputs"],
            self.op_desc["attrs"],
        )

    def calc_time(self):
        if self.rank is None or self.op_desc is None:
            device_gflops = 7800
        else:
            device_gflops = self.get_rank_gflops(
                self.rank, self.op_desc["dtype"]
            )
        flops_count = self.calc_flops()
        utilization_rate = 0.65
        return flops_count / (utilization_rate * device_gflops) * 1e-3


def register_op_cost(cls):
    op_type = cls.OP_TYPE

    def register(op_type):
        global _g_op_cost_factory
        _g_op_cost_factory[op_type] = cls

    register(op_type)
    return cls


def calc_time_by_modeling(op=None, desc=None, cluster=None):
    op_type = op.type if op is not None else desc["op"]
    if op_type in COMM_OP_TYPE:
        op_cost = _g_op_cost_factory[op_type](
            op=op, op_desc=desc, comm_context=CommContext(cluster)
        )
    elif op_type not in NON_COMP_TYPE:
        op_cost = _g_op_cost_factory[op_type](
            op=op, op_desc=desc, cluster=cluster
        )
    time = op_cost.calc_time()
    return time


def calc_time_by_cost_model(op, cluster=None):
    """Calc op time by cost model and the unit is microsecond."""
    if not isinstance(op, paddle.base.framework.Operator):
        raise TypeError(
            f"OP must be paddle.base.framework.Operator, but got {type(op)}."
        )
    if not cluster:
        cluster = get_default_cluster()

    assert cluster._gpu_model in [
        "V100",
        "A100",
    ], "Only A100 and V100 gpu has been supported currently."

    time = 0.0  # microsecond
    op_type = op.type
    # calc comp op time by flops
    if op_type not in NON_COMP_TYPE:
        attrs = op.all_attrs()
        # build comp op inputs desc to calc flops.
        # for example, a matmul op inputs desc will be {"X": [(1024, 1024)], "Y": [(1024, 1024)]}
        inputs = {}
        for input_name in op.input_names:
            var_names = op.input(input_name)
            inputs[input_name] = []
            for var_name in var_names:
                var = op.block._var_recursive(var_name)
                inputs[input_name].append(var.shape)

        # the time of grad operator is twice than its forward operator empirically
        if "_grad" in op_type:
            op_type = op_type[: len(op_type) - 5]
            flops_count = 2 * flops(op_type, inputs, attrs)
        else:
            flops_count = flops(op_type, inputs, attrs)

        # FIXME(Ruibiao): Need a better way to get dtype
        var_name = op.output_arg_names[0]
        dtype = op.block._var_recursive(var_name).dtype
        device = cluster.get_device(0)
        assert device.type == DeviceType.GPU, (
            "Only GPU device is supported currently."
        )

        gflops = 0.0
        if dtype == paddle.float64:
            gflops = device.dp_gflops
        elif dtype == paddle.float32:
            gflops = device.sp_gflops
        elif dtype == paddle.float16 or dtype == paddle.bfloat16:
            gflops = device.hp_gflops
        else:
            raise ValueError(
                f"Unsupported modeling compute time for dtype: {dtype}."
            )

        utilization_rate = 0.98
        time = flops_count / (utilization_rate * gflops) * 1e-3

    # calc comm op time by communication modeling formula
    elif op_type in COMM_OP_TYPE:
        op_cost = _g_op_cost_factory[op_type](
            op=op, comm_context=CommContext(cluster)
        )
        time = op_cost.calc_time()

    else:
        raise ValueError(f"The {op_type} has not been supported now.")

    return time