chore: import upstream snapshot with attribution

python/paddle/distributed/passes/cpp_pass.py

@@ -0,0 +1,251 @@
+# 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 paddle.framework import (
+    _apply_pass as _apply_cpp_pass,
+    core,
+)
+from paddle.static import Executor
+
+from .pass_base import CPPPassWrapper, PassType, register_pass
+
+
+@register_pass("fuse_elewise_add_act")
+class FuseElementwiseAddActPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_elewise_add_act_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_bn_act")
+class FuseBatchNormActPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_bn_act_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_bn_add_act")
+class FuseBatchNormAddActPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_bn_add_act_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_relu_depthwise_conv")
+class FuseReluDepthwiseConvPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_relu_depthwise_conv_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fused_attention")
+class FusedAttentionPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fused_attention_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fused_feedforward")
+class FusedFeedforwardPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fused_feedforward_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_gemm_epilogue")
+class FuseGemmEpiloguePass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_gemm_epilogue_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_adamw")
+class FuseAdamWPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_adamw_op_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_dot_product_attention")
+class FuseDotProductAttentionPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_dot_product_attention_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+@register_pass("fuse_optimizer")
+class FuseOptimizerPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return [
+            "fuse_adam_op_pass",
+            "fuse_sgd_op_pass",
+            "fuse_momentum_op_pass",
+        ]
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+# @register_pass("inplace_addto_op")
+# class InplaceAddtoOpPass(CPPPassWrapper):
+#     def __init__(self):
+#         super().__init__()
+
+#     @property
+#     def cpp_name(self):
+#         return "inplace_addto_op_pass"
+
+#     def _type(self):
+#         return PassType.CALC_OPT
+
+
+@register_pass("fuse_resunit")
+class FuseResUnitPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def cpp_name(self):
+        return "fuse_resunit_pass"
+
+    def _type(self):
+        return PassType.FUSION_OPT
+
+
+def _set_cinn_op_flag(flag_name, extra_ops):
+    values = core.globals()[flag_name]
+    values = [v.strip() for v in values.split(";") if v.strip()]
+    values.extend(extra_ops)
+    core.globals()[flag_name] = ";".join(values)
+
+
+@register_pass("build_cinn")
+class BuildCINNPass(CPPPassWrapper):
+    def __init__(self):
+        super().__init__()
+        self.set_attr("allow_ops", [])
+        self.set_attr("deny_ops", [])
+
+    @property
+    def cpp_name(self):
+        return "build_cinn_pass"
+
+    def _type(self):
+        return PassType.CALC_OPT
+
+    def _apply_single_impl(self, main_program, startup_program, context):
+        assert 'FLAGS_allow_cinn_ops' in core.globals(), (
+            "PaddlePaddle is not compiled with CINN support"
+        )
+        old_allow_ops = core.globals()['FLAGS_allow_cinn_ops']
+        old_deny_ops = core.globals()['FLAGS_deny_cinn_ops']
+        try:
+            _set_cinn_op_flag(
+                'FLAGS_allow_cinn_ops', self.get_attr("allow_ops")
+            )
+            _set_cinn_op_flag('FLAGS_deny_cinn_ops', self.get_attr("deny_ops"))
+
+            feed = self.get_attr('feed', [])
+            fetch_list = self.get_attr('fetch_list', [])
+            prune_program = self.get_attr('prune_program', True)
+
+            if prune_program:
+                tmp_main_program = Executor._prune_program(
+                    main_program, feed, fetch_list, []
+                )
+
+                tmp_main_program = Executor._add_fetch_ops(
+                    tmp_main_program, fetch_list, 'fetch'
+                )
+
+            else:
+                tmp_main_program = Executor._add_fetch_ops(
+                    main_program, fetch_list, 'fetch'
+                )
+
+            _apply_cpp_pass(
+                tmp_main_program,
+                startup_program,
+                self.cpp_name,
+                {},
+                self.cpp_attr_types,
+            )
+
+            tmp_main_program = Executor._remove_fetch_ops(tmp_main_program)
+
+            tmp_main_program = core.ProgramDesc(tmp_main_program.desc)
+
+            main_program._rebuild_from_desc(tmp_main_program)
+
+        finally:
+            core.globals()['FLAGS_allow_cinn_ops'] = old_allow_ops
+            core.globals()['FLAGS_deny_cinn_ops'] = old_deny_ops