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

[Peiming]

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

738

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?

738

I let Thomas comment on that.

747

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

747

Yeah, love the new format! Thanks!

[ThomasRaoux]

please add a test

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

738

I think this is fine.

751

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(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)
+    auto alb = op.getLowerBound().getDefiningOp<AffineApplyOp>();
+    auto aub = op.getUpperBound().getDefiningOp<AffineApplyOp>();
+
+    // ID map: (d0)->d0
+    auto id = AffineMap::getMultiDimIdentityMap(1, op.getContext());
+    auto lb =
+        alb ? alb.getAffineValueMap() : AffineValueMap(id, op.getLowerBound());
+    auto ub =
+        aub ? aub.getAffineValueMap() : AffineValueMap(id, op.getUpperBound());
+
+    AffineValueMap diffAffine;
+    AffineValueMap::difference(ub, lb, &diffAffine);
+    AffineExpr res = diffAffine.getResult(0);
+    if (!res.isa<AffineConstantExpr>())
       return failure();
 
+    int64_t diff = res.cast<AffineConstantExpr>().getValue();
     // 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",