[PirPass]Refine ApplyPirPass logic in @to_static (#59702)

Aurelius84 · web-flow · commit c4bc9e15e929 · 2023-12-29T20:24:47.000+08:00
* [PirPass]Refine ApplyPirPass logic in @to_static

* fix shared_ptr

* fix code

* del usless code

* fix codestyle

* fix UT
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
@@ -211,7 +211,7 @@ class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
   }
 
  private:
-  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_{nullptr};
 };
 
 class CinnGroupLoweringPass : public pir::PatternRewritePass {
@@ -237,7 +237,7 @@ class CinnGroupLoweringPass : public pir::PatternRewritePass {
   }
 
  private:
-  const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis_;
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_{nullptr};
 };
 
 }  // namespace
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -520,6 +520,7 @@ std::vector<ir::Expr> OpLowererImpl::LowerOps(
   auto& strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");
   std::vector<Expr> func_bodies;
   for (auto* op : ops) {
+    VLOG(4) << "start lowering op:" << op->name();
     // 1.Select Op impl
     std::vector<ir::Tensor> op_func_arg_tensors =
         CollectInputTensor(group, op, group_func_arg_tensors, tensor_map);
@@ -891,6 +892,7 @@ ir::Tensor OpLowererImpl::GetTensor(const GroupPtr& group,
   auto in_shape = ::common::vectorize<int>(type_info.dims());
   auto dtype = type_info.dtype();
   std::string input_id = ValueName(value);
+  VLOG(3) << "group->shape_analysis:" << group->shape_analysis;
   if (group->shape_analysis != nullptr) {
     auto sym_vec =
         group->shape_analysis->GetOrCreateSymbolicDimsForRankedValue(value);
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1607,45 +1607,44 @@ static bool HasDynamicShape(const Program &program) {
   return false;
 }
 
-void ApplyPirPass(Program &forward_program) {  // NOLINT
+void AddCinnPass(std::shared_ptr<PassManager> &pass_manager,  // NOLINT
+                 Program &program) {                          // NOLINT
 #ifdef PADDLE_WITH_CINN
   pir::IrContext *ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
   ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
   ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
 
-  bool has_dynamic_shape = HasDynamicShape(forward_program);
+  bool has_dynamic_shape = HasDynamicShape(program);
 
   auto shape_analysis =
       has_dynamic_shape ? std::make_shared<pir::ShapeConstraintIRAnalysis>(ctx)
                         : nullptr;
 
-  pir::PassManager pass_manager(ctx);
-  pass_manager.AddPass(pir::CreateShapeOptimizationPass());
-  cinn::dialect::ir::PdOp2CinnOpConverter(&forward_program);
+  pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+  cinn::dialect::ir::PdOp2CinnOpConverter(&program);
 
-  pass_manager.AddPass(
+  pass_manager->AddPass(
       std::make_unique<cinn::dialect::ir::AddBroadcastToElementwisePass>());
-  pass_manager.AddPass(pir::CreateDeadCodeEliminationPass());
-  pass_manager.AddPass(pir::CreateBuildCinnPass());
+  pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
+  pass_manager->AddPass(pir::CreateBuildCinnPass());
 
   if (has_dynamic_shape) {
-    pass_manager.AddPass(pir::CreateInferSymbolicShapePass(shape_analysis));
+    pass_manager->AddPass(pir::CreateInferSymbolicShapePass(shape_analysis));
   }
 
-  pass_manager.AddPass(
+  pass_manager->AddPass(
       cinn::dialect::ir::CreateCinnGroupLoweringPass(shape_analysis));
-
-  pass_manager.Run(&forward_program);
-  VLOG(3) << "after BuildCinnPass, forward_program:\n" << forward_program;
+  VLOG(4) << "has_dynamic_shape :" << has_dynamic_shape
+          << ", shape_analysis: " << shape_analysis;
 #else
   PADDLE_THROW(platform::errors::Unimplemented(
       "Currently we only support CINN Pass for Pir under @to_static, please "
       "compile PaddlePaddle with CINN"));
 #endif
 }
 void BindIrPass(pybind11::module *m) {
-  m->def("apply_pir_pass", ApplyPirPass);
+  m->def("add_cinn_pass", AddCinnPass);
 
   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
@@ -280,11 +280,12 @@ def apply_pir_program_pass(self, pass_fn):
         def pass_fn(forward_program, backward_program):
             return forward_program, backward_program
         """
-        program_name_attr = self.program_name_attr
         origin_fwd = self.forward_program
         origin_bwd = self.backward_program
+        # NOTE(dev): Add this line to trigger program_name_attr logic
+        program_name_attr = self.program_name_attr
         self.forward_program, self.backward_program = pass_fn(
-            self.forward_program, self.backward_program, program_name_attr
+            origin_fwd, origin_bwd
         )
 
     # cached property can ensure program is splited only once.
@@ -382,55 +383,6 @@ def backward_program(self):
         return self._forward_backward_program[0][1]
 
 
-class PirPassContext:
-    """
-    PirPassContext is a class that only has staticmethod currently.
-    It will create a new RunableProgram after calling apply method.
-    """
-
-    INPUT_OP_NAME = "pd_op.data"
-    PARAM_OP_NAME = "builtin.parameter"
-    OUTPUT_OP_NAME = "builtin.shadow_output"
-
-    @classmethod
-    def apply(cls, runable_program, build_strategy):
-        # TODO(Aurelius84): Currently only support infer mode,
-        # and we just use forward_program because backward_program
-        # is empty.
-        if not build_strategy.build_cinn_pass:
-            return runable_program
-        elif not paddle.is_compiled_with_cinn():
-            raise RuntimeError(
-                "Please install PaddlePaddle compiled with CINN while setting build_strategy.build_cinn_pass = True."
-            )
-        fwd_program, _ = paddle.base.libpaddle.pir.clone_program(
-            runable_program.forward_program
-        )
-        paddle.base.libpaddle.pir.apply_pir_pass(fwd_program)
-        in_out_values = cls._prepare_attr(fwd_program)
-        return RunableProgram(fwd_program, in_out_values)
-
-    @classmethod
-    def _prepare_attr(cls, program):
-        """
-        After applying Pass, we need to update the Input/Parameter/Output Value
-        that refer to the new program.
-
-        NOTE: We assume that Inputs come from INPUT_OP, Params come from
-              PARM_OP and Output come from OUTPUT_OP.
-        """
-        inputs, params, outputs = [], [], []
-        for op in program.global_block().ops:
-            op_name = op.name()
-            if op_name == cls.INPUT_OP_NAME:
-                inputs.append(op.result(0))
-            elif op_name == cls.PARAM_OP_NAME:
-                params.append(op.result(0))
-            elif op_name == cls.OUTPUT_OP_NAME:
-                outputs.append(op.operand(0).source())
-        return inputs, params, outputs
-
-
 class PartialProgramLayerHook:
     def before_append_backward(self, forward_program, src_vars):
         ...
@@ -596,37 +548,40 @@ def _get_scope(self, program_id=None, use_scope_cache=False):
     @switch_to_static_graph
     def _create_program(self, is_infer_mode=False):
         if is_infer_mode:
+
+            def pass_fn(forward_program, backward_program):
+                pm = paddle.base.libpaddle.pir.PassManager()
+                if self._build_strategy.build_cinn_pass:
+                    paddle.base.libpaddle.pir.add_cinn_pass(pm, forward_program)
+                    pm.run(forward_program)
+                return forward_program, backward_program
+
             # TODO(xiongkun) who to transfer the pruning program?
             infer_program = self.origin_runable_program.clone()
             if self._hooker:
                 self._hooker.after_infer(infer_program)
-            infer_program = PirPassContext.apply(
-                infer_program, self._build_strategy
-            )
+            infer_program.apply_pir_program_pass(pass_fn)
             return infer_program
         else:
             train_program: RunableProgram = self.origin_runable_program.clone()
             train_program = self._append_backward_desc(train_program)
             # Note: Only set grad type once after initializing train program. So we put it here.
             self._set_grad_type(self._params, train_program)
 
-            # (NOTE:@xiongkun) HOW TO APPLY PASS: this is a example for forward/backward clone pass, just replace with your cases.
-            def pass_fn(forward_program, backward_program, name_attr):
-                fwd, _ = paddle.base.libpaddle.pir.clone_program(
-                    forward_program
-                )
-
-                if self._build_strategy.build_cinn_pass:
-                    paddle.base.libpaddle.pir.apply_pir_pass(fwd)
-
-                bwd, _ = paddle.base.libpaddle.pir.clone_program(
-                    backward_program
-                )
+            def pass_fn(forward_program, backward_program):
+                fwd_pm = paddle.base.libpaddle.pir.PassManager()
+                bwd_pm = paddle.base.libpaddle.pir.PassManager()
 
                 if self._build_strategy.build_cinn_pass:
-                    paddle.base.libpaddle.pir.apply_pir_pass(bwd)
-
-                return fwd, bwd
+                    paddle.base.libpaddle.pir.add_cinn_pass(
+                        fwd_pm, forward_program
+                    )
+                    paddle.base.libpaddle.pir.add_cinn_pass(
+                        bwd_pm, backward_program
+                    )
+                    fwd_pm.run(forward_program)
+                    bwd_pm.run(backward_program)
+                return forward_program, backward_program
 
             train_program.apply_pir_program_pass(pass_fn)
             return train_program
@@ -748,7 +703,7 @@ def _insert_aggregation_ops_for_var(target_program, var):
                 shape=var.shape,
             )
             # step2: rename the var.name@GRAD to var.name@GRAD@dy2static
-            for idx, op in finded_ops:
+            for _, op in finded_ops:
                 op._rename_input(var_grad_name, new_grad_name)
                 op._rename_output(var_grad_name, new_grad_name)
             # step3: insert sum op to aggregate the gradient.
diff --git a/test/ir/pir/cinn/test_cinn_sub_graph.py b/test/ir/pir/cinn/test_cinn_sub_graph.py
@@ -203,10 +203,9 @@ def test_forward(self):
         cinn_out = self.train(use_cinn=True)
         dy_out = self.train(use_cinn=False)
 
-        # TODO(zhangliujie) fix precision error
-        # np.testing.assert_allclose(
-        #     cinn_out.numpy(), dy_out.numpy(), atol=1e-8, rtol=1e-4
-        # )
+        np.testing.assert_allclose(
+            cinn_out.numpy(), dy_out.numpy(), atol=1e-8, rtol=1e-4
+        )
 
 
 class TestCinnDropout(TestCinnSubGraphBase):

Original file line number	Diff line number	Diff line change
`@@ -211,7 +211,7 @@ class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {`
`211`	`211`	`}`
`212`	`212`
`213`	`213`	`private:`
`214`		`- std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;`
	`214`	`+ std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_{nullptr};`
`215`	`215`	`};`
`216`	`216`
`217`	`217`	`class CinnGroupLoweringPass : public pir::PatternRewritePass {`
`@@ -237,7 +237,7 @@ class CinnGroupLoweringPass : public pir::PatternRewritePass {`
`237`	`237`	`}`
`238`	`238`
`239`	`239`	`private:`
`240`		`- const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis_;`
	`240`	`+ std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_{nullptr};`
`241`	`241`	`};`
`242`	`242`
`243`	`243`	`} // namespace`