[mlir] scf::ForOp: Propagate constants in loop body to trigger simplifications

[ftynse]

Would you mind providing a commit description in addition to the one-line summary? A good description explains _why_ this new functionality is useful, rather than what the change does.

mlir/lib/Dialect/SCF/SCF.cpp

719

Nit: could you put the results of all std::get<N>(it) into named variables? It would make the code more readable.

mlir/test/Dialect/SCF/canonicalize.mlir

523

Nit: could this match .* instead of the specific regexp?

526

I don't think it's important for this code to be on the next line, so let's just use CHECK and avoid over-constraining the tests.

536

It looks like this tests for a sequence of canonicalizations. Could we in addition have a test that exercises only the new canonicalization you are adding? This helps bisecting problems.

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

720

Is there a reason why this has to be constant for this to apply. This looks like LICM to me and otherwise doesn't need the constantop req?

[chelini]

This is a wrong approach in using the pattern rewrite see: https://mlir.llvm.org/docs/PatternRewriter/#restrictions. I will close this.

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(ForOp forOp,
     }
     unsigned idx = 0;
     rewriter.replaceOpWithIf(forOp, newResults, [&](OpOperand &op) {
       return op.get() != newResults[idx++];
     });
     return success();
   }
 };
+
+/// Propagate constants in scf::For to trigger later simplifications.
+struct PropagateConstantsInLoopBody : public OpRewritePattern<scf::ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForOp forOp,
+                                PatternRewriter &rewriter) const override {
+    if (!forOp.hasIterOperands())
+      return failure();
+
+    Block &block = forOp.getRegion().front();
+    auto yieldOp = cast<scf::YieldOp>(block.getTerminator());
+
+    for (auto it : llvm::zip(forOp.getIterOperands(), forOp.getRegionIterArgs(),
+                             yieldOp.getOperands())) {
+      ConstantOp cst = std::get<0>(it).getDefiningOp<ConstantOp>();
+      if (cst && (std::get<0>(it) == std::get<2>(it)))
+        std::get<1>(it).replaceAllUsesWith(cst.getResult());
+    }
+    return success();
+  }
+};
 } // namespace
 
 void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) {
   results.add<ForOpIterArgsFolder, SimplifyTrivialLoops,
-              LastTensorLoadCanonicalization>(context);
+              LastTensorLoadCanonicalization, PropagateConstantsInLoopBody>(
+      context);
 }
 
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {

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, %1 : i32, i32
   }
   
   // CHECK: return %[[FOR_RES]] : i32
   return %0#0 : i32
 }
+
+// -----
+
+// CHECK-LABEL: constant_prop_in_for
+//  CHECK-SAME: %[[A0:[0-9a-z]*]]: i32
+func @constant_prop_in_for(%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) {
+    // CHECK-NEXT: %{{.*}} = addi %{{.*}}, %[[C32]] : i32
+    %1 = addi %arg2, %arg3 : i32
+    %2 = addi %1, %cst : i32
+    scf.yield %2, %cst : i32, i32
+  }
+  // CHECK: return %[[FOR_RES]] : i32
+  return %0#0 : i32
+}