[PIR] Print CSE program when FLAGS_print_ir enabled (#64695)

SigureMo · web-flow · commit a299904cb3f6 · 2024-05-30T15:02:21.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1813,18 +1813,6 @@ std::unordered_map<::pir::Value, std::string> GetNameMap(
   return value2name;
 }
 
-std::shared_ptr<Program> ApplyFusedBnAddActPass(
-    std::shared_ptr<Program> program) {
-  pir::PassManager pm(pir::IrContext::Instance(), 3);
-  pm.AddPass(pir::CreateFusedBnAddActPass());
-  pm.Run(program.get());
-  if (FLAGS_print_ir) {
-    std::cout << "IR After FusedBnAddActPass -------------------" << std::endl;
-    std::cout << *program << std::endl;
-  }
-  return program;
-}
-
 SplitedResult SplitForwardBackward(
     const Program &program,
     const std::vector<pir::Value> &forward_inputs,
@@ -2160,7 +2148,6 @@ void BindUtils(pybind11::module *m) {
   m->def("get_op_inplace_info", GetOpInplaceInfo);
   m->def("reset_shadow_output_name", ResetShadowOutputName);
   m->def("split_program", SplitForwardBackward);
-  m->def("apply_bn_add_act_pass", ApplyFusedBnAddActPass);
   m->def("append_shadow_outputs", AppendShadowOutputs);
   m->def("append_shadow_output", AppendShadowOutput);
   m->def("fake_value", FakeValue);
@@ -2391,18 +2378,38 @@ void InferSymbolicShapePass(
   }
 }
 
-void CommonSubexpressionEliminationPass(
-    std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
-    pir::Program &program) {                          // NOLINT
-  pass_manager->AddPass(pir::CreateCommonSubexpressionEliminationPass());
+std::shared_ptr<Program> ApplyCommonSubexpressionEliminationPass(
+    std::shared_ptr<Program> program) {
+  pir::PassManager pm(pir::IrContext::Instance(), 2);
+  pm.AddPass(pir::CreateCommonSubexpressionEliminationPass());
+  pm.Run(program.get());
+  if (FLAGS_print_ir) {
+    std::cout
+        << "IR After CommonSubexpressionEliminationPass -------------------"
+        << std::endl;
+    std::cout << *program << std::endl;
+  }
+  return program;
+}
+
+std::shared_ptr<Program> ApplyFusedBnAddActPass(
+    std::shared_ptr<Program> program) {
+  pir::PassManager pm(pir::IrContext::Instance(), 3);
+  pm.AddPass(pir::CreateFusedBnAddActPass());
+  pm.Run(program.get());
+  if (FLAGS_print_ir) {
+    std::cout << "IR After FusedBnAddActPass -------------------" << std::endl;
+    std::cout << *program << std::endl;
+  }
+  return program;
 }
 
 void BindIrPass(pybind11::module *m) {
   m->def("apply_cinn_pass", ApplyCinnPass);
   m->def("check_infer_symbolic_if_need", CheckInferSymbolicIfNeed);
   m->def("infer_symbolic_shape_pass", InferSymbolicShapePass);
-  m->def("common_subexpression_elimination_pass",
-         CommonSubexpressionEliminationPass);
+  m->def("apply_cse_pass", ApplyCommonSubexpressionEliminationPass);
+  m->def("apply_bn_add_act_pass", ApplyFusedBnAddActPass);
 
   py::class_<Pass, std::shared_ptr<Pass>> pass(*m,
                                                "Pass",
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -30,7 +30,12 @@
 from paddle.pir import Value, fake_value, is_fake_value
 
 from .logging_utils import TranslatorLogger
-from .utils import RETURN_NO_VALUE_MAGIC_NUM, backend_guard, cinn_is_enabled
+from .utils import (
+    RETURN_NO_VALUE_MAGIC_NUM,
+    backend_guard,
+    cinn_is_enabled,
+    cse_is_enabled,
+)
 
 __all__ = []
 
@@ -729,17 +734,6 @@ def _get_scope(self, program_id=None, use_scope_cache=False):
     # whole
     @switch_to_static_graph
     def _create_program(self, is_infer_mode=False):
-        def apply_cse(program):
-            if not paddle.get_flags(["FLAGS_enable_cse_in_dy2st"])[
-                "FLAGS_enable_cse_in_dy2st"
-            ]:
-                return
-            pm = paddle.base.libpaddle.pir.PassManager()
-            paddle.base.libpaddle.pir.common_subexpression_elimination_pass(
-                pm, program
-            )
-            pm.run(program)
-
         if is_infer_mode:
 
             def pass_fn(forward_program, backward_program):
@@ -749,7 +743,8 @@ def pass_fn(forward_program, backward_program):
                     pm, forward_program
                 )
                 pm.run(forward_program)
-                apply_cse(forward_program)
+                if cse_is_enabled():
+                    paddle.base.libpaddle.pir.apply_cse_pass(forward_program)
 
                 # if-else pass
                 if cinn_is_enabled(self._build_strategy, self._backend):
@@ -776,8 +771,9 @@ def pass_fn(forward_program, backward_program):
             self._set_grad_type(self._params, train_program)
 
             def pass_fn(forward_program, backward_program):
-                apply_cse(forward_program)
-                apply_cse(backward_program)
+                if cse_is_enabled():
+                    paddle.base.libpaddle.pir.apply_cse_pass(forward_program)
+                    paddle.base.libpaddle.pir.apply_cse_pass(backward_program)
                 if cinn_is_enabled(self._build_strategy, self._backend):
                     paddle.base.libpaddle.pir.apply_cinn_pass(forward_program)
                     paddle.base.libpaddle.pir.apply_cinn_pass(backward_program)
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -545,6 +545,12 @@ def cinn_is_enabled(build_strategy, backend):
     return False
 
 
+def cse_is_enabled():
+    return paddle.get_flags(["FLAGS_enable_cse_in_dy2st"])[
+        "FLAGS_enable_cse_in_dy2st"
+    ]
+
+
 def prim_is_enabled():
     core.check_and_set_prim_all_enabled()
     return core._is_bwd_prim_enabled() or core._is_fwd_prim_enabled()
@@ -588,8 +594,8 @@ def backend_guard(backend):
 
 def construct_grad_names(grad_info_map, x_vars, param_vars, out_vars):
     grad_var_names = {}
-    fn = (
-        lambda grad_var: grad_var.name
+    fn = lambda grad_var: (
+        grad_var.name
         if isinstance(grad_var, framework.Variable)
         else framework.EMPTY_VAR_NAME
     )
diff --git a/test/ir/pir/test_cse_pass.py b/test/ir/pir/test_cse_pass.py
@@ -31,12 +31,6 @@ def program_scope_guard():
             yield main_program
 
 
-def apply_cse(program):
-    pm = paddle.base.libpaddle.pir.PassManager()
-    paddle.base.libpaddle.pir.common_subexpression_elimination_pass(pm, program)
-    pm.run(program)
-
-
 def walk_block(block, fn):
     for op in block.ops:
         fn(op)
@@ -81,7 +75,7 @@ def test_basic(self):
             b = x1 + x2
 
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.add": 1})
 
     def test_basic2(self):
@@ -100,7 +94,7 @@ def expr(x1, x2, x3):
             out2 = expr(x1, x2, x3)
 
             self.assert_op_count_equal(main_program, {"pd_op.add": 4})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
 
     def test_replace_full_with_same_attr(self):
@@ -110,7 +104,7 @@ def test_replace_full_with_same_attr(self):
             x2 = paddle.full([2, 2], 0, dtype="float32")
 
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.full": 1})
 
     def test_skip_replace_full_with_different_value(self):
@@ -120,7 +114,7 @@ def test_skip_replace_full_with_different_value(self):
             x2 = paddle.full([2, 2], 1, dtype="float32")
 
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
 
     def test_skip_replace_full_with_different_dtype(self):
@@ -130,7 +124,7 @@ def test_skip_replace_full_with_different_dtype(self):
             x2 = paddle.full([2, 2], 0, dtype="int32")
 
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
 
     def test_skip_replace_full_with_different_shape(self):
@@ -140,7 +134,7 @@ def test_skip_replace_full_with_different_shape(self):
             x2 = paddle.full([3, 3], 0, dtype="float32")
 
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.full": 2})
 
     def test_complex_computation(self):
@@ -172,7 +166,7 @@ def expr(x1, x2, x3):
                     "pd_op.subtract": repeat_n,
                 },
             )
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(
                 main_program,
                 {
@@ -199,7 +193,7 @@ def test_commutative(self):
             b = x2 + x1
 
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.add": 1})
 
     def test_complex_computation_commutative(self):
@@ -247,7 +241,7 @@ def expr(x1, x2, x3):
                     "pd_op.bitwise_and": repeat_n * expr_per_repeat,
                 },
             )
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(
                 main_program,
                 {
@@ -293,7 +287,7 @@ def get_cond(loop_var):
             paddle.static.nn.while_loop(get_cond, loop_body, [loop_var])
 
             self.assert_op_count_equal(main_program, {"pd_op.multiply": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.multiply": 1})
 
     def test_nested_subblock(self):
@@ -334,7 +328,7 @@ def outer_get_cond(loop_var):
             )
 
             self.assert_op_count_equal(main_program, {"pd_op.multiply": 3})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.multiply": 1})
 
 
@@ -347,7 +341,7 @@ def test_can_not_replace_random(self):
             x2 = paddle.rand([2, 2], dtype="float32")
 
             self.assert_op_count_equal(main_program, {"pd_op.uniform": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.uniform": 2})
 
     def test_can_not_replace_print(self):
@@ -360,7 +354,7 @@ def test_can_not_replace_print(self):
             paddle.static.Print(x1)
 
             self.assert_op_count_equal(main_program, {"pd_op.print": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.print": 2})
 
     def test_can_not_replace_used_by_inplace_op(self):
@@ -376,7 +370,7 @@ def test_can_not_replace_used_by_inplace_op(self):
             paddle.assign(paddle.full([2, 2], 0, dtype="float32"), b)
 
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
-            apply_cse(main_program)
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
 
 
