[Mlir-commits] [mlir] 795e726 - [MLIR] Fix affine.for empty loop body folder

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-07-22T08:11:40+05:30
New Revision: 795e726f5f15a6f07d146255721fe3d392c6a618

URL: https://github.com/llvm/llvm-project/commit/795e726f5f15a6f07d146255721fe3d392c6a618
DIFF: https://github.com/llvm/llvm-project/commit/795e726f5f15a6f07d146255721fe3d392c6a618.diff

LOG: [MLIR] Fix affine.for empty loop body folder

Fix affine.for empty loop body folder in the presence of yield values.
The existing pattern ignored iter_args/yield values and thus crashed
when yield values had uses.

Reviewed By: mehdi_amini

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index b1721cd2f4640..370084d116c6e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1628,7 +1628,8 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
 }
 
 namespace {
-/// This is a pattern to fold trivially empty loops.
+/// This is a pattern to fold trivially empty loop bodies.
+/// TODO: This should be moved into the folding hook.
 struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
   using OpRewritePattern<AffineForOp>::OpRewritePattern;
 
@@ -1637,7 +1638,8 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
     // Check that the body only contains a yield.
     if (!llvm::hasSingleElement(*forOp.getBody()))
       return failure();
-    rewriter.eraseOp(forOp);
+    // The initial values of the iteration arguments would be the op's results.
+    rewriter.replaceOp(forOp, forOp.getIterOperands());
     return success();
   }
 };

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index f12a0ca032d51..3c55431423ed8 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -460,14 +460,18 @@ func @constant_fold_bounds(%N : index) {
 
 // -----
 
-// CHECK-LABEL:  func @fold_empty_loop() {
-func @fold_empty_loop() {
-  // CHECK-NOT: affine.for
+// CHECK-LABEL:  func @fold_empty_loops()
+func @fold_empty_loops() -> index {
+  %c0 = constant 0 : index
   affine.for %i = 0 to 10 {
   }
-  return
+  %res = affine.for %i = 0 to 10 iter_args(%arg = %c0) -> index {
+    affine.yield %arg : index
+  }
+  // CHECK-NEXT: %[[zero:.*]] = constant 0
+  // CHECK-NEXT: return %[[zero]]
+  return %res : index
 }
-// CHECK: return
 
 // -----