[PIR] fix the insertion point error when create parameter in subblock. (#59858)

winter-wang · web-flow · commit 94899e6a798b · 2023-12-12T15:42:32.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.cc b/paddle/fluid/pir/dialect/operator/ir/api_builder.cc
@@ -27,19 +27,15 @@ ApiBuilder::ApiBuilder()
 
 void ApiBuilder::SetProgram(pir::Program* program) {
   IR_ENFORCE(program != nullptr, "argument of program is nullptr");
-  builder_->SetInsertionPointToEnd(program->block());
-}
-
-void ApiBuilder::set_insertion_point(pir::Operation* op) {
-  builder_->set_insertion_point(op);
+  builder_->SetInsertionPointToBlockEnd(program->block());
 }
 
 void ApiBuilder::ResetInsertionPointToStart() {
   builder_->SetInsertionPointToStart(builder_->block());
 }
 
 void ApiBuilder::ResetInsertionPointToEnd() {
-  builder_->SetInsertionPointToEnd(builder_->block());
+  builder_->SetInsertionPointToBlockEnd(builder_->block());
 }
 
 pir::Parameter* ApiBuilder::GetParameter(const std::string& name) const {
@@ -53,16 +49,10 @@ void ApiBuilder::SetParameter(const std::string& name,
   program->SetParameter(name, std::move(parameter));
 }
 
-void ApiBuilder::PushInsertionPoint(
-    const pir::InsertionPoint& insertion_point) {
-  insertion_point_stack_.push(this->insertion_point());
-  set_insertion_point(insertion_point);
-}
-
-void ApiBuilder::PopInsertionPoint() {
+void ApiBuilder::LoadInsertionPoint() {
   IR_ENFORCE(!insertion_point_stack_.empty(),
              "insertion_point_stack_ is empty.");
-  set_insertion_point(insertion_point_stack_.top());
+  builder_->set_insertion_point(insertion_point_stack_.top());
   insertion_point_stack_.pop();
 }
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.h b/paddle/fluid/pir/dialect/operator/ir/api_builder.h
@@ -34,10 +34,6 @@ class ApiBuilder {
   }
   void SetProgram(pir::Program* program);
 
-  /// Set the insertion point to the specified operation, which will cause
-  /// subsequent insertions to go right before it.
-  void set_insertion_point(pir::Operation* op);
-
   void ResetInsertionPointToStart();
 
   void ResetInsertionPointToEnd();
@@ -49,15 +45,31 @@ class ApiBuilder {
 
   std::shared_ptr<pir::Builder> GetBuilder() { return builder_; }
 
-  const pir::InsertionPoint& insertion_point() const {
+  const pir::InsertionPoint& GetCurrentInsertionPoint() const {
     return builder_->insertion_point();
   }
 
-  void set_insertion_point(const pir::InsertionPoint& insertion_point) {
+  /// Set the insertion point to the specified insertion_point.
+  void SetInsertionPoint(const pir::InsertionPoint& insertion_point) {
     builder_->set_insertion_point(insertion_point);
   }
-  void PushInsertionPoint(const pir::InsertionPoint& insertion_point);
-  void PopInsertionPoint();
+
+  /// Set the insertion point to the specified operation, which will cause
+  /// subsequent insertions to go right before it.
+  void SetInsertionPoint(pir::Operation* op) {
+    builder_->set_insertion_point(op);
+  }
+  /// Set the insertion point to the end of specified block.
+  void SetInsertionPointToBlockEnd(pir::Block* block) {
+    builder_->SetInsertionPointToBlockEnd(block);
+  }
+
+  // push current insertion point to the stack.
+  void PushInsertionPoint() {
+    insertion_point_stack_.push(builder_->insertion_point());
+  }
+  // pop the insertion point and set it to the current insertion point.
+  void LoadInsertionPoint();
 
  private:
   ApiBuilder();
diff --git a/paddle/fluid/pir/transforms/sub_graph_detector.cc b/paddle/fluid/pir/transforms/sub_graph_detector.cc
@@ -507,7 +507,7 @@ void ReplaceWithGroupOp(pir::Block* block,
   }
 
   // step 4: Insert YieldOp for outputs
-  builder.SetInsertionPointToEnd(group_block);
+  builder.SetInsertionPointToBlockEnd(group_block);
   builder.Build<::pir::YieldOp>(outputs);
 }
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -316,13 +316,14 @@ void BindBlock(py::module *m) {
       .def(
           "__enter__",
           [](Block &self) -> Block & {
-            ApiBuilder::Instance().PushInsertionPoint({&self, self.end()});
+            ApiBuilder::Instance().PushInsertionPoint();
+            ApiBuilder::Instance().SetInsertionPointToBlockEnd(&self);
             return self;
           },
           return_value_policy::reference)
       .def("__exit__",
            [](Block &self, py::object, py::object, py::object) {
-             ApiBuilder::Instance().PopInsertionPoint();
+             ApiBuilder::Instance().LoadInsertionPoint();
            })
       .def("__len__", [](Block &self) { return self.size(); })
       .def("args", &Block::args, return_value_policy::reference)
@@ -849,6 +850,14 @@ void BindAttribute(py::module *m) {
   });
 }
 
+struct PyInsertionPoint {
+  pir::InsertionPoint value;
+};
+void BindInsertionPoint(pybind11::module *m) {
+  py::class_<PyInsertionPoint> ir_insertion_point(*m, "InsertionPoint", R"DOC(
+    InsertionPoint class represents the insertion point in the Builder.)DOC");
+}
+
 Operation *BuildOpFrom(
     Operation *to_copy_op,
     std::unordered_map<pir::Value, pir::Value> &value_map) {  // NOLINT
@@ -1345,25 +1354,24 @@ void BindUtils(pybind11::module *m) {
   m->def("append_set_parameters", AppendSetParameters);
   m->def("fake_op_result", FakeOpResult);
   m->def("is_fake_op_result", IsFakeOpResult);
-  m->def("set_global_program",
-         [](Program *program) { ApiBuilder::Instance().SetProgram(program); });
-  m->def(
-      "cur_insertion_point",
-      []() { return ApiBuilder::Instance().insertion_point(); },
-      return_value_policy::reference);
-  m->def("set_insertion_point", [](const pir::InsertionPoint &insertion_point) {
-    ApiBuilder::Instance().set_insertion_point(insertion_point);
+  m->def("get_current_insertion_point", []() -> PyInsertionPoint {
+    return {ApiBuilder::Instance().GetCurrentInsertionPoint()};
+  });
+  m->def("set_insertion_point", [](const PyInsertionPoint &insertion_point) {
+    ApiBuilder::Instance().SetInsertionPoint(insertion_point.value);
   });
   m->def("set_insertion_point",
-         [](Operation *op) { ApiBuilder::Instance().set_insertion_point(op); });
+         [](Operation *op) { ApiBuilder::Instance().SetInsertionPoint(op); });
+  m->def("set_insertion_point_to_block_end", [](Block *block) {
+    ApiBuilder::Instance().SetInsertionPointToBlockEnd(block);
+  });
   m->def("reset_insertion_point_to_start",
          []() { ApiBuilder::Instance().ResetInsertionPointToStart(); });
   m->def("reset_insertion_point_to_end",
          []() { ApiBuilder::Instance().ResetInsertionPointToEnd(); });
   m->def("register_paddle_dialect", []() {
     pir::IrContext::Instance()
         ->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
-    pir::IrContext::Instance()->GetOrRegisterDialect<pir::ControlFlowDialect>();
   });
   m->def("create_selected_rows_type_by_dense_tensor",
          CreateSelectedRowsTypeByDenseTensor);
@@ -1620,6 +1628,7 @@ void BindPir(pybind11::module *module) {
   BindOpResult(&ir_module);
   BindType(&ir_module);
   BindAttribute(&ir_module);
+  BindInsertionPoint(&ir_module);
   BindUtils(&ir_module);
   BindIrPass(&ir_module);
   BindPassManager(&ir_module);
diff --git a/paddle/pir/core/builder.h b/paddle/pir/core/builder.h
@@ -103,7 +103,8 @@ class Builder {
   }
 
   /// Set the insertion point to the end of the specified block.
-  void SetInsertionPointToEnd(Block *block) {
+  void SetInsertionPointToBlockEnd(Block *block) {
+    IR_ENFORCE(block != nullptr, "argument of block is nullptr");
     set_insertion_point(block, block->end());
   }
 
diff --git a/python/paddle/pir/__init__.py b/python/paddle/pir/__init__.py
@@ -28,8 +28,9 @@
     register_paddle_dialect,
     reset_insertion_point_to_end,
     reset_insertion_point_to_start,
-    set_global_program,
+    get_current_insertion_point,
     set_insertion_point,
+    set_insertion_point_to_block_end,
     translate_to_pir,
     translate_to_pir_with_param_map,
     create_shaped_type,
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -17,7 +17,12 @@
 
 from paddle.base.core import VarDesc
 from paddle.base.libpaddle import DataType
-from paddle.base.libpaddle.pir import Program, set_global_program
+from paddle.base.libpaddle.pir import (
+    Program,
+    get_current_insertion_point,
+    set_insertion_point,
+    set_insertion_point_to_block_end,
+)
 
 from .._pir_ops import parameter, set_parameter
 from ..base import unique_name
@@ -83,7 +88,7 @@ def convert_np_dtype_to_dtype_(np_dtype):
 # program is a global instance.
 _main_program_ = Program()
 # set the global program for c++ and this program will be used to build ops in c++
-set_global_program(_main_program_)
+set_insertion_point_to_block_end(_main_program_.global_block())
 
 _startup_program_ = Program()
 
@@ -153,7 +158,7 @@ def default_main_program():
     return _main_program_
 
 
-def switch_main_program(program):
+def switch_main_program(program, insertion_point=None):
     """
     Switch the main program to a new program.
 
@@ -165,9 +170,13 @@ def switch_main_program(program):
     """
     global _main_program_
     prev_program = _main_program_
+    prev_insertion_point = get_current_insertion_point()
     _main_program_ = program
-    set_global_program(_main_program_)
-    return prev_program
+    if insertion_point is None:
+        set_insertion_point_to_block_end(_main_program_.global_block())
+    else:
+        set_insertion_point(insertion_point)
+    return prev_program, prev_insertion_point
 
 
 def switch_startup_program(program):
@@ -234,7 +243,7 @@ def program_guard(main_program, startup_program=None):
     check_type(
         main_program, 'main_program', Program, 'paddle.static.program_guard'
     )
-    main_program = switch_main_program(main_program)
+    main_program, prev_insertion_point = switch_main_program(main_program)
     if startup_program is not None:
         check_type(
             startup_program,
@@ -246,7 +255,7 @@ def program_guard(main_program, startup_program=None):
     try:
         yield
     finally:
-        switch_main_program(main_program)
+        switch_main_program(main_program, prev_insertion_point)
         if startup_program is not None:
             switch_startup_program(startup_program)
 
diff --git a/test/cpp/new_executor/standalone_executor_pir_test.cc b/test/cpp/new_executor/standalone_executor_pir_test.cc
@@ -251,7 +251,7 @@ TEST(StandaloneExecutor, if_op) {
   builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_2.out()});
 
   std::string out_name = "if_out";
-  builder.SetInsertionPointToEnd(block);
+  builder.SetInsertionPointToBlockEnd(block);
   builder.Build<pir::ShadowOutputOp>(if_op->result(0), out_name);
 
   auto kernel_program = paddle::dialect::PdOpLowerToKernelPass(&program);
diff --git a/test/cpp/pir/cinn/group_op_test.cc b/test/cpp/pir/cinn/group_op_test.cc
@@ -58,16 +58,16 @@ std::shared_ptr<::pir::Program> BuildGroupProgram() {
   auto group_op1 = builder.Build<cinn::dialect::GroupOp>(
       CreateDenseTensorTypes(common::make_ddim(shape)));
   pir::Block* block1 = group_op1.block();
-  builder.SetInsertionPointToEnd(block1);
+  builder.SetInsertionPointToBlockEnd(block1);
   auto full_op_x = builder.Build<paddle::dialect::FullOp>(
       shape, value_one, phi::DataType::FLOAT32, phi::GPUPlace());
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{full_op_x.out()});
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto group_op2 = builder.Build<cinn::dialect::GroupOp>(
       CreateDenseTensorTypes(common::make_ddim(shape)));
   pir::Block* block2 = group_op2.block();
-  builder.SetInsertionPointToEnd(block2);
+  builder.SetInsertionPointToBlockEnd(block2);
 
   auto tan_op_x = builder.Build<paddle::dialect::TanOp>(group_op1->result(0));
   auto relu_op_x = builder.Build<paddle::dialect::ReluOp>(tan_op_x->result(0));
@@ -109,24 +109,24 @@ std::shared_ptr<::pir::Program> BuildGroupProgramByBlock() {
   const float value_one = 1.0;
   const std::vector<int64_t> shape = {64, 128};
   std::unique_ptr<::pir::Block> block1(new ::pir::Block());
-  builder.SetInsertionPointToEnd(block1.get());
+  builder.SetInsertionPointToBlockEnd(block1.get());
   auto full_op_x = builder.Build<paddle::dialect::FullOp>(
       shape, value_one, phi::DataType::FLOAT32, phi::GPUPlace());
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{full_op_x.out()});
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto group_op1 = builder.Build<cinn::dialect::GroupOp>(std::move(block1));
 
   // ------- Group op 2 ---------
   std::unique_ptr<::pir::Block> block2(new ::pir::Block());
-  builder.SetInsertionPointToEnd(block2.get());
+  builder.SetInsertionPointToBlockEnd(block2.get());
   auto tan_op_x = builder.Build<paddle::dialect::TanOp>(group_op1->result(0));
   auto relu_op_x = builder.Build<paddle::dialect::ReluOp>(tan_op_x->result(0));
   auto tan_op_y = builder.Build<paddle::dialect::TanOp>(relu_op_x->result(0));
   auto relu_op_y = builder.Build<paddle::dialect::ReluOp>(tan_op_y->result(0));
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{relu_op_y.out()});
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto group_op2 = builder.Build<cinn::dialect::GroupOp>(std::move(block2));
 
   return program;
@@ -170,31 +170,31 @@ std::shared_ptr<::pir::Program> BuildGroupProgramForLowering() {
   auto group_op1 = builder.Build<cinn::dialect::GroupOp>(
       CreateDenseTensorTypes(common::make_ddim(shape)));
   pir::Block* block1 = group_op1.block();
-  builder.SetInsertionPointToEnd(block1);
+  builder.SetInsertionPointToBlockEnd(block1);
   auto sin = builder.Build<paddle::dialect::SinOp>(full_x->result(0));
 
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{
       sin.out(),
   });
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto group_op2 = builder.Build<cinn::dialect::GroupOp>(
       CreateDenseTensorTypes(common::make_ddim(shape)));
   pir::Block* block2 = group_op2.block();
-  builder.SetInsertionPointToEnd(block2);
+  builder.SetInsertionPointToBlockEnd(block2);
   auto cos_op = builder.Build<paddle::dialect::CosOp>(full_y->result(0));
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{cos_op.out()});
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto group_op3 = builder.Build<cinn::dialect::GroupOp>(
       CreateDenseTensorTypes(common::make_ddim(shape)));
   pir::Block* block3 = group_op3.block();
-  builder.SetInsertionPointToEnd(block3);
+  builder.SetInsertionPointToBlockEnd(block3);
   auto add = builder.Build<paddle::dialect::AddOp>(group_op1->result(0),
                                                    group_op2->result(0));
   builder.Build<::pir::YieldOp>(std::vector<::pir::Value>{add.out()});
 
-  builder.SetInsertionPointToEnd(program->block());
+  builder.SetInsertionPointToBlockEnd(program->block());
   auto exp = builder.Build<paddle::dialect::ExpOp>(group_op3->result(0));
 
   builder.Build<paddle::dialect::FetchOp>(exp.out(), "out", 0);
diff --git a/test/cpp/pir/cinn/pir_compiler_test.cc b/test/cpp/pir/cinn/pir_compiler_test.cc
@@ -188,7 +188,7 @@ TEST(PirCompier, CompileSoftmax) {
 
   new_program->block()->push_back(cinn_op);
 
-  builder.SetInsertionPointToEnd(new_program->block());
+  builder.SetInsertionPointToBlockEnd(new_program->block());
   builder.Build<paddle::dialect::FetchOp>(
       cinn_op->result(cinn_op->num_results() - 1), "out", 0);
 
diff --git a/test/cpp/pir/control_flow_dialect/if_op_test.cc b/test/cpp/pir/control_flow_dialect/if_op_test.cc
@@ -94,7 +94,7 @@ TEST(if_op_test, build_by_block) {
       std::vector<int64_t>{2}, true, phi::DataType::BOOL);
   builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_2.out()});
 
-  builder.SetInsertionPointToEnd(block);
+  builder.SetInsertionPointToBlockEnd(block);
 
   auto if_op = builder.Build<paddle::dialect::IfOp>(
       full_op.out(), std::move(true_block), std::move(false_block));
@@ -149,7 +149,7 @@ TEST(if_op_test, network_with_backward) {
                                   std::initializer_list<pir::Value>{local2_z});
   builder.Build<pir::YieldOp>(std::vector<pir::Value>{local2_w});
 
-  builder.SetInsertionPointToEnd(block);
+  builder.SetInsertionPointToBlockEnd(block);
 
   // build backward network
   auto out_grad = builder.Build<FullOp>(std::vector<int64_t>{2, 2}, 1.0f).out();
@@ -201,7 +201,7 @@ TEST(if_op_test, network_with_backward) {
   builder.Build<pir::YieldOp>(
       std::vector<pir::Value>{local2_x_grad, local2_y_grad});
 
-  builder.SetInsertionPointToEnd(block);
+  builder.SetInsertionPointToBlockEnd(block);
 
   LOG(INFO) << program;
 
diff --git a/test/cpp/pir/core/block_operand_test.cc b/test/cpp/pir/core/block_operand_test.cc
diff --git a/test/cpp/pir/shape_dialect/shape_struct_test.cc b/test/cpp/pir/shape_dialect/shape_struct_test.cc

Original file line number	Diff line number	Diff line change
`@@ -507,7 +507,7 @@ void ReplaceWithGroupOp(pir::Block* block,`
`507`	`507`	`}`
`508`	`508`
`509`	`509`	`// step 4: Insert YieldOp for outputs`
`510`		`- builder.SetInsertionPointToEnd(group_block);`
	`510`	`+ builder.SetInsertionPointToBlockEnd(group_block);`
`511`	`511`	`builder.Build<::pir::YieldOp>(outputs);`
`512`	`512`	`}`
`513`	`513`
Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,8 @@ class Builder {`
`103`	`103`	`}`
`104`	`104`
`105`	`105`	`/// Set the insertion point to the end of the specified block.`
`106`		`- void SetInsertionPointToEnd(Block *block) {`
	`106`	`+ void SetInsertionPointToBlockEnd(Block *block) {`
	`107`	`+ IR_ENFORCE(block != nullptr, "argument of block is nullptr");`
`107`	`108`	`set_insertion_point(block, block->end());`
`108`	`109`	`}`
`109`	`110`