add comm cost (#42784)

python/paddle/distributed/auto_parallel/cost/base_cost.py

-#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# 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.
@@ -153,6 +153,7 @@ class CommContext:
             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
@@ -338,8 +339,6 @@ class Cost:
 
 class OpCost:
     def __init__(self, op=None, op_desc=None):
-        assert (op is not None and op_desc is None) or (op is None and
-                                                        op_desc is not None)
         self._op = op
         self._op_desc = op_desc
         self._cost = None

python/paddle/distributed/auto_parallel/cost/comp_op_cost.py

@@ -12,7 +12,628 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
-from .base_cost import Cost, register_op_cost, CompOpCost
+from .base_cost import Cost, register_op_cost, CompOpCost, _g_op_cost_factory
+
+
+@register_op_cost
+class AssignOpCost(CompOpCost):
+    OP_TYPE = "assign"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(AssignOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class AssignValueOpCost(CompOpCost):
+    OP_TYPE = "assign_value"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(AssignValueOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class BeamSearchOpCost(CompOpCost):
+    OP_TYPE = "beam_search"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(BeamSearchOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class BeamSearchDecodeOpCost(CompOpCost):
+    OP_TYPE = "beam_search_decode"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(BeamSearchDecodeOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class CastOpCost(CompOpCost):
+    OP_TYPE = "cast"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(CastOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ConcatOpCost(CompOpCost):
+    OP_TYPE = "concat"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ConcatOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseAddOpCost(CompOpCost):
+    OP_TYPE = "elementwise_add"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseAddOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseAddGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_add_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseAddGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseDivOpCost(CompOpCost):
+    OP_TYPE = "elementwise_div"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseDivOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseDivGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_div_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseDivGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseMulOpCost(CompOpCost):
+    OP_TYPE = "elementwise_mul"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseMulOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseMulGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_mul_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseMulGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class ElementwiseSubOpCost(CompOpCost):
+    OP_TYPE = "elementwise_sub"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(ElementwiseSubOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class EmbeddingOpCost(CompOpCost):
+    OP_TYPE = "c_embedding"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(EmbeddingOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class EmbeddingGradOpCost(CompOpCost):
+    OP_TYPE = "c_embedding_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(EmbeddingGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class FillConstantOpCost(CompOpCost):
+    OP_TYPE = "fill_constant"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(FillConstantOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class FillConstantBatchSizeLikeOpCost(CompOpCost):
+    OP_TYPE = "fill_constant_batch_size_like"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(FillConstantBatchSizeLikeOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class FillConstantBatchSizeLikeGradOpCost(CompOpCost):
+    OP_TYPE = "fill_constant_batch_size_like_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(FillConstantBatchSizeLikeGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class GatherOpCost(CompOpCost):
+    OP_TYPE = "gather"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(GatherOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class GeluOpCost(CompOpCost):
+    OP_TYPE = "gelu"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(GeluOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class GeluGradOpCost(CompOpCost):
+    OP_TYPE = "gelu_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(GeluGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class GreaterEqualOpCost(CompOpCost):
+    OP_TYPE = "greater_equal"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(GreaterEqualOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class IncrementOpCost(CompOpCost):
+    OP_TYPE = "increment"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(IncrementOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class IsEmptyOpCost(CompOpCost):
+    OP_TYPE = "is_empty"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(IsEmptyOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LayerNormOpCost(CompOpCost):
+    OP_TYPE = "layer_norm"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LayerNormOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LayerNormGradOpCost(CompOpCost):
+    OP_TYPE = "layer_norm_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LayerNormGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LessThanOpCost(CompOpCost):
+    OP_TYPE = "less_than"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LessThanOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LogicalNotOpCost(CompOpCost):
+    OP_TYPE = "logical_not"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LogicalNotOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LogicalAndOpCost(CompOpCost):
+    OP_TYPE = "logical_and"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LogicalAndOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LodResetOpCost(CompOpCost):
+    OP_TYPE = "lod_reset"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LodResetOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LogOpCost(CompOpCost):
+    OP_TYPE = "log"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LogOpCost, self).__init__(op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LookupTableV2OpCost(CompOpCost):
+    OP_TYPE = "lookup_table_v2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LookupTableV2OpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class LookupTableV2GradOpCost(CompOpCost):
+    OP_TYPE = "lookup_table_v2_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(LookupTableV2GradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class MatmulOpCost(CompOpCost):
+    OP_TYPE = "matmul"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(MatmulOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+
+@register_op_cost
+class MatmulGradOpCost(CompOpCost):
+    OP_TYPE = "matmul_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None):
+        super(MatmulGradOpCost, self).__init__(
+            op=op, op_desc=op_desc, cluster=cluster)
+
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
+    def calc_flops(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
+
+    def calc_time(self):
+        # NOTE: The actual formula will be filled in the future
+        return 0
 
 
 @register_op_cost
@@ -23,7 +644,7 @@ class MatmulV2OpCost(CompOpCost):
         super(MatmulV2OpCost, self).__init__(
             op=op, op_desc=op_desc, cluster=cluster)
 
-    # For a concrete COMP OP, the calc_time and calc_flops function needs to be overrided
+    #  For a concrete COMP OP, the calc_time and calc_flops function need to be overrided
     def calc_flops(self):
         # NOTE: The actual formula will be filled in the future
         return 0

python/paddle/fluid/tests/unittests/auto_parallel/CMakeLists.txt

@@ -30,4 +30,5 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     py_test_modules(test_dist_slice MODULES test_dist_slice ENVS ${dist_ENVS})
     py_test_modules(test_cluster MODULES test_cluster ENVS ${dist_ENVS})
     py_test_modules(test_comm_cost MODULES test_comm_cost ENVS ${dist_ENVS})
+    py_test_modules(test_comp_cost MODULES test_comp_cost ENVS ${dist_ENVS})
 endif()

python/paddle/fluid/tests/unittests/auto_parallel/test_comm_cost.py

-# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 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.

python/paddle/fluid/tests/unittests/auto_parallel/test_comp_cost.py

+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import json
+
+import paddle
+from paddle.distributed.auto_parallel.cluster import Cluster
+from paddle.distributed.auto_parallel.cost.comp_op_cost import AssignOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import AssignValueOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import BeamSearchOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import BeamSearchDecodeOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import CastOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ConcatOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseAddOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseAddGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseDivOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseDivGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseMulOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseMulGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import ElementwiseSubOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import EmbeddingOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import EmbeddingGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import FillConstantOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import FillConstantBatchSizeLikeOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import FillConstantBatchSizeLikeGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import GatherOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import GeluOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import GeluGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import GreaterEqualOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import IncrementOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import IsEmptyOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LayerNormOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LayerNormGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LessThanOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LogicalNotOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LogicalAndOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LodResetOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LogOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LookupTableV2OpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import LookupTableV2GradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import MatmulOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import MatmulGradOpCost
+from paddle.distributed.auto_parallel.cost.comp_op_cost import MatmulV2OpCost
+
+from test_cluster import cluster_json
+
+
+class TestCompOpCost(unittest.TestCase):
+    def test_comp_cost(self):
+        # Build cluster
+        file_dir = os.path.dirname(os.path.abspath(__file__))
+        cluster_json_path = os.path.join(file_dir, "auto_parallel_cluster.json")
+        cluster_json_object = json.loads(cluster_json)
+        with open(cluster_json_path, "w") as cluster_json_file:
+            json.dump(cluster_json_object, cluster_json_file)
+        cluster = Cluster()
+        cluster.build_from_file(cluster_json_path)
+
+        op_cost = AssignOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = AssignValueOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = BeamSearchOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = BeamSearchDecodeOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = CastOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ConcatOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseAddOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseAddGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseDivOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseDivGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseMulOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseMulGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = ElementwiseSubOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = EmbeddingOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = EmbeddingGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = FillConstantOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = FillConstantBatchSizeLikeOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = FillConstantBatchSizeLikeGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = GatherOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = GeluOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = GeluGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = GreaterEqualOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = IncrementOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = IsEmptyOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LayerNormOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LayerNormGradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LessThanOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LogicalNotOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LogicalAndOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LodResetOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LogOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LookupTableV2OpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = LookupTableV2GradOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = MatmulOpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        op_cost = MatmulV2OpCost(cluster=cluster)
+        self.assertTrue(op_cost.flops >= 0)
+        self.assertTrue(op_cost.time >= 0)
+        self.assertTrue(op_cost.memory >= 0)
+
+        # Remove unnecessary files
+        if os.path.exists(cluster_json_path):
+            os.remove(cluster_json_path)
+
+
+if __name__ == "__main__":
+    unittest.main()