[mlir] scf::ForOp: Drop iter arguments (and corresponding result) with no use

[chelini]

Summary

'ForOpIterArgsFolder' can now remove iterator arguments (and corresponding
results) with no use.

Example:

%cst = constant 32 : i32

%0:2 = scf.for %arg1 = %lb to %ub step %step iter_args(%arg2 = %arg0, %arg3 = %cst)
  -> (i32, i32) {
  %1 = addu %arg2, %cst : i32
  scf.yield %1, %1 : i32, i32
}

use(%0#0)

%arg3 is not used in the block, and its corresponding result %0#1 has no use,
thus remove the iter argument.

[chelini]

Minor update for test case.

[chelini]

rebase.

[chelini]

@nicolasvasilache thanks for landing it!

mlir/lib/Dialect/SCF/SCF.cpp

static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
   rewriter.eraseOp(terminator);
 }
 
 namespace {
 // Fold away ForOp iter arguments when:
 // 1) The op yields the iter arguments.
 // 2) The iter arguments have no use and the corresponding outer region
 // iterators (inputs) are yielded.
+// 3) The iter arguments have no use and the corresponding (operation) results
+// have no use.
 //
 // These arguments must be defined outside of
 // the ForOp region and can just be forwarded after simplifying the op inits,
 // yields and returns.
 //
 // The implementation uses `mergeBlockBefore` to steal the content of the
 // original ForOp and avoid cloning.
 struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
SmallVector<Value, 4> newBlockTransferArgs, newIterArgs, newYieldValues,
         newResultValues;
     newBlockTransferArgs.reserve(1 + forOp.getNumIterOperands());
     newBlockTransferArgs.push_back(Value()); // iv placeholder with null value
     newIterArgs.reserve(forOp.getNumIterOperands());
     newYieldValues.reserve(yieldOp.getNumOperands());
     newResultValues.reserve(forOp.getNumResults());
     for (auto it : llvm::zip(forOp.getIterOperands(),   // iter from outside
                              forOp.getRegionIterArgs(), // iter inside region
+                             forOp.getResults(),        // op results
                              yieldOp.getOperands()      // iter yield
                              )) {
       // Forwarded is `true` when:
       // 1) The region `iter` argument is yielded.
-      // 2) The region `iter` argument has zero use, and the corresponding iter
+      // 2) The region `iter` argument has no use, and the corresponding iter
       // operand (input) is yielded.
-      bool forwarded =
-          ((std::get<1>(it) == std::get<2>(it)) ||
-           (std::get<1>(it).use_empty() && std::get<0>(it) == std::get<2>(it)));
+      // 3) The region `iter` argument has no use, and the corresponding op
+      // result has no use.
+      bool forwarded = ((std::get<1>(it) == std::get<3>(it)) ||
+                        (std::get<1>(it).use_empty() &&
+                         (std::get<0>(it) == std::get<3>(it) ||
+                          std::get<2>(it).use_empty())));
       keepMask.push_back(!forwarded);
       canonicalize |= forwarded;
       if (forwarded) {
         newBlockTransferArgs.push_back(std::get<0>(it));
         newResultValues.push_back(std::get<0>(it));
         continue;
       }
       newIterArgs.push_back(std::get<0>(it));
-      newYieldValues.push_back(std::get<2>(it));
+      newYieldValues.push_back(std::get<3>(it));
       newBlockTransferArgs.push_back(Value()); // placeholder with null value
       newResultValues.push_back(Value());      // placeholder with null value
     }
 
     if (!canonicalize)
       return failure();
 
     scf::ForOp newForOp = rewriter.create<scf::ForOp>(

mlir/test/Dialect/SCF/canonicalize.mlir

%0:2 = scf.for %arg1 = %lb to %ub step %step iter_args(%arg2 = %arg0, %arg3 = %cst)
     -> (i32, i32) {
     %1 = addi %arg2, %cst : i32
     scf.yield %1, %cst : i32, i32
   }
 
   // CHECK: return %[[FOR_RES]], %[[C32]] : i32, i32
   return %0#0, %0#1 : i32, i32
 }
+
+// -----
+
+// CHECK-LABEL: fold_away_iter_and_result_with_no_use
+//  CHECK-SAME:   %[[A0:[0-9a-z]*]]: i32
+func @fold_away_iter_and_result_with_no_use(%arg0 : i32,
+                    %ub : index, %lb : index, %step : index) -> (i32) {
+  // CHECK-NEXT: %[[C32:.*]] = constant 32 : i32
+  %cst = constant 32 : i32
+  // CHECK-NEXT: %[[FOR_RES:.*]] = scf.for {{.*}} iter_args({{.*}} = %[[A0]]) -> (i32) {
+  %0:2 = scf.for %arg1 = %lb to %ub step %step iter_args(%arg2 = %arg0, %arg3 = %cst)
+    -> (i32, i32) {
+    %1 = addi %arg2, %cst : i32
+    scf.yield %1, %1 : i32, i32
+  }
+  
+  // CHECK: return %[[FOR_RES]] : i32
+  return %0#0 : i32
+}