[Mlir-commits] [mlir] 564ece9 - [MLIR] Insert loop.yield to IfOp regions only if it's void.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2020-03-23T23:21:10+01:00
New Revision: 564ece93b8362e2334edbd59e1d7e4a63253c7b8

URL: https://github.com/llvm/llvm-project/commit/564ece93b8362e2334edbd59e1d7e4a63253c7b8
DIFF: https://github.com/llvm/llvm-project/commit/564ece93b8362e2334edbd59e1d7e4a63253c7b8.diff

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

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LoopOps/LoopOps.td
    mlir/lib/Dialect/LoopOps/LoopOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
index a20c44fa5e16..7b010729831b 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 5106cc63cb46..3d7ee3846c22 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) {