[mlir][SCF] Fix crash in loop peeling

Upper bound and step size should be symbols instead of dims.

Differential Revision: https://reviews.llvm.org/D142136

diff --git a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
index 0b09dc5..4ee27e4 100644 (file)
--- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
@@ -169,7 +169,8 @@ LogicalResult scf::rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op,
                                          Value iv, Value ub, Value step,
                                          bool insideLoop) {
   FlatAffineValueConstraints constraints;
-  constraints.appendDimVar({iv, ub, step});
+  constraints.appendDimVar({iv});
+  constraints.appendSymbolVar({ub, step});
   if (auto constUb = getConstantIntValue(ub))
     constraints.addBound(IntegerPolyhedron::EQ, 1, *constUb);
   if (auto constStep = getConstantIntValue(step))

diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index b1cf1f2..9a6d1c8 100644 (file)
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -275,3 +275,17 @@ func.func @nested_loops(%lb0: index, %lb1 : index, %ub0: index, %ub1: index,
   }
   return %r0 : i32
 }
+
+// -----
+
+// CHECK-LABEL: func @regression
+func.func @regression(%arg0: memref<i64>, %arg1: index) {
+  %c0 = arith.constant 0 : index
+  %0 = affine.apply affine_map<()[s0] -> (s0 * s0)>()[%arg1]
+  scf.for %arg2 = %c0 to %0 step %arg1 {
+    %1 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + s0 * s0)>(%arg2)[%arg1]
+    %2 = arith.index_cast %0 : index to i64
+    memref.store %2, %arg0[] : memref<i64>
+  }
+  return
+}