chore: import upstream snapshot with attribution

python/paddle/distributed/passes/pass_base.py

@@ -0,0 +1,378 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+
+from paddle.framework import _apply_pass as _apply_cpp_pass
+
+
+class PassContext:
+    def __init__(self):
+        self._applied_passes = []
+        self._attrs = {}
+
+    def set_attr(self, key, value):
+        self._attrs[key] = value
+
+    def get_attr(self, key, default=None):
+        return self._attrs.get(key, default)
+
+    @property
+    def passes(self):
+        return self._applied_passes
+
+    def _add_pass(self, pass_obj):
+        self._applied_passes.append(pass_obj)
+
+    def _pop_pass(self):
+        del self._applied_passes[-1]
+
+
+class PassType:
+    UNKNOWN = 0
+    COMM_OPT = 1
+    CALC_OPT = 2
+    PARALLEL_OPT = 3
+    FUSION_OPT = 4
+
+
+class PassBase(ABC):
+    _REGISTERED_PASSES = {}
+    _COMMON_RULES = []
+
+    _BEFORE_WHITE_LISTS_DICT = {}
+    _AFTER_WHITE_LISTS_DICT = {}
+    _PASS_PROCESS_ORDER_LIST = []
+
+    name = None
+
+    @staticmethod
+    def _register(pass_name, pass_class):
+        assert issubclass(pass_class, PassBase)
+        PassBase._REGISTERED_PASSES[pass_name] = pass_class
+
+    def __init__(self):
+        self._attrs = {}
+
+    def set_attr(self, key, value):
+        self._attrs[key] = value
+        return self
+
+    def get_attr(self, key, default=None):
+        return self._attrs.get(key, default)
+
+    @abstractmethod
+    def _check_self(self):
+        pass
+
+    @abstractmethod
+    def _check_conflict(self, other_pass):
+        pass
+
+    def _type(self):
+        return PassType.UNKNOWN
+
+    def _check_conflict_including_common_rules(self, other_pass):
+        return self._check_conflict(other_pass) and all(
+            r(other_pass, self) for r in PassBase._COMMON_RULES
+        )
+
+    def apply(self, main_programs, startup_programs, context=None):
+        if context is None:
+            context = PassContext()
+
+        if not self._check_self():
+            return context
+
+        if not all(
+            self._check_conflict_including_common_rules(p)
+            for p in context.passes
+        ):
+            return context
+
+        assert isinstance(main_programs, list)
+        assert isinstance(startup_programs, list)
+        assert len(main_programs) == len(startup_programs)
+        self._apply_impl(main_programs, startup_programs, context)
+        context._add_pass(self)
+        return context
+
+    def _apply_impl(self, main_programs, startup_programs, context):
+        for main_program, startup_program in zip(
+            main_programs, startup_programs
+        ):
+            self._apply_single_impl(main_program, startup_program, context)
+
+    @abstractmethod
+    def _apply_single_impl(self, main_program, startup_program, context):
+        pass
+
+
+def register_pass(name):
+    def impl(cls):
+        PassBase._register(name, cls)
+        cls.name = name
+        return cls
+
+    return impl
+
+
+def new_pass(name, pass_attrs={}):
+    pass_class = PassBase._REGISTERED_PASSES.get(name)
+    assert pass_class is not None, f"Pass {name} is not registered"
+    pass_obj = pass_class()
+    for k, v in pass_attrs.items():
+        pass_obj.set_attr(k, v)
+    return pass_obj
+
+
+class CPPPassWrapper(PassBase):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        raise NotImplementedError
+
+    @property
+    def cpp_attr_types(self):
+        return {}
+
+    def _check_self(self):
+        return True
+
+    def _check_conflict(self, other_pass):
+        return True
+
+    def _apply_single_impl(self, main_program, startup_program, context):
+        _apply_cpp_pass(
+            main_program,
+            startup_program,
+            self.cpp_name,
+            self._attrs,
+            self.cpp_attr_types,
+        )
+
+
+def _fusion_opt_last_rule(pass_before, pass_after):
+    if (
+        pass_before._type() == PassType.FUSION_OPT
+        and pass_after._type() != PassType.FUSION_OPT
+    ):
+        return False
+    else:
+        return True
+
+
+def _fusion_opt_list_rule(pass_before, pass_after):
+    if (
+        pass_before._type() == PassType.FUSION_OPT
+        and pass_after._type() == PassType.FUSION_OPT
+    ):
+        return _get_list_index(pass_before) < _get_list_index(pass_after)
+    else:
+        return True
+
+
+def _make_rule_from_white_lists_dict(
+    before_white_lists_dict, after_white_lists_dict
+):
+    def collect_pass_names(white_lists_dict, result):
+        for k, v in white_lists_dict.items():
+            result.add(k)
+            assert isinstance(v, (list, tuple))
+            for pass_name in v:
+                assert isinstance(pass_name, (bytes, str))
+                result.add(pass_name)
+
+    all_pass_names = set()
+    collect_pass_names(before_white_lists_dict, all_pass_names)
+    collect_pass_names(after_white_lists_dict, all_pass_names)
+
+    compatible_pass_dict = {}
+    for pass_name in all_pass_names:
+        compatible_pass_dict[pass_name] = set()
+
+    for k, v in before_white_lists_dict.items():
+        for pass_name in v:
+            compatible_pass_dict[k].add(pass_name)
+
+    for k, v in after_white_lists_dict.items():
+        for pass_name in v:
+            compatible_pass_dict[pass_name].add(k)
+
+    def rule(pass_before, pass_after):
+        all_passes_after = compatible_pass_dict.get(pass_before.name)
+        if (
+            all_passes_after is None
+            or pass_after.name not in compatible_pass_dict
+        ):
+            return True
+        else:
+            return pass_after.name in all_passes_after
+
+    return rule
+
+
+def _get_list_index(in_pass):
+    assert in_pass.name in PassBase._PASS_PROCESS_ORDER_LIST, (
+        f"Pass {in_pass.name} is not in _PASS_PROCESS_ORDER_LIST"
+    )
+    return PassBase._PASS_PROCESS_ORDER_LIST.index(in_pass.name)
+
+
+# The key-value pair (k, [v1, v2, ..., vn]) means the pass k can be
+# applied before any of pass [v1, v2, ..., vn] is applied
+PassBase._BEFORE_WHITE_LISTS_DICT = {
+    "fuse_gradient_merge": ["fuse_all_reduce"],
+    # Add more white lists here
+}
+
+# The key-value pair (k, [v1, v2, ..., vn]) means the pass k can be
+# applied after any of pass [v1, v2, ..., vn] is applied
+PassBase._AFTER_WHITE_LISTS_DICT = {
+    # Add more white lists here
+}
+
+# The index of pass in this list represent the order in which the pass is processed.
+PassBase._PASS_PROCESS_ORDER_LIST = [
+    "fuse_resunit",
+    "fuse_relu_depthwise_conv",
+    "fuse_bn_add_act",
+    "fuse_bn_act",
+    "fused_attention",
+    "fused_feedforward",
+    "fuse_gemm_epilogue",
+    "fuse_adamw",
+    "fuse_optimizer",
+]
+
+PassBase._COMMON_RULES = [
+    _fusion_opt_last_rule,
+    _fusion_opt_list_rule,
+    lambda pass_before, pass_after: type(pass_before) != type(pass_after),
+    _make_rule_from_white_lists_dict(
+        PassBase._BEFORE_WHITE_LISTS_DICT, PassBase._AFTER_WHITE_LISTS_DICT
+    ),
+    # Add more common rules here
+]
+
+
+def _find_longest_path(edges):
+    n = len(edges)
+    paths = [None] * n
+    dists = [None] * n
+
+    min_path = []
+    min_dist = 0
+    for i in range(n):
+        paths[i] = [None] * n
+        dists[i] = [None] * n
+        for j in range(n):
+            assert isinstance(edges[i][j], bool)
+            if not edges[i][j]:
+                dists[i][j] = n  # inf
+                paths[i][j] = []
+            else:
+                assert edges[i][j] is True
+                dists[i][j] = -1
+                paths[i][j] = [i, j]
+                if dists[i][j] < min_dist:
+                    min_dist = -1
+                    min_path = paths[i][j]
+
+    for k in range(n):
+        for i in range(n):
+            for j in range(n):
+                if dists[i][j] > dists[i][k] + dists[k][j]:
+                    dists[i][j] = dists[i][k] + dists[k][j]
+                    if paths[i][k]:
+                        assert paths[i][k][-1] == k
+                    else:
+                        continue
+                    if paths[k][j]:
+                        assert paths[k][j][0] == k
+                    else:
+                        continue
+                    paths[i][j] = (
+                        paths[i][k] + paths[k][j][1:] if paths[k][j] else []
+                    )
+                    if dists[i][j] < min_dist:
+                        min_dist = dists[i][j]
+                        min_path = paths[i][j]
+
+    return min_path if min_path else [0]
+
+
+def _solve_pass_conflict(passes, context):
+    passes = [p for p in passes if p._check_self()]
+    if not passes:
+        return []
+
+    old_passes = passes
+    passes = []
+    for p in old_passes:
+        if all(
+            p._check_conflict_including_common_rules(applied_p)
+            for applied_p in context.passes
+        ):
+            passes.append(p)
+
+    if not passes:
+        return []
+
+    n = len(passes)
+    adjacent_matrix = []
+    for _ in range(n):
+        adjacent_matrix.append([None] * n)
+
+    for i in range(n):
+        for j in range(n):
+            adjacent_matrix[i][j] = passes[
+                j
+            ]._check_conflict_including_common_rules(passes[i])
+
+    longest_path = _find_longest_path(adjacent_matrix)
+    return [passes[idx] for idx in longest_path]
+
+
+class PassManager:
+    def __init__(self, passes, context=None, auto_solve_conflict=True):
+        if context is None:
+            context = PassContext()
+        self._context = context
+
+        if auto_solve_conflict:
+            self._passes = _solve_pass_conflict(passes, context)
+        else:
+            self._passes = list(passes)
+
+    def apply(self, main_programs, startup_programs):
+        context = self._context
+        for p in self._passes:
+            context = p.apply(main_programs, startup_programs, context)
+        self._context = context
+        return context
+
+    @property
+    def context(self):
+        return self._context
+
+    @property
+    def names(self):
+        return [p.name for p in self.passes]
+
+    @property
+    def passes(self):
+        return tuple(self._passes)