diff options
| -rw-r--r-- | mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 6 | ||||
| -rw-r--r-- | mlir/test/Dialect/Affine/canonicalize.mlir | 14 |
2 files changed, 13 insertions, 7 deletions
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index b1721cd..370084d 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 f12a0ca..3c55431 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 // ----- |