scf::ForOp: Fold away iterator arguments with no use and for which the corresponding input is yielded

[chelini]

Summary

Enhance 'ForOpIterArgsFolder' to remove unused iteration arguments in a
scf::ForOp. If the block argument corresponding to the given iterator has no
use and the yielded value equals the input, we fold it away.

[nicolasvasilache]

Did you consider extending the existing ForOpIterArgsFolder to capture this case ?
The existing one will drop if you yield the bbarg directly (i.e. scf.yield %1, %arg3 : i32, i32 in your example).
It should be pretty straightforward to extend the condition here:

// Forwarded is `true` when the region `iter` argument is yielded.
bool forwarded = (std::get<1>(it) == std::get<2>(it));
keepMask.push_back(!forwarded);
canonicalize |= forwarded;
if (forwarded) {

to capture the additional condition you want ?

[rriddle]

mlir/include/mlir/IR/BlockAndValueMapping.h

80 ↗

(On Diff #330237)

These methods should match the traditional datastructure methods, i.e. empty()

[chelini]

Updating D98503: Canonicalization for scf::ForOp

• Use 'forOpIterArgsFolder' as suggested.

[chelini]

@rriddle and @nicolasvasilache thanks for the suggestions. @nicolasvasilache I updated the code as you suggested. Indeed it simplified a lot.

In D98503#2622317, @nicolasvasilache wrote:

Did you consider extending the existing ForOpIterArgsFolder to capture this case ?
The existing one will drop if you yield the bbarg directly (i.e. scf.yield %1, %arg3 : i32, i32 in your example).
It should be pretty straightforward to extend the condition here:

// Forwarded is `true` when the region `iter` argument is yielded.
bool forwarded = (std::get<1>(it) == std::get<2>(it));
keepMask.push_back(!forwarded);
canonicalize |= forwarded;
if (forwarded) {

to capture the additional condition you want ?

[nicolasvasilache]

Thanks!

mlir/lib/Dialect/SCF/SCF.cpp

411

Can you please also use a list here like you did below?

1) ...
2) ...

It helps see the case disjunction more easily and when we add more cases it will be grow naturally.

451

the iterator from .. -> the corresponding iter operand ?

[chelini]

Update comment.

[chelini]

Update missed comment.

[chelini]

In D98503#2626519, @nicolasvasilache wrote:

Thanks!

@nicolasvasilache I have just updated the comments. Could you please land this revision? (I don't have commit access) Thanks!

mlir/lib/Dialect/SCF/SCF.cpp

static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
   Operation *terminator = block->getTerminator();
   ValueRange results = terminator->getOperands();
   rewriter.mergeBlockBefore(block, op, blockArgs);
   rewriter.replaceOp(op, results);
   rewriter.eraseOp(terminator);
 }
 
 namespace {
-// Fold away ForOp iter arguments that are also yielded by the op.
-// These arguments must be defined outside of the ForOp region and can just be
-// forwarded after simplifying the op inits, yields and returns.
+// 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.
+//
+// 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> {
   using OpRewritePattern<scf::ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(scf::ForOp forOp,
                                 PatternRewriter &rewriter) const final {
LogicalResult matchAndRewrite(scf::ForOp forOp,
     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
                              yieldOp.getOperands()      // iter yield
                              )) {
-      // Forwarded is `true` when the region `iter` argument is yielded.
-      bool forwarded = (std::get<1>(it) == std::get<2>(it));
+      // Forwarded is `true` when:
+      // 1) The region `iter` argument is yielded.
+      // 2) The region `iter` argument has zero 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)));
       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));
for (unsigned idx = 0, collapsedIdx = 0, e = newResultValues.size();
       }
     }
 
     Block &oldBlock = forOp.region().front();
     assert(oldBlock.getNumArguments() == newBlockTransferArgs.size() &&
            "unexpected argument size mismatch");
 
     // No results case: the scf::ForOp builder already created a zero
-    // reult terminator. Merge before this terminator and just get rid of the
+    // result terminator. Merge before this terminator and just get rid of the
     // original terminator that has been merged in.
     if (newIterArgs.empty()) {
       auto newYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
       rewriter.mergeBlockBefore(&oldBlock, newYieldOp, newBlockTransferArgs);
       rewriter.eraseOp(newBlock.getTerminator()->getPrevNode());
       rewriter.replaceOp(forOp, newResultValues);
       return success();
     }

mlir/test/Dialect/SCF/canonicalize.mlir

%0 = scf.parallel (%i, %j, %k) = (%lb, %ub, %lb) to (%ub, %ub, %ub) step (%s, %s, %s) init(%init) -> f32 {
     scf.yield
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (f32) -> ()
   return
 }
 
 // -----
+
 func private @process(%0 : memref<128x128xf32>)
 func private @process_tensor(%0 : tensor<128x128xf32>) -> memref<128x128xf32>
 
 // CHECK-LABEL: last_value
 //  CHECK-SAME:   %[[T0:[0-9a-z]*]]: tensor<128x128xf32>
 //  CHECK-SAME:   %[[T1:[0-9a-z]*]]: tensor<128x128xf32>
 //  CHECK-SAME:   %[[T2:[0-9a-z]*]]: tensor<128x128xf32>
 //  CHECK-SAME:   %[[M0:[0-9a-z]*]]: memref<128x128xf32>
func @last_value(%t0: tensor<128x128xf32>, %t1: tensor<128x128xf32>,
   // CHECK-NEXT: }
   }
 
   // CHECK-NEXT: %[[R0:.*]] = memref.tensor_load %[[M0]] : memref<128x128xf32>
   // CHECK-NEXT: %[[R1:.*]] = memref.tensor_load %[[M1]] : memref<128x128xf32>
   // CHECK-NEXT: return %[[R0]], %[[R1]], %[[FOR_RES]] : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
   return %0#0, %0#1, %0#2 : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
 }
+
+// -----
+
+// CHECK-LABEL: fold_away_iter_with_no_use_and_yielded_input
+//  CHECK-SAME:   %[[A0:[0-9a-z]*]]: i32
+func @fold_away_iter_with_no_use_and_yielded_input(%arg0 : i32,
+                    %ub : index, %lb : index, %step : index) -> (i32, 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, %cst : i32, i32
+  }
+
+  // CHECK: return %[[FOR_RES]], %[[C32]] : i32, i32
+  return %0#0, %0#1 : i32, i32
+}