diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp --- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp @@ -268,7 +268,8 @@ unsigned IntegerRelation::appendVar(VarKind kind, unsigned num) { unsigned pos = getNumVarKind(kind); - return insertVar(kind, pos, num); + insertVar(kind, pos, num); + return pos; } void IntegerRelation::addEquality(ArrayRef eq) { diff --git a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp --- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp @@ -211,24 +211,26 @@ unsigned dimIv = constraints.appendDimVar(iv); auto lbv = lb.dyn_cast(); - unsigned dimLb = - lbv ? constraints.appendDimVar(lbv) : constraints.appendDimVar(/*num=*/1); + unsigned symLb = lbv ? constraints.appendSymbolVar(lbv) + : constraints.appendSymbolVar(/*num=*/1); auto ubv = ub.dyn_cast(); - unsigned dimUb = - ubv ? constraints.appendDimVar(ubv) : constraints.appendDimVar(/*num=*/1); + unsigned symUb = ubv ? constraints.appendSymbolVar(ubv) + : constraints.appendSymbolVar(/*num=*/1); // If loop lower/upper bounds are constant: Add EQ constraint. Optional lbInt = getConstantIntValue(lb); Optional ubInt = getConstantIntValue(ub); if (lbInt) - constraints.addBound(IntegerPolyhedron::EQ, dimLb, *lbInt); + constraints.addBound(IntegerPolyhedron::EQ, + symLb + constraints.getNumDimVars(), *lbInt); if (ubInt) - constraints.addBound(IntegerPolyhedron::EQ, dimUb, *ubInt); + constraints.addBound(IntegerPolyhedron::EQ, + symUb + constraints.getNumDimVars(), *ubInt); // Lower bound: iv >= lb (equiv.: iv - lb >= 0) SmallVector ineqLb(constraints.getNumCols(), 0); ineqLb[dimIv] = 1; - ineqLb[dimLb] = -1; + ineqLb[symLb + constraints.getNumDimVars()] = -1; constraints.addInequality(ineqLb); // Upper bound @@ -238,14 +240,14 @@ // iv < lb + 1 // TODO: Try to derive this constraint by simplifying the expression in // the else-branch. - ivUb = rewriter.getAffineDimExpr(dimLb) + 1; + ivUb = rewriter.getAffineSymbolExpr(symLb) + 1; } else { // The loop may have more than one iteration. // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1 AffineExpr exprLb = lbInt ? rewriter.getAffineConstantExpr(*lbInt) - : rewriter.getAffineDimExpr(dimLb); + : rewriter.getAffineSymbolExpr(symLb); AffineExpr exprUb = ubInt ? rewriter.getAffineConstantExpr(*ubInt) - : rewriter.getAffineDimExpr(dimUb); + : rewriter.getAffineSymbolExpr(symUb); ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt))); } auto map = AffineMap::get( diff --git a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir --- a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir +++ b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir @@ -368,3 +368,26 @@ } return } + +// ----- + +// This is a regression test to ensure that the no assertions are failing. + +// CHECK: #[[$map:.+]] = affine_map<(d0)[s0] -> (-(d0 * (5 ceildiv s0)) + 5, 3)> +// CHECK-LABEL: func @regression_multiplication_with_sym +func.func @regression_multiplication_with_sym(%A : memref) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c4 = arith.constant 4 : index + // CHECK: %[[dummy:.*]] = "test.dummy" + %ub = "test.dummy"() : () -> (index) + // CHECK: scf.for %[[iv:.*]] = + scf.for %i = %c0 to %ub step %c1 { + // CHECK: affine.min #[[$map]](%[[iv]])[%[[dummy]]] + %1 = affine.min affine_map<(d0)[s0] -> (-(d0 * (5 ceildiv s0)) + 5, 3)>(%i)[%ub] + %2 = arith.index_cast %1: index to i64 + memref.store %2, %A[]: memref + } + return +} \ No newline at end of file