[SCF] Handle lowering of Execute region to LLVM

[wsmoses]

Summary

Lower SCF.executeregionop to llvm by essentially inlining the region and replacing the return

[mehdi_amini]

Handle lowering of Execute region to LLVM

I think it should say "to standard CFG".

Can you add tests?

[nicolasvasilache]

some test please ?

[wsmoses]

Add test

[mehdi_amini]

(please rename the title)

mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp

 //
 struct IfLowering : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(IfOp ifOp,
                                 PatternRewriter &rewriter) const override;
 };
 
+struct ExecuteRegionLowering : public OpRewritePattern<ExecuteRegionOp> {
+  using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExecuteRegionOp op,
+                                PatternRewriter &rewriter) const override;
+};
+
 struct ParallelLowering : public OpRewritePattern<mlir::scf::ParallelOp> {
   using OpRewritePattern<mlir::scf::ParallelOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(mlir::scf::ParallelOp parallelOp,
                                 PatternRewriter &rewriter) const override;
 };
 
 /// Create a CFG subgraph for this loop construct. The regions of the loop need
rewriter.create<CondBranchOp>(loc, ifOp.condition(), thenBlock,
                                 /*falseArgs=*/ArrayRef<Value>());
 
   // Ok, we're done!
   rewriter.replaceOp(ifOp, continueBlock->getArguments());
   return success();
 }
 
 LogicalResult
+ExecuteRegionLowering::matchAndRewrite(ExecuteRegionOp op,
+                                       PatternRewriter &rewriter) const {
+  auto loc = op.getLoc();
+
+  auto *condBlock = rewriter.getInsertionBlock();
+  auto opPosition = rewriter.getInsertionPoint();
+  auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);
+
+  auto &region = op.region();
+  rewriter.setInsertionPointToEnd(condBlock);
+  rewriter.create<BranchOp>(loc, &region.front());
+
+  for (Block &block : region) {
+    if (auto terminator = dyn_cast<scf::YieldOp>(block.getTerminator())) {
+      ValueRange terminatorOperands = terminator->getOperands();
+      rewriter.setInsertionPointToEnd(&block);
+      rewriter.create<BranchOp>(loc, remainingOpsBlock, terminatorOperands);
+      rewriter.eraseOp(terminator);
+    }
+  }
+
+  rewriter.inlineRegionBefore(region, remainingOpsBlock);
+
+  SmallVector<Value> vals;
+  for (auto arg : remainingOpsBlock->addArguments(op->getResultTypes())) {
+    vals.push_back(arg);
+  }
+  rewriter.replaceOp(op, vals);
+  return success();
+}
+
+LogicalResult
 ParallelLowering::matchAndRewrite(ParallelOp parallelOp,
                                   PatternRewriter &rewriter) const {
   Location loc = parallelOp.getLoc();
 
   // For a parallel loop, we essentially need to create an n-dimensional loop
   // nest. We do this by translating to scf.for ops and have those lowered in
   // a further rewrite. If a parallel loop contains reductions (and thus returns
   // values), forward the initial values for the reductions down the loop
DoWhileLowering::matchAndRewrite(WhileOp whileOp,
   // Replace the op with values "yielded" from the "before" region, which are
   // visible by dominance.
   rewriter.replaceOp(whileOp, condOp.args());
 
   return success();
 }
 
 void mlir::populateLoopToStdConversionPatterns(RewritePatternSet &patterns) {
-  patterns.add<ForLowering, IfLowering, ParallelLowering, WhileLowering>(
-      patterns.getContext());
+  patterns.add<ForLowering, IfLowering, ParallelLowering, WhileLowering,
+               ExecuteRegionLowering>(patterns.getContext());
   patterns.add<DoWhileLowering>(patterns.getContext(), /*benefit=*/2);
 }
 
 void SCFToStandardPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateLoopToStdConversionPatterns(patterns);
   // Configure conversion to lower out scf.for, scf.if, scf.parallel and
   // scf.while. Anything else is fine.
   ConversionTarget target(getContext());
-  target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp, scf::WhileOp>();
+  target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp, scf::WhileOp,
+                      scf::ExecuteRegionOp>();
   target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<Pass> mlir::createLowerToCFGPass() {
   return std::make_unique<SCFToStandardPass>();
 }

mlir/test/Conversion/SCFToStandard/convert-to-cfg.mlir

scf.parallel (%i) = (%arg1) to (%arg2) step (%arg3) {
     }
     scf.yield
   }
 
   // CHECK: ^[[LOOP_CONT]]:
   // CHECK:   return
   return
 }
+
+// CHECK-LABEL: func @func_execute_region_elim_multi_yield
+func @func_execute_region_elim_multi_yield() {
+    "test.foo"() : () -> ()
+    %v = scf.execute_region -> i64 {
+      %c = "test.cmp"() : () -> i1
+      cond_br %c, ^bb2, ^bb3
+    ^bb2:
+      %x = "test.val1"() : () -> i64
+      scf.yield %x : i64
+    ^bb3:
+      %y = "test.val2"() : () -> i64
+      scf.yield %y : i64
+    }
+    "test.bar"(%v) : (i64) -> ()
+  return
+}
+
+// CHECK-NOT: execute_region
+// CHECK:     "test.foo"
+// CHECK:     br ^[[rentry:.+]]
+// CHECK:   ^[[rentry]]
+// CHECK:     %[[cmp:.+]] = "test.cmp"
+// CHECK:     cond_br %[[cmp]], ^[[bb1:.+]], ^[[bb2:.+]]
+// CHECK:   ^[[bb1]]:
+// CHECK:     %[[x:.+]] = "test.val1"
+// CHECK:     br ^[[bb3:.+]](%[[x]] : i64)
+// CHECK:   ^[[bb2]]:
+// CHECK:     %[[y:.+]] = "test.val2"
+// CHECK:     br ^[[bb3]](%[[y:.+]] : i64)
+// CHECK:   ^[[bb3]](%[[z:.+]]: i64):
+// CHECK:     "test.bar"(%[[z]])
+// CHECK:     return