base_cost.py 29.4 KB
Newer Older
C
caozhou 已提交
1
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15
#
# 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
16 17
from functools import reduce

18 19
import paddle

20 21
from ..cluster import LinkType
from ..dist_tensor import DistributedTensor
22 23
from ..process_group import get_process_group
from ..utils import _get_comm_group, _get_idx_in_axis
24

25
COMM_OP_TYPE = [
26 27 28 29 30 31
    "send_v2",
    "recv_v2",
    "c_broadcast",
    "c_allgather",
    "c_allreduce_sum",
    "c_identity",
32 33
]
NON_COMP_TYPE = ["while"] + COMM_OP_TYPE
C
caozhou 已提交
34
_g_op_cost_factory = {}
35 36


37 38 39 40
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
41

42
    desc = {}
43
    # 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}}}
44
    vars = op.block.vars
45
    desc["op"] = op.type
46 47 48 49 50
    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:
51 52
            var = get_var_with_recursion(var_name, op.block, op.block.program)
            shape = var.shape
53 54 55 56 57 58 59 60 61
            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:
62 63
            var = get_var_with_recursion(var_name, op.block, op.block.program)
            shape = var.shape
64 65 66 67 68 69 70 71 72 73
            var_desc.append((var.dtype, shape))
        output_desc[out_name] = var_desc
    desc["outputs"] = output_desc

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

    return desc


74 75 76 77 78 79 80 81 82
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."
83
    processes = process_mesh.process_ids
84 85 86 87 88 89 90 91 92 93 94 95 96 97
    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
        for input_name in op.input_names:
            var_name_list = op.input(input_name)
            var_desc = []
            for var_name in var_name_list:
98 99 100
                var = get_var_with_recursion(
                    var_name, op.block, op.block.program
                )
101 102 103 104 105
                # 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
106
                topology = process_mesh.shape
107
                shape = DistributedTensor.get_local_sizes(
108 109 110 111 112 113 114
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
115 116 117
                var_desc.append((var.dtype, shape))

                # For special op such as embedding and its grad op
118 119 120 121 122 123
                if (
                    op.type == "c_embedding"
                    or op.type == "lookup_table_v2"
                    or op.type == "c_embedding_grad"
                    or op.type == "lookup_table_v2_grad"
                ):
124
                    if input_name == "W":
125 126 127 128 129
                        embedding_row_dim_mapping = (
                            dist_attr.get_input_dims_mapping(
                                op.input(input_name)[0]
                            )[0]
                        )
130
                        relative_idx = _get_idx_in_axis(
131
                            processes,
132
                            dist_attr.process_mesh.shape,
133 134 135
                            embedding_row_dim_mapping,
                            process,
                        )
136 137 138 139 140 141 142 143 144 145 146 147
                        per_part_size = shape[0]
                        relative_idx = relative_idx * per_part_size
                        desc["attrs"]["start_index"] = relative_idx

            input_desc[input_name] = var_desc
        desc["inputs"] = input_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
148 149 150
                var = get_var_with_recursion(
                    var_name, op.block, op.block.program
                )
151 152 153 154 155
                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
156 157
                processes = process_mesh.process_ids
                topology = process_mesh.shape
158
                shape = DistributedTensor.get_local_sizes(
159 160 161 162 163 164 165
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
166 167 168 169 170 171 172
                var_desc.append((var.dtype, shape))

                # 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)
173
                    process_mesh_shape = dist_attr.process_mesh.shape
174 175 176 177
                    shape_list = op.attr("shape")
                    # Modify target shape
                    for idx, axis in enumerate(dims_mapping):
                        if axis >= 0:
178 179 180
                            shape_list[idx] = (
                                shape_list[idx] // process_mesh_shape[axis]
                            )
181 182 183 184 185 186 187 188 189 190 191 192
                    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.
193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232
    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("Unsupported dtype {}".format(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)
233 234 235 236 237
    attrs = desc["attrs"]
    parse_result = (desc_str, attrs)
    return parse_result


238 239 240 241 242 243 244 245 246
def build_comm_desc_from_dist_op(
    op_type,
    dist_op,
    ctx,
    var_names,
    attrs=None,
    parallel_axis=None,
    group_ranks=None,
):
247 248 249 250 251 252 253 254 255
    """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."

256
    processes = process_mesh.process_ids
257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
    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
287 288 289
                var = get_var_with_recursion(
                    var_name, serial_op.block, serial_op.block.program
                )
290

291 292 293 294 295
                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)
                )
296 297
                global_sizes = var.shape
                shard_sizes = None
298
                topology = process_mesh.shape
299
                shape = DistributedTensor.get_local_sizes(
300 301 302 303 304 305 306
                    global_sizes,
                    dims_mapping,
                    topology,
                    processes,
                    process,
                    shard_sizes,
                )
307 308 309 310 311 312 313
                input_list.append((var.dtype, shape))

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

            # Get comm group by parallel_axis or the given group_ranks.
            if parallel_axis is not None:
314 315
                process_mesh_shape = process_mesh.shape
                process_mesh_group = process_mesh.process_ids
316 317 318 319 320 321
                comm_group_ranks = _get_comm_group(
                    process_mesh_group,
                    process_mesh_shape,
                    parallel_axis,
                    rank_id,
                )
322 323 324 325 326 327 328 329 330 331 332 333
            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 = {}
334

335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362
            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
    desc["inputs"] = {"X": [(dtype, shape)]}
    desc["attrs"] = attrs
    return desc


def build_comm_costs_from_descs(op_cost_class, ctx, processes, descs, cluster):
    """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)
363 364 365
            comm_op_cost = op_cost_class(
                op_desc=desc, comm_context=comm_context
            )
366 367 368 369 370 371 372 373 374 375 376 377
            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)
    return costs


378 379 380
def build_dp_costs(
    result, dist_op, ctx, var_names, attrs, parallel_axis, cluster
):
381 382 383 384 385 386
    """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
387
    processes = process_mesh.process_ids
388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412
    assert len(var_names) == 1
    vars = dist_op.serial_op.block.vars
    var_name = var_names[0]
    has_found = False
    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
                break
    if not has_found:
        return

    c_allreduce_sum_descs = build_comm_desc_from_dist_op(
        "c_allreduce_sum",
        dist_op,
        ctx,
        var_names,
        attrs=attrs,
413 414
        parallel_axis=parallel_axis,
    )
415
    comm_cost_list = build_comm_costs_from_descs(
416 417 418 419 420 421
        _g_op_cost_factory["c_allreduce_sum"],
        ctx,
        processes,
        c_allreduce_sum_descs,
        cluster,
    )
422 423 424 425 426 427 428 429 430 431 432 433
    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"] = {}
434 435 436 437 438
            if var_name in dist_attr.inputs_dist_attrs:
                dims_mapping = dist_attr.get_input_dims_mapping(var_name)
            elif var_name in dist_attr.outputs_dist_attrs:
                dims_mapping = dist_attr.get_output_dims_mapping(var_name)
            else:
439 440 441 442
                raise AssertionError(
                    "cannot find dims_mapping for {} in {}".format(
                        var_name, dist_attr
                    )
443 444 445 446 447 448 449
                )

            # dims_mapping = (
            #     dist_attr.get_input_dims_mapping(var_name)
            #     if dist_attr.get_input_dims_mapping(var_name) is not None
            #     else dist_attr.get_output_dims_mapping(var_name)
            # )
450 451 452 453 454
            var = get_var_with_recursion(
                var_name,
                dist_op.serial_op.block,
                dist_op.serial_op.block.program,
            )
455 456
            global_sizes = var.shape
            shard_sizes = None
457
            topology = process_mesh.shape
458 459 460 461 462 463 464 465
            shape = DistributedTensor.get_local_sizes(
                global_sizes,
                dims_mapping,
                topology,
                processes,
                rank,
                shard_sizes,
            )
466 467 468
            desc["inputs"]["X"] = [(var.dtype, shape)]
            attrs = {"scale": 1.0 / dp_degree}
            desc["attrs"] = attrs
469 470 471
            scale_op_cost = _g_op_cost_factory["scale"](
                op_desc=desc, cluster=cluster
            )
472 473
            scale_costs[rank] = scale_op_cost
        result.append(scale_costs)
474 475 476 477 478 479 480 481


class CommContext:
    _instance = None
    _has_instance = False

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
C
caozhou 已提交
482
            cls._instance = super().__new__(cls)
483 484 485
            _has_instance = True
        return cls._instance

C
caozhou 已提交
486 487 488 489 490
    def __init__(self, cluster):
        if CommContext._has_instance:
            return
        self.beta = {}
        self.hops = {}
C
caozhou 已提交
491
        assert cluster is not None
C
caozhou 已提交
492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509
        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.base_inter_ring = None
        # self.base_inter_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
510
            self.base_tree = 0.0
511 512
            # self.base_inter_ring = 9.6
            # self.base_inter_tree = 28
C
caozhou 已提交
513 514 515 516 517 518 519 520 521 522 523 524
            # 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
525
            self.base_tree = base_tree if base_tree is not None else 0.0
C
caozhou 已提交
526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580

            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)
581 582
        key = ','.join(map(str, sorted(ranks)))
        max_beta = None
C
caozhou 已提交
583 584
        if key in self.beta:
            max_beta = self.beta[key]
585 586 587
        else:
            for i in range(len(ranks)):
                for j in range(i + 1, len(ranks)):
588
                    forward_order_beta = self.cluster.get_beta(
589 590
                        ranks[i], ranks[j]
                    )
591
                    backward_order_beta = self.cluster.get_beta(
592 593 594 595 596 597 598
                        ranks[j], ranks[i]
                    )
                    beta = (
                        forward_order_beta
                        if forward_order_beta > backward_order_beta
                        else backward_order_beta
                    )
599
                    if max_beta is None:
C
caozhou 已提交
600
                        max_beta = beta
601 602 603
                    else:
                        if beta > max_beta:
                            max_beta = beta
C
caozhou 已提交
604
            self.beta[key] = max_beta
605 606 607

        return max_beta

C
caozhou 已提交
608 609 610 611 612 613 614 615 616 617 618
    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

619 620 621 622 623 624 625 626 627 628 629

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):
630 631 632
        assert (
            isinstance(val, int) and val >= 0
        ), "Memory must be int and greater than equal to 0."
633 634

    def _check_flops(self, val):
635 636 637
        assert (
            isinstance(val, int) and val >= 0
        ), "FLOPs must be int and greater than equal to 0."
638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670

    @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
671
        assert time >= 0 and memory >= 0 and flops >= 0
672 673 674 675 676 677 678
        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
679
        assert time >= 0 and memory >= 0 and flops >= 0
680 681 682 683 684 685 686
        return Cost(time, memory, flops)


class OpCost:
    def __init__(self, op=None, op_desc=None):
        self._op = op
        self._op_desc = op_desc
C
caozhou 已提交
687
        self._cost = None
688 689 690 691 692 693 694 695 696

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

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

C
caozhou 已提交
697 698 699 700 701 702 703 704 705 706 707 708
    @property
    def time(self):
        return self.cost.time

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

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

709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728
    @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

C
caozhou 已提交
729 730 731 732 733 734 735 736 737
    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
738
            assert time >= 0 and memory >= 0 and flops >= 0
C
caozhou 已提交
739 740 741 742
        elif isinstance(rhs, Cost):
            time = self.time + rhs.time
            memory = self.memory + rhs.memory
            flops = self.flops + rhs.flops
743
            assert time >= 0 and memory >= 0 and flops >= 0
C
caozhou 已提交
744 745 746 747 748 749 750 751 752 753 754
        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
755
            assert time >= 0 and memory >= 0 and flops >= 0
C
caozhou 已提交
756 757 758 759
        elif isinstance(rhs, Cost):
            time = self.time - rhs.time
            memory = self.memory - rhs.memory
            flops = self.flops - rhs.flops
760
            assert time >= 0 and memory >= 0 and flops >= 0
C
caozhou 已提交
761 762
        return Cost(time, memory, flops)

763 764 765 766 767

class CommOpCost(OpCost):
    OP_TYPE = "COMM"

    def __init__(self, op=None, op_desc=None, comm_context=None):
768
        super().__init__(op=op, op_desc=op_desc)
769 770
        self._check_comm_op_type()
        self._comm_context = comm_context
C
caozhou 已提交
771 772 773 774 775 776
        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()
777 778 779 780 781

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

C
caozhou 已提交
782 783
    @property
    def comm_count(self):
784 785
        from ..reshard import get_var_with_recursion

C
caozhou 已提交
786 787 788 789 790 791 792
        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 overrided
                var_name = self.op.input("X")[0]
793 794 795
                var = get_var_with_recursion(
                    var_name, self.op.block, self.program
                )
C
caozhou 已提交
796 797 798 799 800 801 802 803 804 805 806 807 808 809 810
                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
811 812
            elif dtype == paddle.bool:
                factor = 8
C
caozhou 已提交
813
            else:
814
                raise ValueError("Unsupported comm dtype {}".format(dtype))
C
caozhou 已提交
815 816 817 818 819 820 821 822 823 824 825 826 827 828
            comm_count = reduce(lambda x, y: x * y, 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(
829 830
                self.group_ranks
            )
C
caozhou 已提交
831 832 833 834 835 836 837 838 839 840 841 842 843 844
        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:
845
                ring_id = self.op.attrs("ring_id")
C
caozhou 已提交
846 847 848
                process_group = get_process_group(ring_id)
                if process_group is None:
                    raise ValueError(
849 850 851 852
                        "There not exists process group whose ring_id is {}.".format(
                            ring_id
                        )
                    )
C
caozhou 已提交
853 854 855
                self._group_ranks = process_group.ranks
        return self._group_ranks

856 857 858 859
    @classmethod
    def _check_comm_op_type(cls):
        if cls.OP_TYPE != "COMM":
            if cls.OP_TYPE not in COMM_OP_TYPE:
860 861
                raise TypeError(
                    "Please Check op type in {}, but got {}.".format(
862 863 864
                        COMM_OP_TYPE, cls.OP_TYPE
                    )
                )
865 866 867 868 869 870


class CompOpCost(OpCost):
    OP_TYPE = "COMP"

    def __init__(self, op=None, op_desc=None, cluster=None):
871
        super().__init__(op=op, op_desc=op_desc)
872
        self._check_comp_op_type()
C
caozhou 已提交
873
        self._cost = self.calc_cost()
874 875 876 877 878 879
        self.cluster = cluster

    @classmethod
    def _check_comp_op_type(cls):
        if cls.OP_TYPE != "COMP":
            if cls.OP_TYPE in NON_COMP_TYPE:
880 881
                raise TypeError(
                    "Please Check op type not in {}, but got {}.".format(
882 883 884
                        NON_COMP_TYPE, cls.OP_TYPE
                    )
                )
885 886 887 888 889 890


def register_op_cost(cls):
    op_type = cls.OP_TYPE

    def register(op_type):
C
caozhou 已提交
891 892
        global _g_op_cost_factory
        _g_op_cost_factory[op_type] = cls
893

C
caozhou 已提交
894 895
    register(op_type)
    return cls
896 897


C
caozhou 已提交
898
def calc_time_by_modeling(op=None, desc=None, cluster=None):
899 900
    op_type = op.type if op is not None else desc["op"]
    if op_type in COMM_OP_TYPE:
901 902 903
        op_cost = _g_op_cost_factory[op_type](
            op=op, op_desc=desc, comm_context=CommContext(cluster)
        )
904
    elif op_type not in NON_COMP_TYPE:
905 906 907
        op_cost = _g_op_cost_factory[op_type](
            op=op, op_desc=desc, cluster=cluster
        )
908 909
    time = op_cost.calc_time()
    return time