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

412

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.

448

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.
+// Fold away ForOp iter arguments that are also yielded by the op or
+// iter arguments that have no use and the corresponding outer region
+// iterator (input) is 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,
     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 iterator from
+      // outside the region (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
+}