[PIR] support if_op build after subblock has been completed. (#57958)

winter-wang · web-flow · commit 9bea183aa515 · 2023-10-10T10:21:13.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/CMakeLists.txt b/paddle/fluid/pir/dialect/operator/ir/CMakeLists.txt
@@ -191,7 +191,7 @@ cc_library(
 cc_library(
   pd_op_dialect_op
   SRCS ${op_source_file} manual_op.cc control_flow_op.cc
-  DEPS pd_op_dialect_core)
+  DEPS pd_op_dialect_core pir_control_flow)
 cc_library(
   api_builder
   SRCS api_builder.cc
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -17,9 +17,11 @@ paddle::dialect::IfOp, paddle::dialect::WhileOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 
+#include "paddle/phi/core/enforce.h"
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/ir_printer.h"
 #include "paddle/pir/core/operation_utils.h"
+#include "paddle/pir/dialect/control_flow/ir/cf_ops.h"
 
 namespace paddle {
 namespace dialect {
@@ -33,6 +35,50 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
   argument.AddInput(cond);
   argument.output_types.swap(output_types);
 }
+
+void IfOp::Build(pir::Builder &builder,             // NOLINT
+                 pir::OperationArgument &argument,  // NOLINT
+                 pir::Value cond,
+                 std::unique_ptr<pir::Block> &&true_block,
+                 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());
+    }
+  }
+  if (false_block && !false_block->empty() &&
+      false_block->back()->isa<pir::YieldOp>()) {
+    auto *op = false_block->back();
+    PADDLE_ENFORCE_EQ(op->num_operands(),
+                      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) {
+      PADDLE_ENFORCE_EQ(
+          op->operand(i).type(),
+          argument.output_types[i],
+          phi::errors::PreconditionNotMet("The output[%d] type of true block "
+                                          "and false block must be equal.",
+                                          i));
+    }
+  } else {
+    PADDLE_ENFORCE(argument.output_types.empty(),
+                   phi::errors::PreconditionNotMet(
+                       "The output size of true block and false block must be "
+                       "equal. but they are %u and 0, respectively",
+                       argument.output_types.size()));
+  }
+  argument.AddRegion()->push_back(true_block.release());
+  argument.AddRegion()->push_back(false_block.release());
+  argument.AddInput(cond);
+}
+
 pir::Block *IfOp::true_block() {
   pir::Region &true_region = (*this)->region(0);
   if (true_region.empty()) true_region.emplace_back();
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -31,11 +31,11 @@ class IfOp : public pir::Op<IfOp> {
                     pir::Value cond,
                     std::vector<pir::Type> &&output_types);
 
-  // static void Build(pir::Builder &builder,             // NOLINT
-  //                   pir::OperationArgument &argument,  // NOLINT
-  //                   pir::Value cond,
-  //                   std::unique_ptr<pir::Block>&& true_block,
-  //                   std::unique_ptr<pir::Block>&& false_block);
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value cond,
+                    std::unique_ptr<pir::Block> &&true_block,
+                    std::unique_ptr<pir::Block> &&false_block);
 
   pir::Value cond() { return operand_source(0); }
   pir::Block *true_block();
diff --git a/test/cpp/pir/control_flow_dialect/if_op_test.cc b/test/cpp/pir/control_flow_dialect/if_op_test.cc
@@ -59,3 +59,43 @@ TEST(if_op_test, base) {
 
   LOG(INFO) << ss.str();
 }
+
+TEST(if_op_test, build_by_block) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::ControlFlowDialect>();
+
+  pir::Program program(ctx);
+  pir::Block* block = program.block();
+  pir::Builder builder(ctx, block);
+  auto full_op = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, true, phi::DataType::BOOL);
+
+  // construct true block
+  std::unique_ptr<pir::Block> true_block(new pir::Block());
+  builder.SetInsertionPointToStart(true_block.get());
+  auto full_op_1 = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2}, true, phi::DataType::BOOL);
+  builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_1.out()});
+
+  // construct false block
+  std::unique_ptr<pir::Block> false_block(new pir::Block());
+  builder.SetInsertionPointToStart(false_block.get());
+  auto full_op_2 = builder.Build<paddle::dialect::FullOp>(
+      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.Build<paddle::dialect::IfOp>(
+      full_op.out(), std::move(true_block), std::move(false_block));
+
+  EXPECT_FALSE(true_block);
+  EXPECT_FALSE(false_block);
+  EXPECT_EQ(full_op_2->GetParentProgram(), &program);
+
+  std::stringstream ss;
+  program.Print(ss);
+
+  LOG(INFO) << ss.str();
+}
