[Mlir-commits] [mlir] 51e3709 - [MLIR] Don't insert YieldOp for non-void loop.for by default.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2020-04-05T11:48:22+02:00
New Revision: 51e3709c2b9055fe9e6fc83fc7f056bd55b2aaa6

URL: https://github.com/llvm/llvm-project/commit/51e3709c2b9055fe9e6fc83fc7f056bd55b2aaa6
DIFF: https://github.com/llvm/llvm-project/commit/51e3709c2b9055fe9e6fc83fc7f056bd55b2aaa6.diff

LOG: [MLIR] Don't insert YieldOp for non-void loop.for by default.

The ForOp::build ensures that there is a block terminator which is great for
the default use case when there are no iter_args and loop.for returns no
results.  In non-zero results case we always need to call replaceOpWithNewOp
which is not the nicest thing in the world. We can stop inserting YieldOp when
iter_args is non-empty. IfOp::build already behaves similarly.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LoopOps/LoopOps.td
    mlir/lib/Conversion/LoopToStandard/LoopToStandard.cpp
    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 c51bf6f3c7b9..1e446efa9758 100644
--- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
+++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
@@ -52,7 +52,7 @@ def ForOp : Loop_Op<"for",
     the lower bound but does not include the upper bound.
 
     The body region must contain exactly one block that terminates with
-    "loop.yield".  Calling ForOp::build will create such a region and insert
+    "loop.yield". Calling ForOp::build will create such a region and insert
     the terminator implicitly if none is defined, so will the parsing even in
     cases when it is absent from the custom format. For example:
 
@@ -71,9 +71,11 @@ def ForOp : Loop_Op<"for",
 
     The region must terminate with a "loop.yield" that passes all the current
     iteration variables to the next iteration, or to the "loop.for" result, if
-    at the last iteration.  "loop.for" results hold the final values after the
-    last iteration.
+    at the last iteration. Note, that when the loop-carried variables are
+    present, calling ForOp::build will not insert the terminator implicitly.
+    The caller must insert "loop.yield" in that case.
 
+    "loop.for" results hold the final values after the last iteration.
     For example, to sum-reduce a memref:
 
     ```mlir

diff  --git a/mlir/lib/Conversion/LoopToStandard/LoopToStandard.cpp b/mlir/lib/Conversion/LoopToStandard/LoopToStandard.cpp
index e971ca4ba4b5..7d20561ab6cf 100644
--- a/mlir/lib/Conversion/LoopToStandard/LoopToStandard.cpp
+++ b/mlir/lib/Conversion/LoopToStandard/LoopToStandard.cpp
@@ -357,15 +357,11 @@ ParallelLowering::matchAndRewrite(ParallelOp parallelOp,
       // the results of the parallel loop when it is fully rewritten.
       loopResults.assign(forOp.result_begin(), forOp.result_end());
       first = false;
-    } else {
-      // A loop is constructed with an empty "yield" terminator by default.
-      // Replace it with another "yield" that forwards the results of the nested
-      // loop to the parent loop. We need to explicitly make sure the new
-      // terminator is the last operation in the block because further
-      // transforms rely on this.
+    } else if (!forOp.getResults().empty()) {
+      // A loop is constructed with an empty "yield" terminator if there are
+      // no results.
       rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
-      rewriter.replaceOpWithNewOp<YieldOp>(
-          rewriter.getInsertionBlock()->getTerminator(), forOp.getResults());
+      rewriter.create<YieldOp>(loc, forOp.getResults());
     }
 
     rewriter.setInsertionPointToStart(forOp.getBody());
@@ -398,9 +394,10 @@ ParallelLowering::matchAndRewrite(ParallelOp parallelOp,
         mapping.lookup(reduceBlock.getTerminator()->getOperand(0)));
   }
 
-  rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
-  rewriter.replaceOpWithNewOp<YieldOp>(
-      rewriter.getInsertionBlock()->getTerminator(), yieldOperands);
+  if (!yieldOperands.empty()) {
+    rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
+    rewriter.create<YieldOp>(loc, yieldOperands);
+  }
 
   rewriter.replaceOp(parallelOp, loopResults);
 

diff  --git a/mlir/lib/Dialect/LoopOps/LoopOps.cpp b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
index c62507fa8aaf..d09620599172 100644
--- a/mlir/lib/Dialect/LoopOps/LoopOps.cpp
+++ b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
@@ -47,7 +47,9 @@ void ForOp::build(Builder *builder, OperationState &result, Value lb, Value ub,
   for (Value v : iterArgs)
     result.addTypes(v.getType());
   Region *bodyRegion = result.addRegion();
-  ForOp::ensureTerminator(*bodyRegion, *builder, result.location);
+  bodyRegion->push_back(new Block());
+  if (iterArgs.empty())
+    ForOp::ensureTerminator(*bodyRegion, *builder, result.location);
   bodyRegion->front().addArgument(builder->getIndexType());
   for (Value v : iterArgs)
     bodyRegion->front().addArgument(v.getType());
@@ -201,7 +203,7 @@ 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,