[mlir][SCF] Fix crash in loop peeling

[springerm]

Summary

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

[mehdi_amini]

Thanks for the quick fix!

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

 /// Note: `ub` is the previous upper bound of the loop (before peeling).
 /// `insideLoop` must be true for min/max ops inside the loop and false for
 /// affine.min ops inside the partial iteration. For an explanation of the other
 /// parameters, see comment of `canonicalizeMinMaxOpInLoop`.
 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))
     constraints.addBound(IntegerPolyhedron::EQ, 2, *constStep);
 
   // Add loop peeling invariant. This is the main piece of knowledge that
   // enables AffineMinOp simplification.
   if (insideLoop) {

mlir/test/Dialect/SCF/for-loop-peeling.mlir

%r1 = scf.for %iv1 = %lb1 to %ub1 step %step iter_args(%arg1 = %arg0) -> i32 {
       %0 = arith.addi %arg1, %casted : i32
       scf.yield %0 : i32
     }
     %1 = arith.addi %arg0, %r1 : i32
     scf.yield %1 : i32
   }
   return %r0 : i32
 }
+
+// -----
+
+// Make sure that there is no crash.
+
+// 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 {
+    // Step size is a symbol, so multiplication is OK.
+    %1 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + s0 * s0)>(%arg2)[%arg1]
+    %2 = arith.index_cast %1 : index to i64
+    memref.store %2, %arg0[] : memref<i64>
+  }
+  return
+}