[PIR] add some backward api for controw flow op (#59231)

winter-wang · web-flow · commit 98706c1f74f6 · 2023-11-23T00:43:36.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc b/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
@@ -40,10 +40,10 @@ void GroupOp::Build(pir::Builder &builder,             // NOLINT
                     std::unique_ptr<pir::Block> &&block) {
   VLOG(4) << "Start build GroupOp";
   if (block && !block->empty()) {
-    IR_ENFORCE(block->back()->isa<pir::YieldOp>());
-    auto *op = block->back();
-    for (size_t i = 0; i < op->num_operands(); ++i) {
-      argument.AddOutput(op->operand(i).type());
+    IR_ENFORCE(block->back().isa<pir::YieldOp>());
+    auto &op = block->back();
+    for (size_t i = 0; i < op.num_operands(); ++i) {
+      argument.AddOutput(op.operand(i).type());
     }
   }
   argument.AddRegion()->push_back(block.release());
@@ -52,7 +52,7 @@ void GroupOp::Build(pir::Builder &builder,             // NOLINT
 pir::Block *GroupOp::block() {
   pir::Region &region = (*this)->region(0);
   if (region.empty()) region.emplace_back();
-  return region.front();
+  return &region.front();
 }
 
 std::vector<pir::Operation *> GroupOp::ops() {
diff --git a/paddle/fluid/framework/new_executor/instruction/instruction_util.cc b/paddle/fluid/framework/new_executor/instruction/instruction_util.cc
@@ -196,10 +196,10 @@ OpFuncType AnalyseOpFuncType(pir::Operation* op, const platform::Place& place) {
 std::vector<pir::Value> GetYiedOpInputs(pir::Block* block) {
   std::vector<pir::Value> vec_res;
 
-  if (block && !block->empty() && block->back()->isa<pir::YieldOp>()) {
-    auto* op = block->back();
-    for (size_t i = 0; i < op->num_operands(); ++i) {
-      vec_res.emplace_back(op->operand_source(i));
+  if (block && !block->empty() && block->back().isa<pir::YieldOp>()) {
+    auto& op = block->back();
+    for (size_t i = 0; i < op.num_operands(); ++i) {
+      vec_res.emplace_back(op.operand_source(i));
     }
   }
   return vec_res;
diff --git a/paddle/fluid/ir_adaptor/translator/program_translator.cc b/paddle/fluid/ir_adaptor/translator/program_translator.cc
@@ -471,7 +471,7 @@ pir::Operation* ProgramTranslator::TranslateCondIfOperation(
                    0,
                    true_sub_block.OpSize(),
                    true_block_context,
-                   true_region.front(),
+                   &true_region.front(),
                    true,
                    cond_ops.TrueBlockOutputVarNames(),
                    cond_ops.TrueBlockInitOps());
@@ -488,7 +488,7 @@ pir::Operation* ProgramTranslator::TranslateCondIfOperation(
                    0,
                    false_sub_block.OpSize(),
                    false_block_context,
-                   false_region.front(),
+                   &false_region.front(),
                    true,
                    cond_ops.FalseBlockOutputVarNames(),
                    cond_ops.FalseBlockInitOps());
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -45,25 +45,26 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
                  std::unique_ptr<pir::Block> &&false_block) {
   VLOG(4) << "Start build IfOp";
   if (true_block && !true_block->empty() &&
-      true_block->back()->isa<pir::YieldOp>()) {
-    auto *op = true_block->back();
-    for (size_t i = 0; i < op->num_operands(); ++i) {
-      argument.AddOutput(op->operand(i).type());
+      true_block->back().isa<pir::YieldOp>()) {
+    auto &op = true_block->back();
+    for (size_t i = 0; i < op.num_operands(); ++i) {
+      argument.AddOutput(op.operand(i).type());
     }
   }
   if (false_block && !false_block->empty() &&
-      false_block->back()->isa<pir::YieldOp>()) {
-    auto *op = false_block->back();
-    PADDLE_ENFORCE_EQ(op->num_operands(),
+      false_block->back().isa<pir::YieldOp>()) {
+    auto &op = false_block->back();
+    auto size = op.num_operands();
+    PADDLE_ENFORCE_EQ(size,
                       argument.output_types.size(),
                       phi::errors::PreconditionNotMet(
                           "The output size of true block and false block must "
                           "be equal. but they are %u and %u, respectively",
                           argument.output_types.size(),
-                          op->num_operands()));
-    for (size_t i = 0; i < op->num_operands(); ++i) {
+                          size));
+    for (size_t i = 0; i < size; ++i) {
       PADDLE_ENFORCE_EQ(
-          op->operand(i).type(),
+          op.operand(i).type(),
           argument.output_types[i],
           phi::errors::PreconditionNotMet("The output[%d] type of true block "
                                           "and false block must be equal.",
@@ -84,12 +85,12 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
 pir::Block *IfOp::true_block() {
   pir::Region &region = true_region();
   if (region.empty()) region.emplace_back();
-  return region.front();
+  return &region.front();
 }
 pir::Block *IfOp::false_block() {
   pir::Region &region = false_region();
   if (region.empty()) region.emplace_back();
-  return region.front();
+  return &region.front();
 }
 
 void IfOp::Print(pir::IrPrinter &printer) {
@@ -158,22 +159,22 @@ void IfOp::VerifyRegion() {
                                         (*this)->region(0).size(),
                                         (*this)->region(1).size()));
 
-    auto *true_last_op = (*this)->region(0).front()->back();
-    auto *false_last_op = (*this)->region(1).front()->back();
-    PADDLE_ENFORCE_EQ(true_last_op->isa<pir::YieldOp>(),
-                      true,
+    auto &true_last_op = (*this)->region(0).front().back();
+    auto &false_last_op = (*this)->region(1).front().back();
+    PADDLE_ENFORCE_EQ(true,
+                      true_last_op.isa<pir::YieldOp>(),
                       phi::errors::PreconditionNotMet(
                           "The last of true block must be YieldOp"));
-    PADDLE_ENFORCE_EQ(true_last_op->num_operands(),
+    PADDLE_ENFORCE_EQ(true_last_op.num_operands(),
                       (*this)->num_results(),
                       phi::errors::PreconditionNotMet(
                           "The size of last of true block op's input must be "
                           "equal to IfOp's outputs num."));
-    PADDLE_ENFORCE_EQ(false_last_op->isa<pir::YieldOp>(),
-                      true,
+    PADDLE_ENFORCE_EQ(true,
+                      false_last_op.isa<pir::YieldOp>(),
                       phi::errors::PreconditionNotMet(
                           "The last of false block must be YieldOp"));
-    PADDLE_ENFORCE_EQ(false_last_op->num_operands(),
+    PADDLE_ENFORCE_EQ(false_last_op.num_operands(),
                       (*this)->num_results(),
                       phi::errors::PreconditionNotMet(
                           "The size of last of false block op's input must be "
@@ -195,7 +196,7 @@ void WhileOp::Build(pir::Builder &builder,             // NOLINT
 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); }
 
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
@@ -27,8 +27,10 @@
 #include "paddle/fluid/platform/enforce.h"
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
+#include "paddle/pir/core/block.h"
 #include "paddle/pir/core/op_result.h"
 #include "paddle/pir/core/operation.h"
+#include "paddle/pir/core/program.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
 
 namespace py = pybind11;
@@ -37,7 +39,10 @@ using paddle::dialect::IfOp;
 using pir::Block;
 using pir::Builder;
 using pir::Operation;
+using pir::Program;
 using pir::Region;
+using pir::StackCreateOp;
+using pir::TuplePushOp;
 using pir::Type;
 using pir::Value;
 using pir::YieldOp;
@@ -142,12 +147,49 @@ std::vector<Value> GetUsedExternalValue(const Operation& op) {
   return used_values;
 }
 
+void BuildPipeForBlock(Block* block) {
+  PADDLE_ENFORCE_NOT_NULL(
+      block,
+      paddle::platform::errors::InvalidArgument(
+          "The block used to hook local value can't be nullptr"));
+  auto& builder = *(ApiBuilder::Instance().GetBuilder());
+  Program* program = block->parent_program();
+  PADDLE_ENFORCE_NOT_NULL(
+      program,
+      paddle::platform::errors::InvalidArgument(
+          "The block used to hook local value must belong to a program"));
+
+  auto original_position = builder.insertion_point();
+
+  builder.SetInsertionPointToStart(program->block());
+  auto inlet = builder.Build<StackCreateOp>().inlet();
+  auto iter = block->end();
+  if (!block->empty() && block->back().isa<YieldOp>()) {
+    --iter;
+  }
+  std::vector<Value> local_values;
+  for (auto arg_value : block->args()) {
+    local_values.push_back(arg_value);
+  }
+  for (auto& op : *block) {
+    for (auto result_value : op.results()) {
+      local_values.push_back(result_value);
+    }
+  }
+  builder.set_insertion_point(block, iter);
+  builder.Build<TuplePushOp>(inlet, local_values);
+  builder.set_insertion_point(original_position);
+}
+
 }  // namespace
 
 namespace paddle {
 namespace pybind {
 void BindControlFlowApi(py::module* m) {
   m->def("get_used_external_value", GetUsedExternalValue);
+  m->def("build_pipe_for_block", BuildPipeForBlock);
+  m->def("cvt_as_if_op",
+         [](Operation& op) { return PyIfOp(op.dyn_cast<IfOp>()); });
   BindIfOp(m);
 }
 }  // namespace pybind
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -266,7 +266,11 @@ void BindBlock(py::module *m) {
         The constructor of Block should not be invoked directly. You can
         use `Program.block()` to get a block.
   )DOC");
-  block.def("front", &Block::front, return_value_policy::reference)
+  block
+      .def(
+          "front",
+          [](Block &self) { return &self.front(); },
+          return_value_policy::reference)
       .def_property_readonly(
           "program",
           [](Block &self) { return self.GetParentOp()->GetParentProgram(); },
@@ -287,6 +291,7 @@ void BindBlock(py::module *m) {
            [](Block &self, py::object, py::object, py::object) {
              ApiBuilder::Instance().PopInsertionPoint();
            })
+      .def("__len__", [](Block &self) { return self.size(); })
       .def(
           "remove_op",
           [](Block &self, Operation *op) {
diff --git a/paddle/pir/core/block.h b/paddle/pir/core/block.h
@@ -25,6 +25,7 @@
 
 namespace pir {
 class Operation;
+class Program;
 
 class IR_API Block {
   using OpListType = std::list<Operation *>;
@@ -42,6 +43,12 @@ class IR_API Block {
   Region *GetParent() const { return parent_; }
   Operation *GetParentOp() const;
 
+  // return the program which contains this block.
+  // if block is not in a program, return nullptr.
+  Program *parent_program() const {
+    return parent_ ? parent_->parent_program() : nullptr;
+  }
+
   bool empty() const { return ops_.empty(); }
   size_t size() const { return ops_.size(); }
 
@@ -54,8 +61,11 @@ class IR_API Block {
   ReverseIterator rbegin() { return ops_.rbegin(); }
   ReverseIterator rend() { return ops_.rend(); }
 
-  Operation *back() const { return ops_.back(); }
-  Operation *front() const { return ops_.front(); }
+  Operation &back() { return *ops_.back(); }
+  Operation &front() { return *ops_.front(); }
+  const Operation &back() const { return *ops_.back(); }
+  const Operation &front() const { return *ops_.front(); }
+
   void push_back(Operation *op);
   void push_front(Operation *op);
   Iterator insert(ConstIterator iterator, Operation *op);
diff --git a/paddle/pir/core/builtin_op.cc b/paddle/pir/core/builtin_op.cc
@@ -88,7 +88,7 @@ Block *ModuleOp::block() {
   assert(operation() != nullptr);
   assert(operation()->num_regions() == 1);
   assert(operation()->region(0).size() == 1);
-  return operation()->region(0).front();
+  return &operation()->region(0).front();
 }
 
 ModuleOp ModuleOp::Create(IrContext *context, Program *pointer) {
diff --git a/paddle/pir/core/parser/ir_parser.cc b/paddle/pir/core/parser/ir_parser.cc
@@ -190,7 +190,7 @@ std::unique_ptr<Program> IrParser::ParseProgram() {
 
 // Region := Block
 void IrParser::ParseRegion(Region& region) {  // NOLINT
-  ParseBlock(*region.front());
+  ParseBlock(region.front());
   IR_ENFORCE(PeekToken().val_ != "{",
              "Only one block in a region is supported");
 }
diff --git a/paddle/pir/core/region.cc b/paddle/pir/core/region.cc
@@ -69,9 +69,11 @@ void Region::clear() {
     blocks_.pop_back();
   }
 }
-
+Program *Region::parent_program() const {
+  return parent_ ? parent_->GetParentProgram() : nullptr;
+}
 IrContext *Region::ir_context() const {
-  IR_ENFORCE(parent_, "Region is not attached to a container.");
+  IR_ENFORCE(parent_, "Region is not attached to a operation.");
   return parent_->ir_context();
 }
 }  // namespace pir
diff --git a/paddle/pir/core/region.h b/paddle/pir/core/region.h
@@ -26,6 +26,7 @@ namespace pir {
 class Block;
 class Operation;
 class IrContext;
+class Program;
 
 class IR_API Region {
  public:
@@ -50,8 +51,12 @@ class IR_API Region {
   ConstReverseIterator rbegin() const { return blocks_.rbegin(); }
   ConstReverseIterator rend() const { return blocks_.rend(); }
 
-  Block *back() const { return blocks_.back(); }
-  Block *front() const { return blocks_.front(); }
+  Block &front() { return *blocks_.front(); }
+  Block &back() { return *blocks_.back(); }
+
+  const Block &front() const { return *blocks_.front(); }
+  const Block &back() const { return *blocks_.back(); }
+
   void push_back(Block *block);
   Block *emplace_back();
   void push_front(Block *block);
@@ -66,6 +71,9 @@ class IR_API Region {
 
   Operation *GetParent() const { return parent_; }
   void set_parent(Operation *parent) { parent_ = parent; }
+  // return the program which contains this region.
+  // if region is not in a program, return nullptr.
+  Program *parent_program() const;
 
   IrContext *ir_context() const;
 
diff --git a/paddle/pir/dialect/shape/ir/shape_op.cc b/paddle/pir/dialect/shape/ir/shape_op.cc
@@ -244,7 +244,7 @@ void FuncOp::Build(Builder &builder, OperationArgument &argument) {
 Block *FuncOp::block() {
   Region &region = (*this)->region(0);
   if (region.empty()) region.emplace_back();
-  return region.front();
+  return &region.front();
 }
 
 void FuncOp::Print(IrPrinter &printer) {
diff --git a/test/cpp/pir/cinn/build_cinn_pass_test.cc b/test/cpp/pir/cinn/build_cinn_pass_test.cc
@@ -60,9 +60,8 @@ TEST(BuildCinnPassTest, AllOpSupportCinn) {
   LOG(INFO) << "after pass: " << *origin_program;
 
   CHECK_EQ(origin_program->block()->size(), 1u);
-  pir::Operation* group_op = origin_program->block()->front();
-  pir::Block* group_block =
-      group_op->dyn_cast<cinn::dialect::GroupOp>().block();
+  pir::Operation& group_op = origin_program->block()->front();
+  pir::Block* group_block = group_op.dyn_cast<cinn::dialect::GroupOp>().block();
   CHECK_EQ(group_block->size(), 6u);
 
   std::vector<std::string> op_names = {
@@ -163,9 +162,8 @@ TEST(BuildCinnPassTest, OneCinnSubgraph) {
   LOG(INFO) << "after pass: " << *origin_program;
 
   CHECK_EQ(origin_program->block()->size(), 4u);
-  pir::Operation* group_op = origin_program->block()->front();
-  pir::Block* group_block =
-      group_op->dyn_cast<cinn::dialect::GroupOp>().block();
+  pir::Operation& group_op = origin_program->block()->front();
+  pir::Block* group_block = group_op.dyn_cast<cinn::dialect::GroupOp>().block();
   CHECK_EQ(group_block->size(), 4u);
 
   std::vector<std::string> op_names = {
@@ -219,7 +217,7 @@ TEST(BuildCinnPassTest, MultiCinnSubgraph) {
   LOG(INFO) << "after pass: " << *origin_program;
 
   CHECK_EQ(origin_program->block()->size(), 6u);
-  pir::Operation* group_op = origin_program->block()->front();
+  pir::Operation* group_op = &origin_program->block()->front();
   pir::Block* group_block =
       group_op->dyn_cast<cinn::dialect::GroupOp>().block();
   CHECK_EQ(group_block->size(), 3u);
@@ -234,7 +232,7 @@ TEST(BuildCinnPassTest, MultiCinnSubgraph) {
     CHECK_EQ(op.name(), op_names_front[index++]);
   }
 
-  group_op = origin_program->block()->back();
+  group_op = &origin_program->block()->back();
   group_block = group_op->dyn_cast<cinn::dialect::GroupOp>().block();
   CHECK_EQ(group_block->size(), 2u);
 
diff --git a/test/cpp/pir/core/ir_op_test.cc b/test/cpp/pir/core/ir_op_test.cc
diff --git a/test/cpp/pir/core/ir_program_test.cc b/test/cpp/pir/core/ir_program_test.cc
diff --git a/test/cpp/pir/core/ir_region_test.cc b/test/cpp/pir/core/ir_region_test.cc
diff --git a/test/cpp/pir/core/op_info_test.cc b/test/cpp/pir/core/op_info_test.cc
diff --git a/test/cpp/pir/shape_dialect/shape_op_test.cc b/test/cpp/pir/shape_dialect/shape_op_test.cc
diff --git a/test/ir/pir/test_if_api.py b/test/ir/pir/test_if_api.py

Original file line number	Diff line number	Diff line change
`@@ -196,10 +196,10 @@ OpFuncType AnalyseOpFuncType(pir::Operation* op, const platform::Place& place) {`
`196`	`196`	`std::vector<pir::Value> GetYiedOpInputs(pir::Block* block) {`
`197`	`197`	`std::vector<pir::Value> vec_res;`
`198`	`198`
`199`		`- if (block && !block->empty() && block->back()->isa<pir::YieldOp>()) {`
`200`		`- auto* op = block->back();`
`201`		`- for (size_t i = 0; i < op->num_operands(); ++i) {`
`202`		`- vec_res.emplace_back(op->operand_source(i));`
	`199`	`+ if (block && !block->empty() && block->back().isa<pir::YieldOp>()) {`
	`200`	`+ auto& op = block->back();`
	`201`	`+ for (size_t i = 0; i < op.num_operands(); ++i) {`
	`202`	`+ vec_res.emplace_back(op.operand_source(i));`
`203`	`203`	`}`
`204`	`204`	`}`
`205`	`205`	`return vec_res;`
Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ Block *ModuleOp::block() {`
`88`	`88`	`assert(operation() != nullptr);`
`89`	`89`	`assert(operation()->num_regions() == 1);`
`90`	`90`	`assert(operation()->region(0).size() == 1);`
`91`		`- return operation()->region(0).front();`
	`91`	`+ return &operation()->region(0).front();`
`92`	`92`	`}`
`93`	`93`
`94`	`94`	`ModuleOp ModuleOp::Create(IrContext context, Program pointer) {`
Original file line number	Diff line number	Diff line change
`@@ -190,7 +190,7 @@ std::unique_ptr<Program> IrParser::ParseProgram() {`
`190`	`190`
`191`	`191`	`// Region := Block`
`192`	`192`	`void IrParser::ParseRegion(Region& region) { // NOLINT`
`193`		`- ParseBlock(*region.front());`
	`193`	`+ ParseBlock(region.front());`
`194`	`194`	`IR_ENFORCE(PeekToken().val_ != "{",`
`195`	`195`	`"Only one block in a region is supported");`
`196`	`196`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,9 +69,11 @@ void Region::clear() {`
`69`	`69`	`blocks_.pop_back();`
`70`	`70`	`}`
`71`	`71`	`}`
`72`		`-`
	`72`	`+Program *Region::parent_program() const {`
	`73`	`+ return parent_ ? parent_->GetParentProgram() : nullptr;`
	`74`	`+}`
`73`	`75`	`IrContext *Region::ir_context() const {`
`74`		`- IR_ENFORCE(parent_, "Region is not attached to a container.");`
	`76`	`+ IR_ENFORCE(parent_, "Region is not attached to a operation.");`
`75`	`77`	`return parent_->ir_context();`
`76`	`78`	`}`
`77`	`79`	`} // namespace pir`