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

[Peiming]

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

756

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?

756

I let Thomas comment on that.

765

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

765

Yeah, love the new format! Thanks!

[ThomasRaoux]

please add a test

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

756

I think this is fine.

758

nit: use dyn_cast instead of isa + cast

[Peiming]

add a test case.

[aartbik]

mlir/test/Dialect/SCF/canonicalize.mlir

734

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

%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/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFDialect
   SCF.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/SCF
 
   DEPENDS
   MLIRSCFOpsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRAffineDialect
   MLIRArithmeticDialect
   MLIRBufferizationDialect
   MLIRControlFlowDialect
   MLIRIR
   MLIRLoopLikeInterface
   MLIRSideEffectInterfaces
   )
 

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);
+  if (auto constDiff = diffMap.getResult(0).dyn_cast<AffineConstantExpr>())
+    return constDiff.getValue();
+  return llvm::None;
+}
+
 /// 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();
     }
 

mlir/test/Dialect/SCF/canonicalize.mlir

 }
 
 // CHECK-LABEL:   func @to_select_with_body
 // CHECK-DAG:       [[C0:%.*]] = arith.constant 0 : index
 // CHECK-DAG:       [[C1:%.*]] = arith.constant 1 : index
 // CHECK:           [[V0:%.*]] = arith.select {{.*}}, [[C0]], [[C1]]
 // CHECK:           scf.if {{.*}} {
 // CHECK:             "test.op"() : () -> ()
-// CHECK:           } 
+// CHECK:           }
 // CHECK:           return [[V0]] : index
 // -----
 
 func.func @to_select2(%cond: i1) -> (index, index) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
scf.if %arg0 {
     %a:2 = scf.if %arg1 -> (i32, f32) {
       %i = "test.inop"() : () -> (i32)
       %i1 = "test.inop1"() : () -> (f32)
       scf.yield %i, %i1 : i32, f32
     } else {
       scf.yield %0, %1 : i32, f32
     }
   }
-  return 
+  return
 }
 
 // CHECK-LABEL: @merge_yielding_nested_if_nv2
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 func.func @merge_yielding_nested_if_nv2(%arg0: i1, %arg1: i1) -> i32 {
 // CHECK: %[[PRE0:.*]] = "test.op"() : () -> i32
 // CHECK: %[[PRE1:.*]] = "test.op1"() : () -> i32
 // CHECK: %[[COND:.*]] = arith.andi %[[ARG0]], %[[ARG1]]
 // CHECK: %[[RES:.*]] = arith.select %[[ARG0]], %[[PRE0]], %[[PRE1]]
-// CHECK: scf.if %[[COND]] 
+// CHECK: scf.if %[[COND]]
 // CHECK:   "test.run"() : () -> ()
 // CHECK: }
 // CHECK: return %[[RES]]
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (i32)
   %r = scf.if %arg0 -> i32 {
     scf.if %arg1 {
       "test.run"() : () -> ()
%0 = scf.for %i = %c42 to %c47 step %c5 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+
+// -----
+
+#map = affine_map<(d0) -> (d0 + 1)>
+// CHECK-LABEL: func @replace_single_iteration_const_diff
+func.func @replace_single_iteration_const_diff(%arg0 : index) {
+  // CHECK-NEXT: %[[CST:.*]] = arith.constant 2
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2.0 : f32
+  %5 = affine.apply #map(%arg0)
+  // CHECK-NOT: scf.for
+  scf.for %arg2 = %arg0 to %5 step %c1 {
+    // CHECK-NEXT: %[[VAL:.*]] = "test.init"()
+    %init = "test.init"() : () -> f32
+    // CHECK-NEXT: %[[MUL:.*]] = arith.mulf %[[VAL]], %[[CST]]
+    %7 = arith.mulf %c2, %init : f32
+    // CHECK-NEXT: "test.consume"(%[[MUL]])
+    "test.consume"(%7) : (f32) -> ()
+  }
+  return
+}
+
 // -----
 
 // CHECK-LABEL: @remove_empty_parallel_loop
 func.func @remove_empty_parallel_loop(%lb: index, %ub: index, %s: index) {
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> f32
   // CHECK-NOT: scf.parallel
   // CHECK-NOT: test.produce
func.func @while_cond_true() -> i1 {
   }
   return %0 : i1
 }
 // CHECK-NEXT:         %[[true:.+]] = arith.constant true
 // CHECK-NEXT:         %{{.+}} = scf.while : () -> i1 {
 // CHECK-NEXT:           %[[cmp:.+]] = "test.condition"() : () -> i1
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[cmp]] : i1
 // CHECK-NEXT:         } do {
-// CHECK-NEXT:         ^bb0(%arg0: i1):  
+// CHECK-NEXT:         ^bb0(%arg0: i1):
 // CHECK-NEXT:           "test.use"(%[[true]]) : (i1) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
 // -----
 
 // CHECK-LABEL: @while_unused_arg
 func.func @while_unused_arg(%x : i32, %y : f64) -> i32 {
   %0 = scf.while (%arg1 = %x, %arg2 = %y) : (i32, f64) -> (i32) {
     %condition = "test.condition"(%arg1) : (i32) -> i1
     scf.condition(%condition) %arg1 : i32
   } do {
   ^bb0(%arg1: i32):
     %next = "test.use"(%arg1) : (i32) -> (i32)
     scf.yield %next, %y : i32, f64
   }
   return %0 : i32
 }
 // CHECK-NEXT:         %[[res:.*]] = scf.while (%[[arg2:.+]] = %{{.*}}) : (i32) -> i32 {
 // CHECK-NEXT:           %[[cmp:.*]] = "test.condition"(%[[arg2]]) : (i32) -> i1
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[arg2]] : i32
 // CHECK-NEXT:         } do {
-// CHECK-NEXT:         ^bb0(%[[post:.+]]: i32):  
+// CHECK-NEXT:         ^bb0(%[[post:.+]]: i32):
 // CHECK-NEXT:           %[[next:.+]] = "test.use"(%[[post]]) : (i32) -> i32
 // CHECK-NEXT:           scf.yield %[[next]] : i32
 // CHECK-NEXT:         }
 // CHECK-NEXT:         return %[[res]] : i32
 
 // -----
 
 // CHECK-LABEL: @invariant_loop_args_in_same_order
func.func @while_unused_result() -> i32 {
   return %0#0 : i32
 }
 // CHECK-NEXT:         %[[res:.*]] = scf.while : () -> i32 {
 // CHECK-NEXT:           %[[cmp:.*]] = "test.condition"() : () -> i1
 // CHECK-NEXT:           %[[val:.*]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:           %{{.*}} = "test.get_some_value"() : () -> i64
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
-// CHECK-NEXT:         ^bb0(%[[arg:.*]]: i32):  
+// CHECK-NEXT:         ^bb0(%[[arg:.*]]: i32):
 // CHECK-NEXT:           "test.use"(%[[arg]]) : (i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 // CHECK-NEXT:         return %[[res]] : i32
 
 // CHECK-LABEL: @while_cmp_lhs
 func.func @while_cmp_lhs(%arg0 : i32) {
   %0 = scf.while () : () -> i32 {

 }
 // CHECK-DAG:         %[[true:.+]] = arith.constant true
 // CHECK-DAG:         %[[false:.+]] = arith.constant false
 // CHECK-DAG:         %{{.+}} = scf.while : () -> i32 {
 // CHECK-NEXT:         %[[val:.+]] = "test.val"
 // CHECK-NEXT:         %[[cmp:.+]] = arith.cmpi ne, %[[val]], %arg0 : i32
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
-// CHECK-NEXT:         ^bb0(%arg1: i32):  
+// CHECK-NEXT:         ^bb0(%arg1: i32):
 // CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
 // CHECK-LABEL: @while_cmp_rhs
 func.func @while_cmp_rhs(%arg0 : i32) {
   %0 = scf.while () : () -> i32 {
     %val = "test.val"() : () -> i32

 }
 // CHECK-DAG:         %[[true:.+]] = arith.constant true
 // CHECK-DAG:         %[[false:.+]] = arith.constant false
 // CHECK-DAG:         %{{.+}} = scf.while : () -> i32 {
 // CHECK-NEXT:         %[[val:.+]] = "test.val"
 // CHECK-NEXT:         %[[cmp:.+]] = arith.cmpi ne, %arg0, %[[val]] : i32
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
-// CHECK-NEXT:         ^bb0(%arg1: i32):  
+// CHECK-NEXT:         ^bb0(%arg1: i32):
 // CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
 // -----
 
 // CHECK-LABEL: @combineIfs
 func.func @combineIfs(%arg0 : i1, %arg2: i64) -> (i32, i32) {

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",