[mlir][scf] Support simple symbolic expression when simplify loops

[Peiming]

mlir/lib/Dialect/SCF/IR/SCF.cpp

758

Or should we use ()[s0]->(s0) for id map?

[aartbik]

LGTM overall, I let Nicolas/Thomas comment more.
Also, perhaps add a CHECK test?

(since my example is not submitted yet)

mlir/lib/Dialect/SCF/IR/SCF.cpp

11

Does this new dep require some bazel changes?

758

I let Thomas comment on that.

767

maybe put this whole block in a computeDiff() method that returns optional difference?
that makes the individual methods a bit easier to read?

[Peiming]

updating bazel file

mlir/lib/Dialect/SCF/IR/SCF.cpp

11

Yeah, you are right.

Added the bazel changes

[Peiming]

address comment from Aart.

[Peiming]

add comments.

[aartbik]

mlir/lib/Dialect/SCF/IR/SCF.cpp

767

Yeah, love the new format! Thanks!

[ThomasRaoux]

please add a test

mlir/lib/Dialect/SCF/IR/SCF.cpp

758

I think this is fine.

760

nit: use dyn_cast instead of isa + cast

[Peiming]

add a test case.

[aartbik]

mlir/test/Dialect/SCF/canonicalize.mlir

734 ↗

(On Diff #461337)

not that c2 x c2 always folds into 4

I think the test would be a bit better if you use

%7 = arith.mulf %arg2, %c2 : f32

and test for 2xarg0 in the folded consume

[ThomasRaoux]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

2681

I believe you need the matching change in CMakeLists.txt

[Peiming]

update test case

[Peiming]

add dependency in CMake file

[Peiming]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

2681

Good catch!

[Peiming]

remove unused parameter in test case

[aartbik]

This comment has been deleted.

mlir/test/Dialect/SCF/canonicalize.mlir

737 ↗

(On Diff #461349)

%init -> %arg2 so that we see %arg2 -> %arg1 change?

[Peiming]

update test case

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

rebase against main

[ftynse]

Adding a dependency from SCF on Affine sounds like a quite unfortunate layering inversion...

[Peiming]

In D134204#3801947, @ftynse wrote:

Adding a dependency from SCF on Affine sounds like a quite unfortunate layering inversion...

Do you want the commit to be reverted?

[Peiming]

In D134204#3801947, @ftynse wrote:

Adding a dependency from SCF on Affine sounds like a quite unfortunate layering inversion...

See https://reviews.llvm.org/D134291

mlir/lib/Dialect/SCF/IR/SCF.cpp

 //===- SCF.cpp - Structured Control Flow Operations -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/FunctionInterfaces.h"
LogicalResult matchAndRewrite(scf::ForOp forOp,
     auto mergedYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
     cloneFilteredTerminator(mergedYieldOp);
     rewriter.eraseOp(mergedYieldOp);
     rewriter.replaceOp(forOp, newResultValues);
     return success();
   }
 };
 
+/// Util function that tries to compute a constant diff between u and l.
+/// Returns llvm::None when the difference between two AffineValueMap is
+/// dynamic.
+static Optional<int64_t> computeConstDiff(Value l, Value u) {
+  auto alb = l.getDefiningOp<AffineApplyOp>();
+  auto aub = u.getDefiningOp<AffineApplyOp>();
+  // ID map: (d0)->d0
+  auto id = AffineMap::getMultiDimIdentityMap(1, l.getContext());
+  auto lb = alb ? alb.getAffineValueMap() : AffineValueMap(id, l);
+  auto ub = aub ? aub.getAffineValueMap() : AffineValueMap(id, u);
+
+  AffineValueMap diffMap;
+  AffineValueMap::difference(ub, lb, &diffMap);
+  AffineExpr diff = diffMap.getResult(0);
+  if (!diff.isa<AffineConstantExpr>())
+    return llvm::None;
+
+  return diff.cast<AffineConstantExpr>().getValue();
+}
+
 /// Rewriting pattern that erases loops that are known not to iterate, replaces
 /// single-iteration loops with their bodies, and removes empty loops that
 /// iterate at least once and only return values defined outside of the loop.
 struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForOp op,
                                 PatternRewriter &rewriter) const override {
     // If the upper bound is the same as the lower bound, the loop does not
     // iterate, just remove it.
     if (op.getLowerBound() == op.getUpperBound()) {
       rewriter.replaceOp(op, op.getIterOperands());
       return success();
     }
 
-    auto lb = op.getLowerBound().getDefiningOp<arith::ConstantOp>();
-    auto ub = op.getUpperBound().getDefiningOp<arith::ConstantOp>();
-    if (!lb || !ub)
+    Optional<int64_t> diff =
+        computeConstDiff(op.getLowerBound(), op.getUpperBound());
+    if (!diff)
       return failure();
 
     // If the loop is known to have 0 iterations, remove it.
-    llvm::APInt lbValue = lb.getValue().cast<IntegerAttr>().getValue();
-    llvm::APInt ubValue = ub.getValue().cast<IntegerAttr>().getValue();
-    if (lbValue.sge(ubValue)) {
+    if (*diff <= 0) {
       rewriter.replaceOp(op, op.getIterOperands());
       return success();
     }
 
     auto step = op.getStep().getDefiningOp<arith::ConstantOp>();
     if (!step)
       return failure();
 
     // If the loop is known to have 1 iteration, inline its body and remove the
     // loop.
     llvm::APInt stepValue = step.getValue().cast<IntegerAttr>().getValue();
-    if ((lbValue + stepValue).sge(ubValue)) {
+    if (stepValue.sge(*diff)) {
       SmallVector<Value, 4> blockArgs;
       blockArgs.reserve(op.getNumIterOperands() + 1);
       blockArgs.push_back(op.getLowerBound());
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     ),
     hdrs = glob(
         [
             "include/mlir/Dialect/SCF/IR/*.h",
         ],
     ),
     includes = ["include"],
     deps = [
+        ":AffineDialect",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":BufferizationDialect",
         ":ControlFlowDialect",
         ":ControlFlowInterfaces",
         ":FuncDialect",
         ":IR",
         ":LoopLikeInterface",