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

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.

446

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/include/mlir/IR/BlockAndValueMapping.h

template <typename T> T lookup(T from) const {
     auto result = lookupOrNull(from);
     assert(result && "expected 'from' to be contained within the map");
     return result;
   }
 
   /// Clears all mappings held by the mapper.
   void clear() { valueMap.clear(); }
 
+  /// Check if the map is empty.
+  bool isEmpty() { return valueMap.empty(); }
+
   /// Returns a new mapper containing the inverse mapping.
   BlockAndValueMapping getInverse() const {
     BlockAndValueMapping result;
     for (const auto &pair : valueMap)
       result.valueMap.try_emplace(pair.second, pair.first);
     return result;
   }
 

mlir/lib/Dialect/SCF/SCF.cpp

static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
   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.
 //
 // 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,
LogicalResult matchAndRewrite(scf::ForOp forOp,
     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));
       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));
       newBlockTransferArgs.push_back(Value()); // placeholder with null value
LogicalResult matchAndRewrite(ForOp forOp,
     }
     unsigned idx = 0;
     rewriter.replaceOpWithIf(forOp, newResults, [&](OpOperand &op) {
       return op.get() != newResults[idx++];
     });
     return success();
   }
 };
+
+/// Canonicalization pattern 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 remap the input value with
+/// the block argument allowing the ForOpIterArgsFolder to do the cleanup.
+struct EnableIterArgsFolderOnUnusedArgs : public OpRewritePattern<ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForOp forOp,
+                                PatternRewriter &rewriter) const override {
+    if (forOp.getNumRegionIterArgs() == 0)
+      return failure();
+
+    BlockAndValueMapping mapping;
+    auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+
+    for (BlockArgument bbArg : forOp.getRegionIterArgs()) {
+      if (!bbArg.use_empty())
+        continue;
+
+      unsigned idx = bbArg.getArgNumber() - 1 /*indVar*/;
+      Value yieldVal = yieldOp.getOperand(idx);
+      Value inputVal = forOp.getOperand(idx + forOp.getNumControlOperands());
+
+      if (yieldVal == inputVal)
+        mapping.map(yieldVal, bbArg);
+    }
+
+    if (mapping.isEmpty())
+      return failure();
+
+    unsigned sizeOperands = yieldOp.getNumOperands();
+    for (unsigned idx = 0; idx < sizeOperands; idx++) {
+      Value newOperand = mapping.lookupOrDefault(yieldOp.getOperand(idx));
+      rewriter.startRootUpdate(yieldOp);
+      yieldOp.setOperand(idx, newOperand);
+      rewriter.finalizeRootUpdate(yieldOp);
+    }
+    return success();
+  }
+};
+
 } // namespace
 
 void ForOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
-  results.insert<ForOpIterArgsFolder, SimplifyTrivialLoops,
-                 LastTensorLoadCanonicalization>(context);
+  results
+      .insert<ForOpIterArgsFolder, SimplifyTrivialLoops,
+              LastTensorLoadCanonicalization, EnableIterArgsFolderOnUnusedArgs>(
+          context);
 }
 
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {

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:.*]] = tensor_load %[[M0]] : memref<128x128xf32>
   // CHECK-NEXT: %[[R1:.*]] = 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: EnableIterArgsFolderOnUnusedSecondArgInFor
+//  CHECK-SAME:   %[[A0:[0-9a-z]*]]: i32
+func @EnableIterArgsFolderOnUnusedSecondArgInFor(%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
+}