[MLIR] Fix affine.for empty loop body folder

[bondhugula]

Summary

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.

[bondhugula]

Gentle ping - reviewers: @rriddle (from the history, you were the last to touch this code).

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
   if (lbMap != prevLbMap)
     forOp.setLowerBound(lbOperands, lbMap);
   if (ubMap != prevUbMap)
     forOp.setUpperBound(ubOperands, ubMap);
   return success();
 }
 
 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;
 
   LogicalResult matchAndRewrite(AffineForOp forOp,
                                 PatternRewriter &rewriter) const override {
     // 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();
   }
 };
 } // end anonymous namespace
 
 void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<AffineForEmptyLoopFolder>(context);

mlir/test/Dialect/Affine/canonicalize.mlir

func @constant_fold_bounds(%N : index) {
   affine.for %k = max affine_map<()[s0] -> (0, s0)> ()[%l] to min affine_map<()[s0] -> (100, s0)> ()[%N] {
     "foo"(%k, %c3) : (index, index) -> ()
   }
   return
 }
 
 // -----
 
-// 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
 
 // -----
 
 // CHECK-LABEL:  func @fold_zero_iter_loops
 // CHECK-SAME: %[[ARG:.*]]: index
 func @fold_zero_iter_loops(%in : index) -> index {
   %c1 = constant 1 : index
   affine.for %i = 0 to 0 {