[PIR] polish the pir interface. (#59392)

winter-wang · web-flow · commit 4061ba744f9e · 2023-11-28T02:45:34.000+08:00
diff --git a/paddle/fluid/framework/new_executor/instruction/while_instruction.cc b/paddle/fluid/framework/new_executor/instruction/while_instruction.cc
@@ -74,7 +74,7 @@ WhileInstruction::WhileInstruction(size_t id,
         parent_exe_info->GetValue2VarName().at(while_op.result(i))));
   }
 
-  body_block_ = while_op.body_block();
+  body_block_ = &while_op.body_block();
 
   std::unordered_map<pir::Value, std::vector<int>> inputs;
   GetInputIds(op, *parent_exe_info, &inputs);
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -237,10 +237,10 @@ void WhileOp::Build(pir::Builder &builder,             // NOLINT
   }
   argument.AddRegion(nullptr);
 }
-pir::Block *WhileOp::body_block() {
+pir::Block &WhileOp::body_block() {
   pir::Region &body_region = (*this)->region(0);
   if (body_region.empty()) body_region.emplace_back();
-  return &body_region.front();
+  return body_region.front();
 }
 pir::Value WhileOp::cond() { return (*this)->operand_source(0); }
 
@@ -259,11 +259,11 @@ void WhileOp::Print(pir::IrPrinter &printer) {
       [&]() { os << ", "; });
   os << "] { \n ^";
   pir::PrintInterleave(
-      body_block()->args_begin(),
-      body_block()->args_end(),
+      body_block().args_begin(),
+      body_block().args_end(),
       [&](pir::Value v) { printer.PrintValue(v); },
       [&]() { os << ", "; });
-  for (auto &item : *body_block()) {
+  for (auto &item : body_block()) {
     os << "\n  ";
     printer.PrintOperation(&item);
   }
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -76,7 +76,7 @@ class WhileOp : public pir::Op<WhileOp> {
                     pir::OperationArgument &argument,  // NOLINT
                     pir::Value cond,
                     const std::vector<pir::Value> &inputs);
-  pir::Block *body_block();
+  pir::Block &body_block();
   pir::Value cond();
   void Print(pir::IrPrinter &printer);  // NOLINT
   void VerifySig() {}
diff --git a/paddle/fluid/pir/transforms/build_cinn_pass.cc b/paddle/fluid/pir/transforms/build_cinn_pass.cc
@@ -191,15 +191,15 @@ class BuildCinnPass : public pir::Pass {
   void Run(pir::Operation* op) override {
     auto module_op = op->dyn_cast<pir::ModuleOp>();
     IR_ENFORCE(module_op, "build_cinn_pass should run on module op.");
-    auto* block = module_op.block();
+    auto& block = module_op.block();
 
     std::vector<GroupOpsVec> groups =
-        ::pir::SubgraphDetector(block, IsSupportCinn)();
+        ::pir::SubgraphDetector(&block, IsSupportCinn)();
     LOG(INFO) << "--- [build_cinn_pass] detected " << groups.size()
               << " cinn supported subgraphs";
     for (auto& group_ops : groups) {
       VLOG(4) << "current group_ops.size(): " << group_ops.size();
-      ::pir::ReplaceWithGroupOp(block, group_ops);
+      ::pir::ReplaceWithGroupOp(&block, group_ops);
     }
   }
 
diff --git a/paddle/fluid/pir/transforms/inplace_pass.cc b/paddle/fluid/pir/transforms/inplace_pass.cc
@@ -447,9 +447,9 @@ class InplacePass : public pir::Pass {
   void Run(pir::Operation* op) override {
     auto module_op = op->dyn_cast<pir::ModuleOp>();
     IR_ENFORCE(module_op, "inplace_pass should run on module op.");
-    auto* block = module_op.block();
+    auto& block = module_op.block();
 
-    auto inplace_ops = details::GetInplaceOps(block);
+    auto inplace_ops = details::GetInplaceOps(&block);
 
     for (auto kv : inplace_ops) {
       VLOG(6) << "Do inplace for: "
@@ -458,8 +458,8 @@ class InplacePass : public pir::Pass {
                      .dyn_cast<pir::StrAttribute>()
                      .AsString();
       pir::Block::Iterator insert_pos =
-          std::find(block->begin(), block->end(), *kv.first);
-      IR_ENFORCE(insert_pos != block->end(),
+          std::find(block.begin(), block.end(), *kv.first);
+      IR_ENFORCE(insert_pos != block.end(),
                  "Operator %s not found in block.",
                  kv.first->name());
 
diff --git a/paddle/fluid/pir/transforms/params_sync_among_devices_pass.cc b/paddle/fluid/pir/transforms/params_sync_among_devices_pass.cc
@@ -45,8 +45,8 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {
         module_op,
         phi::errors::PreconditionNotMet(
             "params_sync_among_devices_pass should run on module op."));
-    auto* block = module_op.block();
-    for (auto& inner_op : *block) {
+    auto& block = module_op.block();
+    for (auto& inner_op : block) {
       if (inner_op.isa<pir::ParameterOp>()) {
         std::string param_name = inner_op.attributes()
                                      .at("parameter_name")
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -926,16 +926,16 @@ void HandleForWhileOp(
   pir::Builder builder(ctx, block);
   auto base_while_op = op_item->dyn_cast<WhileOp>();
   auto new_while_op = builder.Build<WhileOp>(cond_val, vec_in);
-  pir::Block* body_block = new_while_op.body_block();
+  pir::Block& body_block = new_while_op.body_block();
   for (size_t i = 0; i < vec_in.size(); ++i) {
-    auto block_arg = body_block->AddArgument(vec_in[i].type());
-    (*map_value_pair)[base_while_op.body_block()->argument(i)] = block_arg;
+    auto block_arg = body_block.AddArgument(vec_in[i].type());
+    (*map_value_pair)[base_while_op.body_block().argument(i)] = block_arg;
   }
 
   // process body block
   ProcessBlock(place,
-               base_while_op.body_block(),
-               body_block,
+               &base_while_op.body_block(),
+               &body_block,
                ctx,
                map_op_pair,
                map_value_pair);
diff --git a/paddle/fluid/pir/transforms/sub_graph_extract_pass.cc b/paddle/fluid/pir/transforms/sub_graph_extract_pass.cc
@@ -50,15 +50,15 @@ class SubGraphExtractPass : public pir::Pass {
   void Run(pir::Operation* op) override {
     auto module_op = op->dyn_cast<pir::ModuleOp>();
     IR_ENFORCE(module_op, "build_cinn_pass should run on module op.");
-    auto* block = module_op.block();
+    auto& block = module_op.block();
 
     std::vector<GroupOpsVec> groups =
-        ::pir::SubgraphDetector(block, IsSplitOp)();
+        ::pir::SubgraphDetector(&block, IsSplitOp)();
     LOG(INFO) << "--- [build_cinn_pass] detected " << groups.size()
               << " cinn supported subgraphs";
     for (auto& group_ops : groups) {
       VLOG(4) << "current group_ops.size(): " << group_ops.size();
-      ::pir::ReplaceWithGroupOp(block, group_ops);
+      ::pir::ReplaceWithGroupOp(&block, group_ops);
     }
   }
 
diff --git a/paddle/pir/core/builtin_op.cc b/paddle/pir/core/builtin_op.cc
@@ -84,11 +84,13 @@ Program *ModuleOp::program() {
       iter->second.dyn_cast<PointerAttribute>().data());
 }
 
-Block *ModuleOp::block() {
-  assert(operation() != nullptr);
-  assert(operation()->num_regions() == 1);
-  assert(operation()->region(0).size() == 1);
-  return &operation()->region(0).front();
+Block &ModuleOp::block() {
+  IR_ENFORCE(operation()->num_regions(),
+             "The region size of ModuleOp must be equal to 1.");
+  auto &region = (*this)->region(0);
+  IR_ENFORCE(region.size() == 1,
+             "The region size of ModuleOp must be equal to 1.");
+  return region.front();
 }
 
 ModuleOp ModuleOp::Create(IrContext *context, Program *pointer) {
diff --git a/paddle/pir/core/builtin_op.h b/paddle/pir/core/builtin_op.h
@@ -34,7 +34,7 @@ class IR_API ModuleOp : public pir::Op<ModuleOp> {
   static const char *attributes_name[attributes_num];
   void VerifySig() const;
   Program *program();
-  Block *block();
+  Block &block();
 
   //
   // As the top operation, ModuleOp only support create&destroye through
diff --git a/paddle/pir/core/interface_support.h b/paddle/pir/core/interface_support.h
@@ -25,7 +25,7 @@ class ConstructInterfacesOrTraits {
   /// Construct method for interfaces.
   static void interface(InterfaceSet &interface_set) {  // NOLINT
     (void)std::initializer_list<int>{
-        0, (PlacementConstrctInterface<Args>(interface_set), 0)...};
+        0, (ConstrctInterface<Args>(interface_set), 0)...};
   }
 
   /// Construct method for traits.
@@ -38,8 +38,7 @@ class ConstructInterfacesOrTraits {
  private:
   /// Placement new interface.
   template <typename T>
-  static void PlacementConstrctInterface(
-      InterfaceSet &interface_set) {  // NOLINT
+  static void ConstrctInterface(InterfaceSet &interface_set) {  // NOLINT
     InterfaceValue val = InterfaceValue::
         Get<ConcreteT, T, typename T::template Model<ConcreteT>>();
     auto suceess = interface_set.insert(std::move(val)).second;
diff --git a/paddle/pir/core/program.h b/paddle/pir/core/program.h
@@ -54,8 +54,8 @@ class IR_API Program {
 
   static std::unique_ptr<Program> Parse(std::istream& is, IrContext* ctx);
 
-  Block* block() { return module_.block(); }
-  const Block* block() const { return module_op().block(); }
+  Block* block() { return &module_.block(); }
+  const Block* block() const { return &module_op().block(); }
 
   Parameter* GetParameter(const std::string& name) const;
   void SetParameter(const std::string& name,
diff --git a/paddle/pir/dialect/shape/utils/shape_optimization_utils.cc b/paddle/pir/dialect/shape/utils/shape_optimization_utils.cc
@@ -49,13 +49,13 @@ bool CompareSymbolicDimProduct(SymbolicDimProduct& lhs,    // NOLINT
 }
 
 SymbolicDimMgr::SymbolicDimMgr(ModuleOp m) : m_(m) {
-  for (auto& op : *(m.block())) {
+  for (auto& op : m.block()) {
     if (op.isa<shape::FuncOp>()) {
       symbol_table_ = SymbolTable(&op);
       return;
     }
   }
-  Builder builder = Builder(m_.ir_context(), m_.block(), m_.block()->begin());
+  Builder builder = Builder(m_.ir_context(), &m_.block(), m_.block().begin());
   shape::FuncOp func = builder.Build<shape::FuncOp>();
   symbol_table_ = SymbolTable(func);
 }
@@ -473,7 +473,7 @@ bool SymbolicDimMgr::Save() {
   };
 
   // TODO(zhangbopd): update attributes attached in DenseTensorType
-  for (auto& op : *(m_.block())) {
+  for (auto& op : m_.block()) {
     if (!op.HasAttribute(SymbolicDimOp::GetSymbolicDimAttrName())) continue;
     auto attrs =
         op.attribute<ArrayAttribute>(SymbolicDimOp::GetSymbolicDimAttrName());
@@ -499,7 +499,7 @@ bool SymbolicDimMgr::Save() {
         used_symbol_names.push_back(sym.GetSymName());
     }
   };
-  for (auto& op : *(m_.block())) {
+  for (auto& op : m_.block()) {
     if (!op.HasAttribute(SymbolicDimOp::GetSymbolicDimAttrName())) continue;
     auto attrs =
         op.attribute<ArrayAttribute>(SymbolicDimOp::GetSymbolicDimAttrName());
@@ -559,7 +559,7 @@ bool SymbolicDimMgr::Save() {
     name_to_symbol[name] = op;
   }
 
-  for (auto& op : *(m_.block())) {
+  for (auto& op : m_.block()) {
     if (!op.HasAttribute(SymbolicDimOp::GetSymbolicDimAttrName())) continue;
     auto attrs =
         op.attribute<ArrayAttribute>(SymbolicDimOp::GetSymbolicDimAttrName());
diff --git a/paddle/pir/dialect/shape/utils/shape_utils.cc b/paddle/pir/dialect/shape/utils/shape_utils.cc
@@ -49,7 +49,7 @@ bool ShapeAnalysis::IsProductEqual(
 ShapeConstraintIRAnalysis::ShapeConstraintIRAnalysis(ModuleOp m)
     : m_(m), mgr_(m) {
   mgr_.Load();
-  for (auto& op : *(m_.block())) {
+  for (auto& op : m.block()) {
     auto tie_shape_op = op.dyn_cast<shape::TieShapeOp>();
     if (!tie_shape_op) continue;
     Value result = tie_shape_op.input();
diff --git a/test/cpp/new_executor/standalone_executor_pir_test.cc b/test/cpp/new_executor/standalone_executor_pir_test.cc
@@ -303,10 +303,10 @@ TEST(StandaloneExecutor, while_op) {
       builder.Build<WhileOp>(cond_value, std::vector<pir::Value>{i, ten});
 
   // { i = i + 1}
-  pir::Block* body_block = while_op.body_block();
-  auto body_i_argument = body_block->AddArgument(i.type());
-  auto body_ten_argument = body_block->AddArgument(ten.type());
-  builder.SetInsertionPointToStart(body_block);
+  pir::Block& body_block = while_op.body_block();
+  auto body_i_argument = body_block.AddArgument(i.type());
+  auto body_ten_argument = body_block.AddArgument(ten.type());
+  builder.SetInsertionPointToStart(&body_block);
   auto one =
       builder.Build<FullOp>(std::vector<int64_t>{1}, 1, phi::DataType::INT32)
           .out();
diff --git a/test/cpp/pir/control_flow_dialect/while_op_test.cc b/test/cpp/pir/control_flow_dialect/while_op_test.cc
@@ -50,10 +50,10 @@ TEST(while_op_test, base) {
       builder.Build<WhileOp>(cond_value, std::vector<pir::Value>{i, ten});
 
   // { i = i + 1}
-  pir::Block* body_block = while_op.body_block();
-  auto body_i_argument = body_block->AddArgument(i.type());
-  auto body_ten_argument = body_block->AddArgument(ten.type());
-  builder.SetInsertionPointToStart(body_block);
+  pir::Block& body_block = while_op.body_block();
+  auto body_i_argument = body_block.AddArgument(i.type());
+  auto body_ten_argument = body_block.AddArgument(ten.type());
+  builder.SetInsertionPointToStart(&body_block);
   auto one =
       builder.Build<FullOp>(std::vector<int64_t>{1}, 1, phi::DataType::INT32)
           .out();
@@ -104,11 +104,11 @@ TEST(while_op_test, network_with_backward) {
       builder.Build<WhileOp>(cond_value, std::vector<pir::Value>{i, x});
 
   // { return i + 1, x + y}
-  pir::Block* body_block = while_op.body_block();
-  builder.SetInsertionPointToStart(body_block);
+  auto& body_block = while_op.body_block();
+  builder.SetInsertionPointToStart(&body_block);
 
-  auto body_i_argument = body_block->AddArgument(i.type());
-  auto body_x_argument = body_block->AddArgument(x.type());
+  auto body_i_argument = body_block.AddArgument(i.type());
+  auto body_x_argument = body_block.AddArgument(x.type());
 
   auto new_i = builder.Build<AddOp>(body_i_argument, one).out();
   auto new_x = builder.Build<AddOp>(body_x_argument, y).out();
@@ -141,10 +141,10 @@ TEST(while_op_test, network_with_backward) {
   auto bwd_cond = builder.Build<pir::HasElementsOp>(stack).out();
   auto while_grad = builder.Build<WhileOp>(
       bwd_cond, std::vector<pir::Value>{x_out_grad, zero});
-  pir::Block* bwd_body_block = while_grad.body_block();
-  builder.SetInsertionPointToStart(bwd_body_block);
-  auto local_x_out_grad_arg = bwd_body_block->AddArgument(x.type());
-  auto local_y_grad_arg = bwd_body_block->AddArgument(y.type());
+  pir::Block& bwd_body_block = while_grad.body_block();
+  builder.SetInsertionPointToStart(&bwd_body_block);
+  auto local_x_out_grad_arg = bwd_body_block.AddArgument(x.type());
+  auto local_y_grad_arg = bwd_body_block.AddArgument(y.type());
 
   auto pop_op = builder.Build<pir::TuplePopOp>(outlet);
   auto bwd_body_x_argument = pop_op.outlet_element(0);
diff --git a/test/cpp/pir/core/program_translator_test.cc b/test/cpp/pir/core/program_translator_test.cc
@@ -293,10 +293,10 @@ TEST(OperatorDialectTest, WhileOpProgram) {
       EXPECT_TRUE(op.isa<paddle::dialect::WhileOp>());
       EXPECT_EQ(op.num_regions(), 1u);
       // body block
-      pir::Block *body_block =
+      pir::Block &body_block =
           op.dyn_cast<paddle::dialect::WhileOp>().body_block();
       size_t body_id = 0;
-      for (auto &op1 : *body_block) {
+      for (auto &op1 : body_block) {
         if (body_id == 0) {
           EXPECT_TRUE(op1.isa<paddle::dialect::FullOp>());
         }
@@ -307,10 +307,10 @@ TEST(OperatorDialectTest, WhileOpProgram) {
           EXPECT_TRUE(op1.isa<paddle::dialect::LessThanOp>());
         }
         if (body_id == 3) {
-          pir::Block *body_body_block =
+          pir::Block &body_body_block =
               op1.dyn_cast<paddle::dialect::WhileOp>().body_block();
           size_t body_body_id = 0;
-          for (auto &op2 : *body_body_block) {
+          for (auto &op2 : body_body_block) {
             if (body_body_id == 0) {
               EXPECT_TRUE(op2.isa<paddle::dialect::FullOp>());
             }

Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,7 @@ WhileInstruction::WhileInstruction(size_t id,`
`74`	`74`	`parent_exe_info->GetValue2VarName().at(while_op.result(i))));`
`75`	`75`	`}`
`76`	`76`
`77`		`- body_block_ = while_op.body_block();`
	`77`	`+ body_block_ = &while_op.body_block();`
`78`	`78`
`79`	`79`	`std::unordered_map<pir::Value, std::vector<int>> inputs;`
`80`	`80`	`GetInputIds(op, *parent_exe_info, &inputs);`