[MLIR] Insert loop.yield to IfOp regions only if it's void.

Differential Revision: https://reviews.llvm.org/D76603

2 files changed, 16 insertions, 8 deletions

diff --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
index a20c44f..7b01072 100644
--- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
+++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
@@ -223,20 +223,21 @@ def IfOp : Loop_Op<"if",
     OpBuilder<"Builder *builder, OperationState &result, "
               "Value cond, bool withElseRegion">,
     OpBuilder<"Builder *builder, OperationState &result, "
-              "TypeRange resultTypes, Value cond, "
-              "bool withElseRegion">
+              "TypeRange resultTypes, Value cond, bool withElseRegion">
   ];
 
   let extraClassDeclaration = [{
     OpBuilder getThenBodyBuilder() {
       assert(!thenRegion().empty() && "Unexpected empty 'then' region.");
       Block &body = thenRegion().front();
-      return OpBuilder(&body, std::prev(body.end()));
+      return OpBuilder(&body,
+                       results().empty() ? std::prev(body.end()) : body.end());
     }
     OpBuilder getElseBodyBuilder() {
       assert(!elseRegion().empty() && "Unexpected empty 'else' region.");
       Block &body = elseRegion().front();
-      return OpBuilder(&body, std::prev(body.end()));
+      return OpBuilder(&body,
+                       results().empty() ? std::prev(body.end()) : body.end());
     }
   }];
 }
diff --git a/mlir/lib/Dialect/LoopOps/LoopOps.cpp b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
index 5106cc6..3d7ee38 100644
--- a/mlir/lib/Dialect/LoopOps/LoopOps.cpp
+++ b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
@@ -201,18 +201,25 @@ ForOp mlir::loop::getForInductionVarOwner(Value val) {
 
 void IfOp::build(Builder *builder, OperationState &result, Value cond,
                  bool withElseRegion) {
-  build(builder, result, /*resultTypes=*/llvm::None, cond, withElseRegion);
+    build(builder, result, /*resultTypes=*/llvm::None, cond, withElseRegion);
 }
 
 void IfOp::build(Builder *builder, OperationState &result,
                  TypeRange resultTypes, Value cond, bool withElseRegion) {
   result.addOperands(cond);
   result.addTypes(resultTypes);
+
   Region *thenRegion = result.addRegion();
+  thenRegion->push_back(new Block());
+  if (resultTypes.empty())
+    IfOp::ensureTerminator(*thenRegion, *builder, result.location);
+
   Region *elseRegion = result.addRegion();
-  IfOp::ensureTerminator(*thenRegion, *builder, result.location);
-  if (withElseRegion)
-    IfOp::ensureTerminator(*elseRegion, *builder, result.location);
+  if (withElseRegion) {
+    elseRegion->push_back(new Block());
+    if (resultTypes.empty())
+      IfOp::ensureTerminator(*elseRegion, *builder, result.location);
+  }
 }
 
 static LogicalResult verify(IfOp op) {