[mlir] Extend SimplifyTrivialLoops

[ayzhuang]

Summary

Fold away empty loops that iterate at least once and only return
values defined outside of the loop.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

763

Where is the "iterates at least once" guaranteed?

766–767

[dcaballe]

I'll let @bondhugula clarify but my impression about his comment in D120776 is that we should be able to have a single implementation that works for both ScfForOp and AffineForOp by using/extending interfaces like LoopLikeInterface? Maybe I got it wrong. In any case, adding a scf specific implementation is also great. Thanks!

mlir/test/Dialect/SCF/canonicalize.mlir

363

Please, add // ----- before the new test.

[ayzhuang]

mlir/lib/Dialect/SCF/SCF.cpp

763

If the loop is known to have 0 iteration, we would have removed the loop and returned success: line 711-714.

mlir/test/Dialect/SCF/canonicalize.mlir

363

for_yields_2, for_yields_3, and for_yields_4 all use make_i32() declared before for_yields_2, thus no // ----- between them.

[ayzhuang]

Address review comment to modify test.

[bondhugula]

LGTM - modulo the remaining minor comments.

mlir/lib/Dialect/SCF/SCF.cpp

764

.. remove it and replace it with yield values.

767–769

In addition to Mehdi's suggestions, you could use llvm::any_of here.

if (llvm::any_of(...))
  return failure();

[ayzhuang]

Address review comments.

[ayzhuang]

Upload the correct patch.

[ayzhuang]

Thank you for the review. I'll merge it tomorrow if there are no more comments

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

761

This check isn't tight enough:
lbValue + step - 1 <= ubValue?

[ayzhuang]

mlir/lib/Dialect/SCF/SCF.cpp

761

When lbValue >= ubValue, the loop does not iterate, so we do not fold. When lbValue + step - 1 <= ubValue, the loop can still iterate. For example, if lbValue = 0, step = 1, and ubValue = 1, we have 0 + 1 - 1 <= 1, but the loop iterates once.

[ayzhuang]

Add a comment and remove unneeded check.

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

761

Oops... looks like I was thinking about something else. Sounds fine.

mlir/test/Dialect/SCF/canonicalize.mlir

366

Please either add a short comment hereon what's being tested or make the function name descriptive enough.

[ayzhuang]

Address review comment to add a comment to the unit test.

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(scf::ForOp forOp,
     auto mergedYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
     cloneFilteredTerminator(mergedYieldOp);
     rewriter.eraseOp(mergedYieldOp);
     rewriter.replaceOp(forOp, newResultValues);
     return success();
   }
 };
 
-/// Rewriting pattern that erases loops that are known not to iterate and
-/// replaces single-iteration loops with their bodies.
+/// Rewriting pattern that erases loops that are known not to iterate, replaces
+/// single-iteration loops with their bodies, and removes empty loops that
+/// iterate at least once and only return values defined outside of the loop.
 struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForOp op,
                                 PatternRewriter &rewriter) const override {
     // If the upper bound is the same as the lower bound, the loop does not
     // iterate, just remove it.
     if (op.getLowerBound() == op.getUpperBound()) {
if ((lbValue + stepValue).sge(ubValue)) {
       SmallVector<Value, 4> blockArgs;
       blockArgs.reserve(op.getNumIterOperands() + 1);
       blockArgs.push_back(op.getLowerBound());
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }
 
+    // Now we are left with loops that have more than 1 iterations.
+    Block &block = op.getRegion().front();
+    if (!llvm::hasSingleElement(block))
       return failure();
+    // If the loop is empty, iterates at least once, and only returns values
+    // defined outside of the loop, remove it and replace it with yield values.
+    auto yieldOp = cast<scf::YieldOp>(block.getTerminator());
+    auto yieldOperands = yieldOp.getOperands();
+    if (llvm::any_of(yieldOperands,
+                     [&](Value v) { return !op.isDefinedOutsideOfLoop(v); }))
+      return failure();
+    rewriter.replaceOp(op, yieldOperands);
+    return success();
   }
 };
 
 /// Perform a replacement of one iter OpOperand of an scf.for to the
 /// `replacement` value which is expected to be the source of a tensor.cast.
 /// tensor.cast ops are inserted inside the block to account for the type cast.
 static ForOp replaceTensorCastForOpIterArg(PatternRewriter &rewriter,
                                            OpOperand &operand,

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-LABEL:   func @for_yields_3
 //  CHECK-NEXT:     %[[a:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     %[[b:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     %[[r1:.*]] = scf.for {{.*}} iter_args(%arg4 = %[[a]]) -> (i32) {
 //  CHECK-NEXT:       %[[c:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:       scf.yield %[[c]] : i32
 //  CHECK-NEXT:     }
 //  CHECK-NEXT:     return %[[a]], %[[r1]], %[[b]] : i32, i32, i32
 
 // -----
 
+// Test that an empty loop which iterates at least once and only returns
+// values defined outside of the loop is folded away.
+func @for_yields_4() -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %a = arith.constant 3 : i32
+  %b = arith.constant 4 : i32
+  %r = scf.for %i = %c0 to %c2 step %c1 iter_args(%0 = %a) -> i32 {
+    scf.yield %b : i32
+  }
+  return %r : i32
+}
+
+// CHECK-LABEL:   func @for_yields_4
+//  CHECK-NEXT:     %[[b:.*]] = arith.constant 4 : i32
+//  CHECK-NEXT:     return %[[b]] : i32
+
+// -----
+
 // CHECK-LABEL: @replace_true_if
 func @replace_true_if() {
   %true = arith.constant true
   // CHECK-NOT: scf.if
   // CHECK: "test.op"
   scf.if %true {
     "test.op"() : () -> ()
     scf.yield