[Mlir-commits] [mlir] 1d98861 - [mlir][SCF] Fix crash in loop peeling

[Matthias Springer]

Author: Matthias Springer
Date: 2023-01-19T19:06:28+01:00
New Revision: 1d98861a7896236895d467f7e7ab4eadf7dffd82

URL: https://github.com/llvm/llvm-project/commit/1d98861a7896236895d467f7e7ab4eadf7dffd82
DIFF: https://github.com/llvm/llvm-project/commit/1d98861a7896236895d467f7e7ab4eadf7dffd82.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
    mlir/test/Dialect/SCF/for-loop-peeling.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
index 0b09dc5374d1b..4ee27e4d00343 100644
--- 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 b1cf1f2414de0..9a6d1c8c0a14c 100644
--- 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
+}